/src/org/charter_project/graph/compile/statement/ForEach.java

package org.charter_project.graph.compile.statement;

import org.charter_project.graph.client.ExpressionType;
import org.charter_project.graph.client.GenerateException;
import org.charter_project.graph.compile.Code.FailCode;
import org.charter_project.graph.compile.LoopCode;
import org.charter_project.graph.compile.Parsed;
import org.charter_project.graph.compile.RuleFile;
import org.charter_project.graph.compile.block.Block;
import org.charter_project.graph.compile.block.Block.BlockType;
import org.charter_project.graph.compile.expression.Cast;
import org.charter_project.graph.compile.expression.Expression;
import org.charter_project.graph.compile.expression.GetSize;
import org.charter_project.graph.compile.expression.Literal;
import org.charter_project.graph.compile.expression.SelectMField;
import org.charter_project.graph.compile.expression.Var;
import org.charter_project.graph.parse.MyTree;

/**
 * A <code>ForEach</code> represents the 'foreach' alternative of a
 * {@link Statement}.
 *
 * @author Maarten de Mol
 * @version 1.0
 */
public class ForEach extends Statement implements AssignRHS
{
  // The for variable.
  private final Var        loopVar;

  // The collection expression. May be null.
  private final Expression coll;

  // The body.
  private final Block      body;

  // The collects expression. May be null.
  private final Expression collects;

  // Additional flag to disable initial comment line.
  private boolean          disableComment = false;

  /** Constructor for expressions provided by ANTLR. */
  public ForEach(MyTree tree, Var loopVar, Expression coll, Block body,
                 Expression collects)
  {
    super(tree);
    this.loopVar = loopVar;
    this.coll = coll;
    this.body = body;
    this.collects = collects;
  }

  /** Constructor for custom expressions, created during code generation. */
  public ForEach(Parsed source, Var loopVar, Expression coll, Block body,
                 Expression collects)
  {
    super(source);
    this.loopVar = loopVar;
    this.coll = coll;
    this.body = body;
    this.collects = collects;
  }

  /** Getter for the loop variable. */
  public Var getLoopVar()
  {
    return this.loopVar;
  }

  /** Getter for the collection expression. */
  public Expression getCollection()
  {
    return this.coll;
  }

  /** Getter for the body. */
  public Block getBody()
  {
    return this.body;
  }

  /** Getter for the return expression. */
  public Expression getCollects()
  {
    return this.collects;
  }

  /** Setter for the disable comment flag. */
  public void disableComment()
  {
    this.disableComment = true;
  }

  @Override
  public ExpressionType getType()
  {
    return this.collects.getType();
  }

  @Override
  public void show(Indented builder)
  {
    builder.write("foreach (");
    builder.write(this.loopVar.getType().toString());
    builder.write(" ");
    builder.write(this.loopVar.getName());
    if (this.coll != null)
      {
        builder.write(" : ");
        builder.write(this.coll.show(false));
      }
    builder.write(") ");
    builder.write("{");
    builder.incIndent();
    builder.newline();
    this.body.show(builder, false);
    if (this.collects != null)
      {
        builder.newline();
        builder.write("return ");
        builder.write(this.collects.show(false));
        builder.write(";");
      }
    builder.decIndent();
    builder.newline();
    builder.write("}");
  }

  // ========================================================================
  // LOOP CODE
  // ========================================================================

  /** Generates the {@link LoopCode} for this foreach statement. */
  public LoopCode generateLoopCode(RuleFile file)
  {
    if (this.coll == null)
      {
        return new LoopAllCode();
      }
    else if (this.coll instanceof SelectMField &&
             file.inBlockType() != BlockType.SEQUENCE)
      {
        SelectMField select = (SelectMField) this.coll;
        return new LoopFieldCode(select);
      }
    else if (this.coll instanceof GetSize)
      {
        GetSize getSize = (GetSize) this.coll;
        return new LoopIndexCode(getSize);
      }
    else
      {
        return new MyLoopCode();
      }
  }

  /**
   * Helper method for printing a comment for the 'collects' expression at the
   * end of the loop body.
   */
  private static void writeCollectsComment(Expression collects, RuleFile file)
    throws GenerateException
  {
    if (collects != null)
      {
        file.newline();
        file.write("// collect ");
        file.write(collects.show(false));
        file.newline();
      }
  }

  /**
   * Helper class for writing loop case, for the case that no collection
   * expression was provided.
   */
  private class LoopAllCode implements LoopCode
  {

    // Loop code for the bare node visitor.
    private NodeVisitCode baseLoop;

