Implement while loops in backends

8 files changed, 120 insertions, 55 deletions

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 ca1c706..5f6fefb 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -53,6 +53,7 @@ import static org.zwobble.couscous.backends.python.ast.PythonModuleNode.pythonMo
 import static org.zwobble.couscous.backends.python.ast.PythonReturnNode.pythonReturn;
 import static org.zwobble.couscous.backends.python.ast.PythonStringLiteralNode.pythonStringLiteral;
 import static org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNode.pythonVariableReference;
+import static org.zwobble.couscous.backends.python.ast.PythonWhileNode.pythonWhile;
 
 public class PythonCodeGenerator {
     public static PythonModuleNode generateCode(ClassNode classNode) {
@@ -199,7 +200,9 @@ public class PythonCodeGenerator {
 
         @Override
         public PythonStatementNode visit(WhileNode whileLoop) {
-            throw new UnsupportedOperationException();
+            return pythonWhile(
+                generateExpression(whileLoop.getCondition()),
+                generateStatements(whileLoop.getBody()));
         }
     }
     
diff --git a/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java b/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
index 4eea760..6501d1b 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
@@ -2,28 +2,8 @@ package org.zwobble.couscous.backends.python;
 
 import java.util.List;
 import java.util.function.Consumer;
-import org.zwobble.couscous.backends.python.ast.PythonAssignmentNode;
-import org.zwobble.couscous.backends.python.ast.PythonAttributeAccessNode;
-import org.zwobble.couscous.backends.python.ast.PythonBinaryOperation;
-import org.zwobble.couscous.backends.python.ast.PythonBlock;
-import org.zwobble.couscous.backends.python.ast.PythonBooleanLiteralNode;
-import org.zwobble.couscous.backends.python.ast.PythonCallNode;
-import org.zwobble.couscous.backends.python.ast.PythonClassNode;
-import org.zwobble.couscous.backends.python.ast.PythonConditionalExpressionNode;
-import org.zwobble.couscous.backends.python.ast.PythonExpressionNode;
-import org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode;
-import org.zwobble.couscous.backends.python.ast.PythonGetSliceNode;
-import org.zwobble.couscous.backends.python.ast.PythonIfStatementNode;
-import org.zwobble.couscous.backends.python.ast.PythonImportNode;
-import org.zwobble.couscous.backends.python.ast.PythonIntegerLiteralNode;
-import org.zwobble.couscous.backends.python.ast.PythonModuleNode;
-import org.zwobble.couscous.backends.python.ast.PythonNode;
-import org.zwobble.couscous.backends.python.ast.PythonNotNode;
-import org.zwobble.couscous.backends.python.ast.PythonPassNode;
-import org.zwobble.couscous.backends.python.ast.PythonReturnNode;
-import org.zwobble.couscous.backends.python.ast.PythonStatementNode;
-import org.zwobble.couscous.backends.python.ast.PythonStringLiteralNode;
-import org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNode;
+
+import org.zwobble.couscous.backends.python.ast.*;
 import org.zwobble.couscous.backends.python.ast.visitors.PythonNodeVisitor;
 import org.zwobble.couscous.util.Action;
 import static com.google.common.collect.Iterables.skip;
@@ -207,7 +187,17 @@ public class PythonSerializer implements PythonNodeVisitor {
             writeBlock(ifStatement.getFalseBranch()); 
         });
     }
