Read statement blocks

4 files changed, 79 insertions, 14 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/NodeTypes.java b/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
index f7f8bf6..82090aa 100644
--- a/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
+++ b/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
@@ -45,6 +45,7 @@ public class NodeTypes {
     public static final int VARIABLE_REFERENCE = 10029;
     public static final int WHILE = 10030;
     public static final int FOR_EACH = 10034;
+    public static final int STATEMENT_BLOCK = 10035;
 
     private NodeTypes() {
     }
@@ -85,6 +86,7 @@ public class NodeTypes {
         .put(VariableReferenceNode.class, VARIABLE_REFERENCE)
         .put(WhileNode.class, WHILE)
         .put(ForEachNode.class, FOR_EACH)
+        .put(StatementBlockNode.class, STATEMENT_BLOCK)
         .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/ast/StatementBlockNode.java b/src/main/java/org/zwobble/couscous/ast/StatementBlockNode.java
new file mode 100644
index 0000000..0c24295
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/ast/StatementBlockNode.java
@@ -0,0 +1,51 @@
+package org.zwobble.couscous.ast;
+
+import org.zwobble.couscous.ast.visitors.NodeTransformer;
+
+import java.util.List;
+
+public class StatementBlockNode implements StatementNode {
+    private final List<StatementNode> statements;
+
+    public StatementBlockNode(List<StatementNode> statements) {
+        this.statements = statements;
+    }
+
+    @Override
+    public int type() {
+        return NodeTypes.STATEMENT_BLOCK;
+    }
+
+    @Override
+    public Iterable<? extends Node> childNodes() {
+        return statements;
+    }
+
+    @Override
+    public StatementNode transformSubtree(NodeTransformer transformer) {
+        return new StatementBlockNode(transformer.transformStatements(statements));
+    }
+
+    @Override
+    public String toString() {
+        return "StatementBlockNode(" +
+            "statements=" + statements +
+            ')';
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        StatementBlockNode that = (StatementBlockNode) o;
+
+        return statements.equals(that.statements);
+
+    }
+
+    @Override
+    public int hashCode() {
+        return statements.hashCode();
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
index 30adff9..1c4ac81 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
@@ -367,7 +367,7 @@ public class JavaReader {
             tryCast(Initializer.class, declaration)
                 .ifPresent(initializer -> body.addInitializer(
                     Modifier.isStatic(initializer.getModifiers()),
-                    readStatement(scope, initializer.getBody())));
+                    readStatementBody(scope, initializer.getBody())));
 
             tryCast(FieldDeclaration.class, declaration)
                 .ifPresent(field -> readField(body, scope, type, field));
@@ -494,6 +494,11 @@ public class JavaReader {
         }
     }
 
+    List<StatementNode> readStatementBody(Scope scope, Statement statement) {
+        JavaStatementReader statementReader = new JavaStatementReader(scope, expressionReader(scope), Optional.empty());
+        return statementReader.readBody(statement);
+    }
+
     List<StatementNode> readStatement(Scope scope, Statement statement) {
         JavaStatementReader statementReader = new JavaStatementReader(scope, expressionReader(scope), Optional.empty());
         return statementReader.readStatement(statement);
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 df75ec4..72f94d5 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java
@@ -10,7 +10,6 @@ import org.zwobble.couscous.types.Types;
 
 import java.util.List;
 import java.util.Optional;
-import java.util.stream.Collectors;
 
 import static org.zwobble.couscous.ast.ExceptionHandlerNode.exceptionHandler;
 import static org.zwobble.couscous.ast.ExpressionStatementNode.expressionStatement;
@@ -83,10 +82,9 @@ class JavaStatementReader {
 
     private List<StatementNode> readBlock(Block block) {
         @SuppressWarnings("unchecked")
-        List<Statement> statements = block.statements();
-        return statements.stream()
-            .flatMap(statement -> readStatement(statement).stream())
-            .collect(Collectors.toList());
+        List<Statement> javaStatements = block.statements();
+        List<StatementNode> statements = eagerFlatMap(javaStatements, this::readStatement);
+        return list(new StatementBlockNode(statements));
     }
 
     private StatementNode readReturnStatement(ReturnStatement statement) {
@@ -107,10 +105,10 @@ class JavaStatementReader {
     private StatementNode readIfStatement(IfStatement statement) {
         List<StatementNode> falseBranch = statement.getElseStatement() == null
             ? list()
-            : readStatement(statement.getElseStatement());
+            : readBody(statement.getElseStatement());
         return ifStatement(
             readExpression(Types.BOOLEAN, statement.getExpression()),
-            readStatement(statement.getThenStatement()),
+            readBody(statement.getThenStatement()),
             falseBranch);
     }
 
@@ -145,7 +143,7 @@ class JavaStatementReader {
         @SuppressWarnings("unchecked")
         List<VariableDeclarationExpression> resources = statement.resources();
         List<LocalVariableDeclarationNode> resourceDeclarations = eagerFlatMap(resources, this::readVariableDeclarationExpression);
-        List<StatementNode> body = readStatement(statement.getBody());
+        List<StatementNode> body = readBody(statement.getBody());
         for (LocalVariableDeclarationNode resource : Lists.reverse(resourceDeclarations)) {
             body = list(
                 resource,
@@ -155,7 +153,7 @@ class JavaStatementReader {
         @SuppressWarnings("unchecked")
         List<CatchClause> javaCatchClauses = statement.catchClauses();
         List<ExceptionHandlerNode> catchClauses = eagerMap(javaCatchClauses, this::readCatchClause);
-        List<StatementNode> finallyBody = statement.getFinally() == null ? list() : readStatement(statement.getFinally());
+        List<StatementNode> finallyBody = statement.getFinally() == null ? list() : readBody(statement.getFinally());
         if (catchClauses.isEmpty() && finallyBody.isEmpty()) {
             return body;
         } else {
@@ -166,7 +164,7 @@ class JavaStatementReader {
     private ExceptionHandlerNode readCatchClause(CatchClause clause) {
         return exceptionHandler(
             JavaVariableDeclarationReader.read(scope, clause.getException()),
-            readStatement(clause.getBody()));
+            readBody(clause.getBody()));
     }
 
     private List<StatementNode> readStatements(Iterable<Statement> statements) {
@@ -176,7 +174,7 @@ class JavaStatementReader {
     private WhileNode readWhileStatement(WhileStatement statement) {
         return whileLoop(
             readExpression(Types.BOOLEAN, statement.getExpression()),
-            readStatement(statement.getBody()));
+            readBody(statement.getBody()));
     }
 
     private List<StatementNode> readForStatement(ForStatement statement) {
@@ -193,7 +191,7 @@ class JavaStatementReader {
             whileLoop(
                 readExpression(Types.BOOLEAN, statement.getExpression()),
                 concat(
-                    readStatement(statement.getBody()),
+                    readBody(statement.getBody()),
                     eagerMap(updaters, updater -> expressionStatement(readExpressionWithoutBoxing(updater))))));
     }
 
@@ -207,7 +205,7 @@ class JavaStatementReader {
         ForEachNode node = new ForEachNode(
             scope.generateVariable(parameterBinding.getKey(), parameterBinding.getName(), elementType),
             iterableValue,
-            readStatement(statement.getBody())
+            readBody(statement.getBody())
         );
         return list(node);
     }
@@ -243,4 +241,13 @@ class JavaStatementReader {
     private ExpressionNode readExpression(Type targetType, Expression expression) {
         return expressionReader.readExpression(targetType, expression);
     }
+
+    List<StatementNode> readBody(Statement statement) {
+        if (statement instanceof Block) {
+            List<Statement> javaStatements = ((Block) statement).statements();
+            return eagerFlatMap(javaStatements, this::readStatement);
+        } else {
+            return readStatement(statement);
+        }
+    }
 }