    @Override
    public void writeOpen(FailCode fail, RuleFile file)
      throws GenerateException
    {
      this.baseLoop =
        new NodeVisitCode(ForEach.this, ForEach.this.loopVar.getType()
            .getNode());
      if (ForEach.this.collects == null && !ForEach.this.disableComment)
        {
          generateComment(file);
        }
      this.baseLoop.writeOpen(fail, file);
      ForEach.this.loopVar.writeAssignCode(this.baseLoop.getValue(), file);
      ForEach.this.body.write(FailCode.VISITOR_FAIL, file);
      writeCollectsComment(ForEach.this.collects, file);
    }

    @Override
    public void writeClose(boolean ruleCall, boolean endReachable,
                           FailCode fail, RuleFile file)
      throws GenerateException
    {
      this.baseLoop.writeClose(ruleCall || ForEach.this.body.hasRuleCall(),
                               endReachable, fail, file);
    }

    @Override
    public void writeContinue(RuleFile file)
      throws GenerateException
    {
      this.baseLoop.writeContinue(file);
    }

    @Override
    public void writeBreak(RuleFile file)
      throws GenerateException
    {
      this.baseLoop.writeBreak(file);
    }

    @Override
    public Expression getValue()
    {
      return ForEach.this.collects;
    }
  }

  /**
   * Helper class for writing loop case, for the case that the collection
   * expression is a field selection.
   */
  private class LoopFieldCode implements LoopCode
  {

    // The select expression, and its generated loop code.
    private final SelectMField  select;
    private LoopCode            loop;

    /** Default constructor. */
    public LoopFieldCode(SelectMField select)
    {
      this.select = select;
    }

    @Override
    public void writeOpen(FailCode fail, RuleFile file)
      throws GenerateException
    {
      // Comment
      if (ForEach.this.collects == null && !ForEach.this.disableComment)
        {
          generateComment(file);
        }
      // Create index variable, if necessary.
      Var index = null;
      if (file.getRule().getVarsWithIndex()
          .contains(ForEach.this.loopVar.getName()) &&
          select.getType().isList())
        {
          index = file.generateClassVar(ForEach.this, ExpressionType.INT);
          index.writeAssignCode(Literal.MINUS_ONE(ForEach.this).getCode(),
                                file);
          file.getIndexVarMap().put(ForEach.this.loopVar.getName(), index);
        }
      // Open visitor.
      this.loop = select.generateLoopCode();
      this.loop.writeOpen(fail, file);
      // Body statements (assign loopVar)
      this.loop.getValue().writeHelperCode(fail, file);
      if (!ForEach.this.loopVar.getType().equals(this.loop.getValue()
                                                     .getType()))
        {
          // type cast before assign
          file.write("if (!(");
          this.loop.getValue().getCode().write(true, file);
          file.write(" instanceof ");
          file.write(ForEach.this.loopVar.getType());
          file.write(")) ");
          file.openBlock();
          this.loop.writeContinue(file);
          file.closeBlock();
          Cast cast =
            new Cast(ForEach.this, false, ForEach.this.loopVar.getType(),
                     this.loop.getValue());
          cast.writeHelperCode(FailCode.EMPTY, file);
          ForEach.this.loopVar.writeAssignCode(cast.getCode(), file);
        }
      else
        {
          // assign as is
          ForEach.this.loopVar
              .writeAssignCode(this.loop.getValue().getCode(), file);
        }
      // Body statements (increase indexVar).
      if (index != null)
        {
          file.write(index);
          file.write("++;");
          file.newline();
        }
      // Body statements (for body, collects)
      ForEach.this.body.write(FailCode.VISITOR_FAIL, file);
      writeCollectsComment(ForEach.this.collects, file);
    }

    @Override
    public void writeClose(boolean ruleCall, boolean endReachable,
                           FailCode fail, RuleFile file)
      throws GenerateException
    {
      this.loop.writeClose(ruleCall || ForEach.this.body.hasRuleCall(),
                           endReachable, fail, file);
    }

    @Override
    public void writeContinue(RuleFile file)
      throws GenerateException
    {
      this.loop.writeContinue(file);
    }

    @Override
    public void writeBreak(RuleFile file)
      throws GenerateException
    {
      this.loop.writeBreak(file);
    }

    @Override
    public Expression getValue()
    {
      return ForEach.this.collects;
    }
  }

  /**
   * Helper class for writing loop case, for the case that the collection
   * expression is a getSize operation. This alternative cannot be reached
   * directly, and is only supported for the 'restrict to indexOf' alternative
   * in a match block.
   */
  private class LoopIndexCode implements LoopCode
  {

    // The getSize expression.
    private final GetSize getSize;

    // The label of the generated for loop.
    private String        label;

    /** Default constructor. */
    public LoopIndexCode(GetSize getSize)
    {
      this.getSize = getSize;
    }

