Desugar foreach to for rather than while

3 files changed, 16 insertions, 12 deletions

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 28d2a3e..2245782 100644
--- a/src/main/java/org/zwobble/couscous/backends/csharp/CsharpCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/csharp/CsharpCodeGenerator.java
@@ -5,7 +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.transforms.DesugarForEachToFor;
 import org.zwobble.couscous.transforms.DesugarForToWhile;
 import org.zwobble.couscous.types.*;
 
@@ -47,7 +47,7 @@ public class CsharpCodeGenerator {
         return nodeTransformer.transform(
             CsharpPrimitiveMethods.TRANSFORMER.transform(
                 DesugarForToWhile.transformer().transform(
-                    DesugarForEachToWhile.transformer().transform(
+                    DesugarForEachToFor.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 29e7274..2190f66 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -10,7 +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.DesugarForEachToFor;
 import org.zwobble.couscous.transforms.DesugarForToWhile;
 import org.zwobble.couscous.transforms.DesugarSwitchToIfElse;
 import org.zwobble.couscous.types.ScalarType;
@@ -63,7 +63,7 @@ public class PythonCodeGenerator {
     public static PythonModuleNode generateCode(TypeNode sugaredTypeNode) {
         NodeTransformer switchTransformer = DesugarSwitchToIfElse.transformer();
         NodeTransformer forTransformer = DesugarForToWhile.transformer();
-        NodeTransformer forEachTransformer = DesugarForEachToWhile.transformer();
+        NodeTransformer forEachTransformer = DesugarForEachToFor.transformer();
         TypeNode typeNode = forTransformer.transformTypeDeclaration(
             forEachTransformer.transformTypeDeclaration(
                 switchTransformer.transformTypeDeclaration(sugaredTypeNode)
diff --git a/src/main/java/org/zwobble/couscous/transforms/DesugarForEachToWhile.java b/src/main/java/org/zwobble/couscous/transforms/DesugarForEachToFor.java
index 7b6147b..bdae29a 100644
--- a/src/main/java/org/zwobble/couscous/transforms/DesugarForEachToWhile.java
+++ b/src/main/java/org/zwobble/couscous/transforms/DesugarForEachToFor.java
@@ -13,14 +13,13 @@ 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 class DesugarForEachToFor {
     public static NodeTransformer transformer() {
-        Scope scope = Scope.create().temporaryPrefix("_couscous_desugar_foreach_to_while");
+        Scope scope = Scope.create().temporaryPrefix("_couscous_desugar_foreach_to_for");
         return NodeTransformer.builder()
             .transformStatement(statement -> {
                 if (statement.type() == NodeTypes.FOR_EACH) {
@@ -39,11 +38,16 @@ public class DesugarForEachToWhile {
 
         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()
-        ));
+        ForNode loop = new ForNode(
+            list(iterator),
+            hasNext,
+            list(),
+            cons(
+                localVariableDeclaration(forEach.getTarget(), next),
+                forEach.getStatements()
+            )
+        );
 
-        return list(new StatementBlockNode(list(iterator, loop)));
+        return list(loop);
     }
 }