Implement for loops using desugar

10 files changed, 182 insertions, 26 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);
     }
 
diff --git a/src/main/java/org/zwobble/couscous/backends/csharp/CsharpCodeGenerator.java b/src/main/java/org/zwobble/couscous/backends/csharp/CsharpCodeGenerator.java
index f542625..28d2a3e 100644
--- a/src/main/java/org/zwobble/couscous/backends/csharp/CsharpCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/csharp/CsharpCodeGenerator.java
@@ -6,6 +6,7 @@ import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.backends.csharp.primitives.CsharpPrimitiveMethods;
 import org.zwobble.couscous.backends.naming.Naming;
 import org.zwobble.couscous.transforms.DesugarForEachToWhile;
+import org.zwobble.couscous.transforms.DesugarForToWhile;
 import org.zwobble.couscous.types.*;
 
 import java.util.Map;
@@ -45,8 +46,10 @@ public class CsharpCodeGenerator {
     private Node generateCode(Node node) {
         return nodeTransformer.transform(
             CsharpPrimitiveMethods.TRANSFORMER.transform(
-                DesugarForEachToWhile.transformer().transform(
-                    node
+                DesugarForToWhile.transformer().transform(
+                    DesugarForEachToWhile.transformer().transform(
+                        node
+                    )
                 )
             )
         );
diff --git a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
index 477a26f..29e7274 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -11,6 +11,7 @@ import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.backends.naming.Names;
 import org.zwobble.couscous.backends.python.ast.*;
 import org.zwobble.couscous.transforms.DesugarForEachToWhile;
+import org.zwobble.couscous.transforms.DesugarForToWhile;
 import org.zwobble.couscous.transforms.DesugarSwitchToIfElse;
 import org.zwobble.couscous.types.ScalarType;
 import org.zwobble.couscous.types.Type;
@@ -61,8 +62,13 @@ import static org.zwobble.couscous.util.ExtraLists.*;
 public class PythonCodeGenerator {
     public static PythonModuleNode generateCode(TypeNode sugaredTypeNode) {
         NodeTransformer switchTransformer = DesugarSwitchToIfElse.transformer();
+        NodeTransformer forTransformer = DesugarForToWhile.transformer();
         NodeTransformer forEachTransformer = DesugarForEachToWhile.transformer();
-        TypeNode typeNode = forEachTransformer.transformTypeDeclaration(switchTransformer.transformTypeDeclaration(sugaredTypeNode));
+        TypeNode typeNode = forTransformer.transformTypeDeclaration(
+            forEachTransformer.transformTypeDeclaration(
+                switchTransformer.transformTypeDeclaration(sugaredTypeNode)
+            )
+        );
 
         Iterator<PythonImportNode> imports = generateImports(typeNode).iterator();
 
diff --git a/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java b/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java
index 72f94d5..6e9444c 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java
@@ -20,6 +20,7 @@ import static org.zwobble.couscous.ast.VariableReferenceNode.reference;
 import static org.zwobble.couscous.ast.WhileNode.whileLoop;
 import static org.zwobble.couscous.ast.sugar.SwitchCaseNode.switchCase;
 import static org.zwobble.couscous.frontends.java.JavaTypes.typeOf;
+import static org.zwobble.couscous.util.ExtraIterables.only;
 import static org.zwobble.couscous.util.ExtraIterables.takeUntil;
 import static org.zwobble.couscous.util.ExtraLists.*;
 import static org.zwobble.couscous.util.Tails.tail;
@@ -180,19 +181,15 @@ class JavaStatementReader {
     private List<StatementNode> readForStatement(ForStatement statement) {
         @SuppressWarnings("unchecked")
         List<VariableDeclarationExpression> javaInitialisers = statement.initializers();
-        if (javaInitialisers.size() != 1) {
-            throw new UnsupportedOperationException();
-        }
         @SuppressWarnings("unchecked")
         List<Expression> updaters = statement.updaters();
 
-        return append(
-            readVariableDeclarationExpression(javaInitialisers.get(0)),
-            whileLoop(
-                readExpression(Types.BOOLEAN, statement.getExpression()),
-                concat(
-                    readBody(statement.getBody()),
-                    eagerMap(updaters, updater -> expressionStatement(readExpressionWithoutBoxing(updater))))));
+        return list(new ForNode(
+            readVariableDeclarationExpression(only(javaInitialisers)),
+            readExpression(Types.BOOLEAN, statement.getExpression()),
+            eagerMap(updaters, this::readExpressionWithoutBoxing),
+            readBody(statement.getBody())
+        ));
     }
 
     private List<StatementNode> readEnhancedForStatement(EnhancedForStatement statement) {
diff --git a/src/main/java/org/zwobble/couscous/interpreter/JavaProject.java b/src/main/java/org/zwobble/couscous/interpreter/JavaProject.java
index 7b5a0bc..3e4b182 100644
--- a/src/main/java/org/zwobble/couscous/interpreter/JavaProject.java
+++ b/src/main/java/org/zwobble/couscous/interpreter/JavaProject.java
@@ -8,6 +8,7 @@ import org.zwobble.couscous.interpreter.types.IntrinsicInterpreterType;
 import org.zwobble.couscous.interpreter.types.UserDefinedInterpreterType;
 import org.zwobble.couscous.interpreter.values.BoxedIntegerInterpreterValue;
 import org.zwobble.couscous.interpreter.values.InternalCouscousInterpreterValue;
+import org.zwobble.couscous.transforms.DesugarForToWhile;
 import org.zwobble.couscous.transforms.DesugarSwitchToIfElse;
 
 import java.util.List;
@@ -24,10 +25,11 @@ public class JavaProject {
     }
 
     public static Project of(List<TypeNode> classNodes) {
-        NodeTransformer transformer = DesugarSwitchToIfElse.transformer();
+        NodeTransformer switchTransformer = DesugarSwitchToIfElse.transformer();
+        NodeTransformer forTransformer = DesugarForToWhile.transformer();
         Iterable<InterpreterType> concreteTypes = Iterables.transform(
             classNodes,
-            classNode -> new UserDefinedInterpreterType(transformer.transformTypeDeclaration(classNode))
+            classNode -> new UserDefinedInterpreterType(forTransformer.transformTypeDeclaration(switchTransformer.transformTypeDeclaration(classNode)))
         );
         return builder()
                 .addClasses(concreteTypes)
diff --git a/src/main/java/org/zwobble/couscous/transforms/DesugarForToWhile.java b/src/main/java/org/zwobble/couscous/transforms/DesugarForToWhile.java
new file mode 100644
index 0000000..5989838
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/transforms/DesugarForToWhile.java
@@ -0,0 +1,44 @@
+package org.zwobble.couscous.transforms;
+
+import org.zwobble.couscous.ast.ForNode;
+import org.zwobble.couscous.ast.NodeTypes;
+import org.zwobble.couscous.ast.StatementNode;
+import org.zwobble.couscous.ast.visitors.NodeTransformer;
+import org.zwobble.couscous.frontends.java.Scope;
+
+import java.util.List;
+import java.util.Optional;
+
+import static org.zwobble.couscous.ast.ExpressionStatementNode.expressionStatement;
+import static org.zwobble.couscous.ast.WhileNode.whileLoop;
+import static org.zwobble.couscous.util.ExtraIterables.lazyMap;
+import static org.zwobble.couscous.util.ExtraLists.append;
+import static org.zwobble.couscous.util.ExtraLists.concat;
+
+public class DesugarForToWhile {
+    public static NodeTransformer transformer() {
+        Scope scope = Scope.create().temporaryPrefix("_couscous_desugar_for_to_while");
+        return NodeTransformer.builder()
+            .transformStatement(statement -> {
+                if (statement.type() == NodeTypes.FOR) {
+                    return Optional.of(desugar(scope, ((ForNode) statement)));
+                } else {
+                    return Optional.empty();
+                }
+            })
+            .build();
+    }
+
+    public static List<StatementNode> desugar(Scope scope, ForNode forNode) {
+        return append(
+            forNode.getInitializers(),
+            whileLoop(
+                forNode.getCondition(),
+                concat(
+                    forNode.getStatements(),
+                    lazyMap(forNode.getUpdaters(), updater -> expressionStatement(updater))
+                )
+            )
+        );
+    }
+}
diff --git a/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReaderTests.java b/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReaderTests.java
index c7ad81d..fdcfce6 100644
--- a/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReaderTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReaderTests.java
@@ -166,17 +166,17 @@ public class JavaReaderTests {
 
     @Test
     public void canReadForLoops() {
-        List<StatementNode> statements = readStatements("int", "for (int x = 0; true; ++x) { return 1; }");
-        LocalVariableDeclarationNode declaration = (LocalVariableDeclarationNode) statements.get(0);
+        StatementNode statement = readStatement("int", "for (int x = 0; true; ++x) { return 1; }");
+        VariableDeclaration loopVariable = var(MAIN_IDENTIFIER.variable("x"), "x", Types.INT);
         assertEquals(
-            list(
-                localVariableDeclaration(declaration.getDeclaration(), literal(0)),
-                whileLoop(
-                    literal(true),
-                    list(
-                        returns(literal(1)),
-                        expressionStatement(assign(declaration, integerAdd(reference(declaration), literal(1))))))),
-            statements);
+            new ForNode(
+                list(localVariableDeclaration(loopVariable, literal(0))),
+                literal(true),
+                list(assign(reference(loopVariable), integerAdd(reference(loopVariable), literal(1)))),
+                list(returns(literal(1)))
+            ),
+            statement
+        );
     }
 
     @Test
diff --git a/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java b/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java
index e64de11..eb6e839 100644
--- a/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java
+++ b/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java
@@ -38,6 +38,7 @@ public class JavaReading {
     }
 
     static StatementNode readStatement(String returnType, String statementSource) {
+        // TODO: use only()
         return readStatements(returnType, statementSource).get(0);
     }