    @Override
    public void writeOpen(FailCode fail, RuleFile file)
      throws GenerateException
    {
      // Comment
      if (ForEach.this.collects == null && !ForEach.this.disableComment)
        {
          generateComment(file);
        }
      // Variables.
      Var tempVar =
        file.generateLocalVar(ForEach.this, ForEach.this.loopVar.getType());
      this.getSize.writeHelperCode(fail, file);
      // Open for loop.
      this.label = file.freshLabel();
      file.write(label);
      file.write(": for (int ");
      file.write(tempVar);
      file.write("= 0; ");
      file.write(tempVar);
      file.write(" < ");
      this.getSize.getCode().write(true, file);
      file.write("; ");
      file.write(tempVar);
      file.write("++) ");
      file.openBlock();
      // Body statements.
      ForEach.this.loopVar.writeAssignCode(tempVar, file);
      ForEach.this.body.write(FailCode.VISITOR_FAIL, file);
      writeCollectsComment(ForEach.this.collects, file);
    }

    @Override
    public void writeClose(boolean ruleCall, boolean endReachable,
                           FailCode fail, RuleFile file)
      throws GenerateException
    {
      file.closeBlock();
    }

    @Override
    public void writeContinue(RuleFile file)
      throws GenerateException
    {
      file.write("continue ");
      file.write(this.label);
      file.write(";");
      file.newline();
    }

    @Override
    public void writeBreak(RuleFile file)
      throws GenerateException
    {
      file.write("break ");
      file.write(this.label);
      file.write(";");
      file.newline();
    }

    @Override
    public Expression getValue()
    {
      return ForEach.this.collects;
    }
  }

  /**
   * Helper class for writing loop case, for the case that the collection
   * expression cannot be decomposed.
   */
  private class MyLoopCode implements LoopCode
  {

    // Label of the for loop.
    private String label;

    @Override
    public void writeOpen(FailCode fail, RuleFile file)
      throws GenerateException
    {
      // Comment
      if (ForEach.this.collects == null && !ForEach.this.disableComment)
        {
          generateComment(file);
        }
      // Variables.
      Var tempVar =
        file.generateLocalVar(ForEach.this, ForEach.this.coll.getType()
            .getElement());
      // Helper code.
      ForEach.this.coll.writeHelperCode(fail, file);
      // Create index variable, if necessary.
      Var index = null;
      if (file.getRule().getVarsWithIndex()
          .contains(ForEach.this.loopVar.getName()) &&
          ForEach.this.coll.getType().isList())
        {
          index = file.generateClassVar(ForEach.this, ExpressionType.INT);
          index.writeAssignCode(Literal.MINUS_ONE(ForEach.this).getCode(),
                                file);
          file.getIndexVarMap().put(ForEach.this.loopVar.getName(), index);
        }
      // Open for loop
      this.label = file.freshLabel();
      file.write(label);
      file.write(": for (");
      file.write(tempVar.getType());
      file.write(" ");
      file.write(tempVar.getName());
      file.write(" : ");
      ForEach.this.coll.getCode().write(false, file);
      file.write(") ");
      file.openBlock();
      // Body statements (assign loopVar)
      if (!ForEach.this.loopVar.getType().equals(tempVar.getType()))
        {
          // type cast before assign
          file.write("if (!(");
          file.write(tempVar);
          file.write(" instanceof ");
          file.write(ForEach.this.loopVar.getType());
          file.write(")) ");
          file.openBlock();
          writeContinue(file);
          file.closeBlock();
          Cast cast =
            new Cast(ForEach.this, false, ForEach.this.loopVar.getType(),
                     tempVar);
          cast.writeHelperCode(FailCode.EMPTY, file);
          ForEach.this.loopVar.writeAssignCode(cast.getCode(), file);
        }
      else
        {
          ForEach.this.loopVar.writeAssignCode(tempVar, file);
        }

      // Body statements (increase indexVar).
      if (index != null)
        {
          file.write(index);
          file.write("++;");
          file.newline();
        }
      // Body statements (for body, collects)
      ForEach.this.body.write(fail, file);
      writeCollectsComment(ForEach.this.collects, file);
    }

    @Override
    public void writeClose(boolean ruleCall, boolean endReachable,
                           FailCode fail, RuleFile file)
      throws GenerateException
    {
      file.closeBlock();
    }

    @Override
    public void writeContinue(RuleFile file)
      throws GenerateException
    {
      file.write("continue ");
      file.write(this.label);
      file.write(";");
      file.newline();
    }

    @Override
    public void writeBreak(RuleFile file)
      throws GenerateException
    {
      file.write("break ");
      file.write(this.label);
      file.write(";");
      file.newline();
    }

    @Override
    public Expression getValue()
    {
      return ForEach.this.collects;
    }
  }

  // ========================================================================
  // CODE GENERATION
  // ========================================================================

  @Override
  public void write(FailCode fail, RuleFile file)
    throws GenerateException
  {
    LoopCode loop = generateLoopCode(file);
    loop.writeOpen(fail, file);
    // collects is not supported in this mode
    loop.writeClose(false, true, fail, file);
  }
}