Read for-each as sugared statements

7 files changed, 157 insertions, 12 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/ForEachNode.java b/src/main/java/org/zwobble/couscous/ast/ForEachNode.java
new file mode 100644
index 0000000..e30f130
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/ast/ForEachNode.java
@@ -0,0 +1,79 @@
+package org.zwobble.couscous.ast;
+
+import org.zwobble.couscous.ast.visitors.NodeTransformer;
+import org.zwobble.couscous.util.ExtraIterables;
+
+import java.util.List;
+
+public class ForEachNode implements StatementNode {
+    private final VariableDeclaration target;
+    private final ExpressionNode iterable;
+    private final List<StatementNode> statements;
+
+    public ForEachNode(VariableDeclaration target, ExpressionNode iterable, List<StatementNode> statements) {
+        this.target = target;
+        this.iterable = iterable;
+        this.statements = statements;
+    }
+
+    public VariableDeclaration getTarget() {
+        return target;
+    }
+
+    public ExpressionNode getIterable() {
+        return iterable;
+    }
+
+    public List<StatementNode> getStatements() {
+        return statements;
+    }
+
+    @Override
+    public int type() {
+        return NodeTypes.FOR_EACH;
+    }
+
+    @Override
+    public Iterable<? extends Node> childNodes() {
+        return ExtraIterables.of(iterable);
+    }
+
+    @Override
+    public StatementNode transformSubtree(NodeTransformer transformer) {
+        return new ForEachNode(
+            transformer.transform(target),
+            transformer.transformExpression(iterable),
+            transformer.transformStatements(statements)
+        );
+    }
+
+    @Override
+    public String toString() {
+        return "ForEachNode(" +
+            "target=" + target +
+            ", iterable=" + iterable +
+            ", statements=" + statements +
+            ')';
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        ForEachNode that = (ForEachNode) o;
+
+        if (!target.equals(that.target)) return false;
+        if (!iterable.equals(that.iterable)) return false;
+        return statements.equals(that.statements);
+
+    }
+
+    @Override
+    public int hashCode() {
+        int result = target.hashCode();
+        result = 31 * result + iterable.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 9a9a244..f7f8bf6 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_EACH = 10034;
 
     private NodeTypes() {
     }
@@ -83,6 +84,7 @@ public class NodeTypes {
         .put(TypeCoercionNode.class, TYPE_COERCION)
         .put(VariableReferenceNode.class, VARIABLE_REFERENCE)
         .put(WhileNode.class, WHILE)
+        .put(ForEachNode.class, FOR_EACH)
         .put(PrimitiveInstanceMethodCall.class, CsharpNodeTypes.PRIMITIVE_INSTANCE_METHOD_CALL)
         .put(PrimitiveStaticMethodCall.class, CsharpNodeTypes.PRIMITIVE_STATIC_METHOD_CALL)
         .build();
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 75a23ec..f542625 100644
--- a/src/main/java/org/zwobble/couscous/backends/csharp/CsharpCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/csharp/CsharpCodeGenerator.java
@@ -5,6 +5,7 @@ import org.zwobble.couscous.ast.Node;
 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.types.*;
 
 import java.util.Map;
@@ -42,7 +43,13 @@ public class CsharpCodeGenerator {
     }
 
     private Node generateCode(Node node) {
-        return nodeTransformer.transform(CsharpPrimitiveMethods.TRANSFORMER.transform(node));
+        return nodeTransformer.transform(
+            CsharpPrimitiveMethods.TRANSFORMER.transform(
+                DesugarForEachToWhile.transformer().transform(
+                    node
+                )
+            )
+        );
     }
 
     private Type transformType(Type type) {
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 84ffdbd..f0f2e00 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -10,6 +10,7 @@ import org.zwobble.couscous.ast.visitors.DynamicNodeMapper;
 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.DesugarSwitchToIfElse;
 import org.zwobble.couscous.types.ScalarType;
 import org.zwobble.couscous.types.Type;
@@ -60,7 +61,8 @@ import static org.zwobble.couscous.util.ExtraLists.*;
 public class PythonCodeGenerator {
     public static PythonModuleNode generateCode(TypeNode sugaredTypeNode) {
         NodeTransformer switchTransformer = DesugarSwitchToIfElse.transformer();
-        TypeNode typeNode = switchTransformer.transformTypeDeclaration(sugaredTypeNode);
+        NodeTransformer forEachTransformer = DesugarForEachToWhile.transformer();
+        TypeNode typeNode = 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 5692f8a..df75ec4 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java
@@ -21,7 +21,6 @@ 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.types.BoundTypeParameter.boundTypeParameter;
 import static org.zwobble.couscous.util.ExtraIterables.takeUntil;
 import static org.zwobble.couscous.util.ExtraLists.*;
 import static org.zwobble.couscous.util.Tails.tail;
@@ -203,17 +202,14 @@ class JavaStatementReader {
         Type elementType = typeOf(parameter.getType());
 
         ExpressionNode iterableValue = readExpression(JavaTypes.iterable(elementType), statement.getExpression());
-        ExpressionNode iteratorValue = methodCall(iterableValue, "iterator", list(), JavaTypes.iterator(elementType));
-        LocalVariableDeclarationNode iterator = scope.temporaryVariable(iteratorValue);
-
-        ExpressionNode hasNext = methodCall(reference(iterator), "hasNext", list(), Types.BOOLEAN);
-        ExpressionNode next = methodCall(reference(iterator), "next", list(), boundTypeParameter(JavaTypes.ITERATOR_TYPE_PARAMETER, elementType));
         IVariableBinding parameterBinding = parameter.resolveBinding();
-        WhileNode loop = whileLoop(hasNext, cons(
-            scope.localVariable(parameterBinding.getKey(), parameterBinding.getName(), elementType, next),
-            readStatement(statement.getBody())));
 
-        return list(iterator, loop);
+        ForEachNode node = new ForEachNode(
+            scope.generateVariable(parameterBinding.getKey(), parameterBinding.getName(), elementType),
+            iterableValue,
+            readStatement(statement.getBody())
+        );
+        return list(node);
     }
 
     private List<StatementNode> readVariableDeclarationStatement(VariableDeclarationStatement statement) {
diff --git a/src/main/java/org/zwobble/couscous/transforms/DesugarForEachToWhile.java b/src/main/java/org/zwobble/couscous/transforms/DesugarForEachToWhile.java
new file mode 100644
index 0000000..ab6f1e5
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/transforms/DesugarForEachToWhile.java
@@ -0,0 +1,49 @@
+package org.zwobble.couscous.transforms;
+
+import org.zwobble.couscous.ast.*;
+import org.zwobble.couscous.ast.visitors.NodeTransformer;
+import org.zwobble.couscous.frontends.java.JavaTypes;
+import org.zwobble.couscous.frontends.java.Scope;
+import org.zwobble.couscous.types.Type;
+import org.zwobble.couscous.types.Types;
+
+import java.util.List;
+import java.util.Optional;
+
+import static org.zwobble.couscous.ast.LocalVariableDeclarationNode.localVariableDeclaration;
+import static org.zwobble.couscous.ast.MethodCallNode.methodCall;
+import static org.zwobble.couscous.ast.VariableReferenceNode.reference;
+import static org.zwobble.couscous.ast.WhileNode.whileLoop;
+import static org.zwobble.couscous.types.BoundTypeParameter.boundTypeParameter;
+import static org.zwobble.couscous.util.ExtraLists.cons;
+import static org.zwobble.couscous.util.ExtraLists.list;
+
+public class DesugarForEachToWhile {
+    public static NodeTransformer transformer() {
+        Scope scope = Scope.create().temporaryPrefix("_couscous_desugar_foreach_to_while");
+        return NodeTransformer.builder()
+            .transformStatement(statement -> {
+                if (statement.type() == NodeTypes.FOR_EACH) {
+                    return Optional.of(desugar(scope, ((ForEachNode) statement)));
+                } else {
+                    return Optional.empty();
+                }
+            })
+            .build();
+    }
+
+    public static List<StatementNode> desugar(Scope scope, ForEachNode forEach) {
+        Type elementType = forEach.getTarget().getType();
+        ExpressionNode iteratorValue = methodCall(forEach.getIterable(), "iterator", list(), JavaTypes.iterator(elementType));
+        LocalVariableDeclarationNode iterator = scope.temporaryVariable(iteratorValue);
+
+        ExpressionNode hasNext = methodCall(reference(iterator), "hasNext", list(), Types.BOOLEAN);
+        ExpressionNode next = methodCall(reference(iterator), "next", list(), boundTypeParameter(JavaTypes.ITERATOR_TYPE_PARAMETER, elementType));
+        WhileNode loop = whileLoop(hasNext, cons(
+            localVariableDeclaration(forEach.getTarget(), next),
+            forEach.getStatements()
+        ));
+
+        return list(iterator, loop);
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/util/ExtraIterables.java b/src/main/java/org/zwobble/couscous/util/ExtraIterables.java
index a40b409..fcc13d2 100644
--- a/src/main/java/org/zwobble/couscous/util/ExtraIterables.java
+++ b/src/main/java/org/zwobble/couscous/util/ExtraIterables.java
@@ -201,4 +201,14 @@ public class ExtraIterables {
             }
         };
     }
+
+    public static <T> T only(Iterable<T> iterable) {
+        Iterator<T> iterator = iterable.iterator();
+        T value = iterator.next();
+        if (iterator.hasNext()) {
+            throw new RuntimeException("iterable had more than one element");
+        } else {
+            return value;
+        }
+    }
 }