-    
+
+    @Override
+    public void visit(PythonWhileNode whileLoop) {
+        writer.writeStatement(() -> {
+            writer.writeKeyword("while");
+            writer.writeSpace();
+            write(whileLoop.getCondition());
+            writeBlock(whileLoop.getBody());
+        });
+    }
+
     private void writeParenthesised(PythonExpressionNode expression) {
         writer.writeSymbol("(");
         write(expression);
diff --git a/src/main/java/org/zwobble/couscous/backends/python/ast/PythonWhileNode.java b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonWhileNode.java
new file mode 100644
index 0000000..fa373e6
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonWhileNode.java
@@ -0,0 +1,34 @@
+package org.zwobble.couscous.backends.python.ast;
+
+import org.zwobble.couscous.backends.python.ast.visitors.PythonNodeVisitor;
+
+import java.util.List;
+
+public class PythonWhileNode implements PythonStatementNode {
+    public static PythonWhileNode pythonWhile(
+            PythonExpressionNode condition,
+            List<PythonStatementNode> body) {
+        return new PythonWhileNode(condition, new PythonBlock(body));
+    }
+
+    private final PythonExpressionNode condition;
+    private final PythonBlock body;
+
+    private PythonWhileNode(PythonExpressionNode condition, PythonBlock body) {
+        this.condition = condition;
+        this.body = body;
+    }
+
+    public PythonExpressionNode getCondition() {
+        return condition;
+    }
+
+    public PythonBlock getBody() {
+        return body;
+    }
+
+    @Override
+    public void accept(PythonNodeVisitor visitor) {
+        visitor.visit(this);
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/backends/python/ast/visitors/PythonNodeVisitor.java b/src/main/java/org/zwobble/couscous/backends/python/ast/visitors/PythonNodeVisitor.java
index fb8fa71..8f04787 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/ast/visitors/PythonNodeVisitor.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/ast/visitors/PythonNodeVisitor.java
@@ -1,23 +1,6 @@
 package org.zwobble.couscous.backends.python.ast.visitors;
 
-import org.zwobble.couscous.backends.python.ast.PythonAssignmentNode;
-import org.zwobble.couscous.backends.python.ast.PythonAttributeAccessNode;
-import org.zwobble.couscous.backends.python.ast.PythonBinaryOperation;
-import org.zwobble.couscous.backends.python.ast.PythonBooleanLiteralNode;
-import org.zwobble.couscous.backends.python.ast.PythonCallNode;
-import org.zwobble.couscous.backends.python.ast.PythonClassNode;
-import org.zwobble.couscous.backends.python.ast.PythonConditionalExpressionNode;
-import org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode;
-import org.zwobble.couscous.backends.python.ast.PythonGetSliceNode;
-import org.zwobble.couscous.backends.python.ast.PythonIfStatementNode;
-import org.zwobble.couscous.backends.python.ast.PythonImportNode;
-import org.zwobble.couscous.backends.python.ast.PythonIntegerLiteralNode;
-import org.zwobble.couscous.backends.python.ast.PythonModuleNode;
-import org.zwobble.couscous.backends.python.ast.PythonNotNode;
-import org.zwobble.couscous.backends.python.ast.PythonPassNode;
-import org.zwobble.couscous.backends.python.ast.PythonReturnNode;
-import org.zwobble.couscous.backends.python.ast.PythonStringLiteralNode;
-import org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNode;
+import org.zwobble.couscous.backends.python.ast.*;
 
 public interface PythonNodeVisitor {
     void visit(PythonIntegerLiteralNode integerLiteral);
@@ -38,6 +21,7 @@ public interface PythonNodeVisitor {
     void visit(PythonAssignmentNode assignment);
     void visit(PythonImportNode importNode);
     void visit(PythonIfStatementNode ifStatement);
+    void visit(PythonWhileNode whileLoop);
 
     void visit(PythonModuleNode module);
 
diff --git a/src/main/java/org/zwobble/couscous/interpreter/Executor.java b/src/main/java/org/zwobble/couscous/interpreter/Executor.java
index 1d9a327..8909acc 100644
--- a/src/main/java/org/zwobble/couscous/interpreter/Executor.java
+++ b/src/main/java/org/zwobble/couscous/interpreter/Executor.java
@@ -23,7 +23,7 @@ public class Executor implements StatementNodeMapper<Optional<InterpreterValue>>
         }
         return UnitInterpreterValue.UNIT;
     }
-    
+
     private static Environment buildEnvironment(final Environment environment, final CallableNode method, Optional<InterpreterValue> thisValue, PositionalArguments arguments) {
         final org.zwobble.couscous.interpreter.StackFrameBuilder stackFrame = new StackFrameBuilder();
         for (int index = 0; index < method.getArguments().size(); index++) {
@@ -32,19 +32,19 @@ public class Executor implements StatementNodeMapper<Optional<InterpreterValue>>
         findDeclarations(method.getBody()).forEach(declaration -> stackFrame.declare(declaration));
         return environment.withStackFrame(thisValue, stackFrame.build());
     }
-    
+
     private static Stream<VariableNode> findDeclarations(List<StatementNode> body) {
         return body.stream().flatMap(statement -> statement.accept(new StatementNodeMapper<Stream<VariableNode>>(){
             @Override
             public Stream<VariableNode> visit(ReturnNode returnNode) {
                 return Stream.empty();
             }
-            
+
             @Override
             public Stream<VariableNode> visit(ExpressionStatementNode expressionStatement) {
                 return Stream.empty();
             }
-            
+
             @Override
             public Stream<VariableNode> visit(LocalVariableDeclarationNode localVariableDeclaration) {
                 return Stream.of(localVariableDeclaration);
@@ -63,27 +63,27 @@ public class Executor implements StatementNodeMapper<Optional<InterpreterValue>>
             }
         }));
     }
-    
+
     public static Optional<InterpreterValue> exec(Environment environment, StatementNode statement) {
         return statement.accept(new Executor(environment));
     }
     private final Environment environment;
-    
+
     private Executor(Environment environment) {
         this.environment = environment;
     }
-    
+
     @Override
     public Optional<InterpreterValue> visit(ReturnNode returnNode) {
         return Optional.of(eval(environment, returnNode.getValue()));
     }
-    
+
     @Override
     public Optional<InterpreterValue> visit(ExpressionStatementNode expressionStatement) {
         eval(environment, expressionStatement.getExpression());
         return Optional.empty();
     }
-    
+
     @Override
     public Optional<InterpreterValue> visit(LocalVariableDeclarationNode localVariableDeclaration) {
         final org.zwobble.couscous.interpreter.values.InterpreterValue value = eval(environment, localVariableDeclaration.getInitialValue());
@@ -93,8 +93,7 @@ public class Executor implements StatementNodeMapper<Optional<InterpreterValue>>
 
     @Override
     public Optional<InterpreterValue> visit(IfStatementNode ifStatement) {
-        BooleanInterpreterValue condition = (BooleanInterpreterValue) eval(environment, ifStatement.getCondition());
-        List<StatementNode> body = condition.getValue()
+        List<StatementNode> body = evalCondition(ifStatement.getCondition())
             ? ifStatement.getTrueBranch()
             : ifStatement.getFalseBranch();
         return exec(body);
@@ -102,7 +101,18 @@ public class Executor implements StatementNodeMapper<Optional<InterpreterValue>>
 
     @Override
     public Optional<InterpreterValue> visit(WhileNode whileLoop) {
-        throw new UnsupportedOperationException();
+        while (evalCondition(whileLoop.getCondition())) {
+            Optional<InterpreterValue> result = exec(whileLoop.getBody());
+            if (result.isPresent()) {
+                return result;
+            }
+        }
+        return Optional.empty();
+    }
+
+    private boolean evalCondition(ExpressionNode conditionExpression) {
+        BooleanInterpreterValue condition = (BooleanInterpreterValue) eval(environment, conditionExpression);
+        return condition.getValue();
     }
 
     private Optional<InterpreterValue> exec(List<StatementNode> statements) {
diff --git a/src/test/java/org/zwobble/couscous/tests/BackendMethodTests.java b/src/test/java/org/zwobble/couscous/tests/BackendMethodTests.java
index 758832f..a325dcb 100644
--- a/src/test/java/org/zwobble/couscous/tests/BackendMethodTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/BackendMethodTests.java
@@ -6,6 +6,7 @@ import org.zwobble.couscous.ast.FormalArgumentNode;
 import org.zwobble.couscous.ast.LiteralNode;
 import org.zwobble.couscous.ast.LocalVariableDeclarationNode;
 import org.zwobble.couscous.ast.MethodNode;
+import org.zwobble.couscous.values.IntegerValue;
 import org.zwobble.couscous.values.PrimitiveValue;
 import org.zwobble.couscous.values.PrimitiveValues;
 import org.zwobble.couscous.values.StringValue;
@@ -13,15 +14,19 @@ import org.zwobble.couscous.values.StringValue;
 import static java.util.Arrays.asList;
 import static org.junit.Assert.assertEquals;
 import static org.zwobble.couscous.ast.AssignmentNode.assign;
+import static org.zwobble.couscous.ast.AssignmentNode.assignStatement;
 import static org.zwobble.couscous.ast.ExpressionStatementNode.expressionStatement;
 import static org.zwobble.couscous.ast.FormalArgumentNode.formalArg;
 import static org.zwobble.couscous.ast.IfStatementNode.ifStatement;
 import static org.zwobble.couscous.ast.LiteralNode.literal;
 import static org.zwobble.couscous.ast.LocalVariableDeclarationNode.localVariableDeclaration;
+import static org.zwobble.couscous.ast.MethodCallNode.integerAdd;
+import static org.zwobble.couscous.ast.MethodCallNode.notEqual;
 import static org.zwobble.couscous.ast.MethodNode.staticMethod;
 import static org.zwobble.couscous.ast.ReturnNode.returns;
 import static org.zwobble.couscous.ast.VariableDeclaration.var;
 import static org.zwobble.couscous.ast.VariableReferenceNode.reference;
+import static org.zwobble.couscous.ast.WhileNode.whileLoop;
 import static org.zwobble.couscous.tests.TestIds.ANY_ID;
 import static org.zwobble.couscous.values.PrimitiveValues.value;
 
@@ -106,6 +111,36 @@ public abstract class BackendMethodTests {
                 asList(returns(literal("[false]")))));
         assertEquals(value("[false]"), runMethod(method));
     }
+
+    @Test
+    public void whileLoopIsExecutedWhileConditionIsTrue() {
+        LocalVariableDeclarationNode x = localVariableDeclaration("x", "x", IntegerValue.REF, literal(0));
+        LocalVariableDeclarationNode y = localVariableDeclaration("y", "y", IntegerValue.REF, literal(2));
+        MethodNode.Builder method = staticMethod("hello")
+            .statement(x)
+            .statement(y)
+            .statement(whileLoop(
+                notEqual(reference(x), reference(y)),
+                asList(assignStatement(x, integerAdd(reference(x), literal(1))))))
+            .statement(returns(reference(x)));
+        assertEquals(value(2), runMethod(method));
+    }
+
+    @Test
+    public void returnWillExitWhileLoopEarly() {
+        LocalVariableDeclarationNode x = localVariableDeclaration("x", "x", IntegerValue.REF, literal(0));
+        LocalVariableDeclarationNode y = localVariableDeclaration("y", "y", IntegerValue.REF, literal(2));
+        MethodNode.Builder method = staticMethod("hello")
+            .statement(x)
+            .statement(y)
+            .statement(whileLoop(
+                notEqual(reference(x), reference(y)),
+                asList(
+                    assignStatement(x, integerAdd(reference(x), literal(1))),
+                    returns(reference(x)))))
+            .statement(returns(reference(x)));
+        assertEquals(value(1), runMethod(method));
+    }
     
     protected PrimitiveValue runMethod(MethodNode.Builder methodBuilder, PrimitiveValue... arguments) {
         MethodNode method = methodBuilder.build();
diff --git a/src/test/java/org/zwobble/couscous/tests/backends/python/PythonSerializerTests.java b/src/test/java/org/zwobble/couscous/tests/backends/python/PythonSerializerTests.java
index c3782aa..9271d85 100644
--- a/src/test/java/org/zwobble/couscous/tests/backends/python/PythonSerializerTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/backends/python/PythonSerializerTests.java
@@ -25,6 +25,7 @@ import static org.zwobble.couscous.backends.python.ast.PythonPassNode.PASS;
 import static org.zwobble.couscous.backends.python.ast.PythonReturnNode.pythonReturn;
 import static org.zwobble.couscous.backends.python.ast.PythonStringLiteralNode.pythonStringLiteral;
 import static org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNode.pythonVariableReference;
+import static org.zwobble.couscous.backends.python.ast.PythonWhileNode.pythonWhile;
 
 public class PythonSerializerTests {
     @Test
@@ -162,6 +163,14 @@ public class PythonSerializerTests {
             asList(pythonReturn(pythonIntegerLiteral(2)))));
         assertEquals("if True:\n    return 1\nelse:\n    return 2\n", output);
     }
+
+    @Test
+    public void whileIsSerializedWithWhileKeyword() {
+        String output = serialize(pythonWhile(
+            pythonBooleanLiteral(true),
+            asList(pythonReturn(pythonIntegerLiteral(1)))));
+        assertEquals("while True:\n    return 1\n", output);
+    }
     
     @Test
     public void moduleIsSerializedStatements() {
diff --git a/src/test/resources/java/while-factorial/com/example/WhileFactorial.java b/src/test/resources/java/while-factorial/com/example/WhileFactorial.java
index 89e41cc..04bc8b6 100644
--- a/src/test/resources/java/while-factorial/com/example/WhileFactorial.java
+++ b/src/test/resources/java/while-factorial/com/example/WhileFactorial.java
@@ -5,7 +5,7 @@ public class WhileFactorial {
         int result = 1;
         while (n > 1) {
             result *= n;
-            n--;
+            --n;
         }
         return result;
     }