Implement for loops using desugar

3 files changed, 104 insertions, 1 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/ForNode.java b/src/main/java/org/zwobble/couscous/ast/ForNode.java
new file mode 100644
index 0000000..4642b85
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/ast/ForNode.java
@@ -0,0 +1,101 @@
+package org.zwobble.couscous.ast;
+
+import com.google.common.collect.Iterables;
+import org.zwobble.couscous.ast.visitors.NodeTransformer;
+import org.zwobble.couscous.util.ExtraIterables;
+import org.zwobble.couscous.util.ExtraLists;
+
+import java.util.List;
+
+public class ForNode implements StatementNode {
+    private final List<LocalVariableDeclarationNode> initializers;
+    private final ExpressionNode condition;
+    private final List<ExpressionNode> updaters;
+    private final List<StatementNode> statements;
+
+    public ForNode(
+        List<LocalVariableDeclarationNode> initializers,
+        ExpressionNode condition,
+        List<ExpressionNode> updaters,
+        List<StatementNode> statements
+    ) {
+        this.initializers = initializers;
+        this.condition = condition;
+        this.updaters = updaters;
+        this.statements = statements;
+    }
+
+    public List<LocalVariableDeclarationNode> getInitializers() {
+        return initializers;
+    }
+
+    public ExpressionNode getCondition() {
+        return condition;
+    }
+
+    public List<ExpressionNode> getUpdaters() {
+        return updaters;
+    }
+
+    public List<StatementNode> getStatements() {
+        return statements;
+    }
+
+    @Override
+    public int type() {
+        return NodeTypes.FOR;
+    }
+
+    @Override
+    public Iterable<? extends Node> childNodes() {
+        return Iterables.concat(
+            initializers,
+            ExtraIterables.of(condition),
+            updaters,
+            statements
+        );
+    }
+
+    @Override
+    public StatementNode transformSubtree(NodeTransformer transformer) {
+        return new ForNode(
+            ExtraLists.copyOf(ExtraIterables.cast(LocalVariableDeclarationNode.class, transformer.transformStatements(initializers))),
+            transformer.transformExpression(condition),
+            transformer.transformExpressions(updaters),
+            transformer.transformStatements(statements)
+        );
+    }
+
+    @Override
+    public String toString() {
+        return "ForNode(" +
+            "initializers=" + initializers +
+            ", condition=" + condition +
+            ", updaters=" + updaters +
+            ", statements=" + statements +
+            ')';
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        ForNode forNode = (ForNode) o;
+
+        if (!initializers.equals(forNode.initializers)) return false;
+        if (!condition.equals(forNode.condition)) return false;
+        if (!updaters.equals(forNode.updaters)) return false;
+        return statements.equals(forNode.statements);
+
+    }
+
+    @Override
+    public int hashCode() {
+        int result = initializers.hashCode();
+        result = 31 * result + condition.hashCode();
+        result = 31 * result + updaters.hashCode();
+        result = 31 * result + statements.hashCode();
+        return result;
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/ast/NodeTypes.java b/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
index 82090aa..7edaaab 100644
--- a/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
+++ b/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
@@ -44,6 +44,7 @@ public class NodeTypes {
     public static final int TYPE_COERCION = 10028;
     public static final int VARIABLE_REFERENCE = 10029;
     public static final int WHILE = 10030;
+    public static final int FOR = 10036;
     public static final int FOR_EACH = 10034;
     public static final int STATEMENT_BLOCK = 10035;
 
@@ -85,6 +86,7 @@ public class NodeTypes {
         .put(TypeCoercionNode.class, TYPE_COERCION)
         .put(VariableReferenceNode.class, VARIABLE_REFERENCE)
         .put(WhileNode.class, WHILE)
+        .put(ForNode.class, FOR)
         .put(ForEachNode.class, FOR_EACH)
         .put(StatementBlockNode.class, STATEMENT_BLOCK)
         .put(PrimitiveInstanceMethodCall.class, CsharpNodeTypes.PRIMITIVE_INSTANCE_METHOD_CALL)
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java b/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
index a1bbda6..d44bb39 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
@@ -216,7 +216,7 @@ public class NodeTransformer {
         return transformList(methods, this::transformMethod);
     }
 
-    public List<StatementNode> transformStatements(List<StatementNode> body) {
+    public List<StatementNode> transformStatements(List<? extends StatementNode> body) {
         return eagerFlatMap(body, this::transformStatement);
     }