Add if-statement node

4 files changed, 31 insertions, 20 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitor.java b/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitor.java
index 8a88052..ce6bb52 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitor.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitor.java
@@ -6,6 +6,7 @@ import org.zwobble.couscous.ast.ConstructorCallNode;
 import org.zwobble.couscous.ast.ConstructorNode;
 import org.zwobble.couscous.ast.ExpressionStatementNode;
 import org.zwobble.couscous.ast.FieldAccessNode;
+import org.zwobble.couscous.ast.IfStatementNode;
 import org.zwobble.couscous.ast.LiteralNode;
 import org.zwobble.couscous.ast.LocalVariableDeclarationNode;
 import org.zwobble.couscous.ast.MethodCallNode;
@@ -30,6 +31,7 @@ public interface NodeVisitor {
     void visit(ReturnNode returnNode);
     void visit(ExpressionStatementNode expressionStatement);
     void visit(LocalVariableDeclarationNode localVariableDeclaration);
+    void visit(IfStatementNode ifStatement);
     
     void visit(ClassNode classNode);
     void visit(MethodNode methodNode);
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitorWithEmptyDefaults.java b/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitorWithEmptyDefaults.java
index 8408b7a..fe57a28 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitorWithEmptyDefaults.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitorWithEmptyDefaults.java
@@ -6,6 +6,7 @@ import org.zwobble.couscous.ast.ConstructorCallNode;
 import org.zwobble.couscous.ast.ConstructorNode;
 import org.zwobble.couscous.ast.ExpressionStatementNode;
 import org.zwobble.couscous.ast.FieldAccessNode;
+import org.zwobble.couscous.ast.IfStatementNode;
 import org.zwobble.couscous.ast.LiteralNode;
 import org.zwobble.couscous.ast.LocalVariableDeclarationNode;
 import org.zwobble.couscous.ast.MethodCallNode;
@@ -66,6 +67,10 @@ public interface NodeVisitorWithEmptyDefaults extends NodeVisitor {
     }
     
     @Override
+    default void visit(IfStatementNode ifStatement) {
+    }
+    
+    @Override
     default void visit(ClassNode classNode) {
     }
     
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitors.java b/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitors.java
index 9424e39..1bcc226 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitors.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitors.java
@@ -1,19 +1,20 @@
 package org.zwobble.couscous.ast.visitors;
 
+import java.util.List;
+
 import org.zwobble.couscous.ast.AssignmentNode;
 import org.zwobble.couscous.ast.ClassNode;
 import org.zwobble.couscous.ast.ConstructorCallNode;
 import org.zwobble.couscous.ast.ConstructorNode;
-import org.zwobble.couscous.ast.ExpressionNode;
 import org.zwobble.couscous.ast.ExpressionStatementNode;
 import org.zwobble.couscous.ast.FieldAccessNode;
+import org.zwobble.couscous.ast.IfStatementNode;
 import org.zwobble.couscous.ast.LiteralNode;
 import org.zwobble.couscous.ast.LocalVariableDeclarationNode;
 import org.zwobble.couscous.ast.MethodCallNode;
 import org.zwobble.couscous.ast.MethodNode;
 import org.zwobble.couscous.ast.Node;
 import org.zwobble.couscous.ast.ReturnNode;
-import org.zwobble.couscous.ast.StatementNode;
 import org.zwobble.couscous.ast.StaticMethodCallNode;
 import org.zwobble.couscous.ast.TernaryConditionalNode;
 import org.zwobble.couscous.ast.ThisReferenceNode;
@@ -56,25 +57,19 @@ public class NodeVisitors {
             public void visit(MethodCallNode methodCall) {
                 visitor.visit(methodCall);
                 methodCall.getReceiver().accept(this);
-                for (ExpressionNode argument : methodCall.getArguments()) {
-                    argument.accept(this);
-                }
+                visitAll(methodCall.getArguments());
             }
             
             @Override
             public void visit(StaticMethodCallNode staticMethodCall) {
                 visitor.visit(staticMethodCall);
-                for (ExpressionNode argument : staticMethodCall.getArguments()) {
-                    argument.accept(this);
-                }
+                visitAll(staticMethodCall.getArguments());
             }
             
             @Override
             public void visit(ConstructorCallNode call) {
                 visitor.visit(call);
-                for (ExpressionNode argument : call.getArguments()) {
-                    argument.accept(this);
-                }
+                visitAll(call.getArguments());
             }
             
             @Override
@@ -99,27 +94,34 @@ public class NodeVisitors {
                 visitor.visit(localVariableDeclaration);
                 localVariableDeclaration.getInitialValue().accept(this);
             }
-            
+
+            @Override
+            public void visit(IfStatementNode ifStatement) {
+                ifStatement.getCondition().accept(visitor);
+                visitAll(ifStatement.getTrueBranch());
+                visitAll(ifStatement.getFalseBranch());
+            }
+
             @Override
             public void visit(ClassNode classNode) {
                 visitor.visit(classNode);
                 classNode.getConstructor().accept(this);
-                for (MethodNode method : classNode.getMethods()) {
-                    method.accept(this);
-                }
+                visitAll(classNode.getMethods());
             }
             
             @Override
             public void visit(MethodNode methodNode) {
-                for (StatementNode statement : methodNode.getBody()) {
-                    statement.accept(this);
-                }
+                visitAll(methodNode.getBody());
             }
             
             @Override
             public void visit(ConstructorNode constructorNode) {
-                for (StatementNode statement : constructorNode.getBody()) {
-                    statement.accept(this);
+                visitAll(constructorNode.getBody());
+            }
+            
+            private <T extends Node> void visitAll(List<T> nodes) {
+                for (Node node : nodes) {
+                    node.accept(this);
                 }
             }
         });
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/StatementNodeMapper.java b/src/main/java/org/zwobble/couscous/ast/visitors/StatementNodeMapper.java
index 4bae481..fe639bd 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/StatementNodeMapper.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/StatementNodeMapper.java
@@ -1,6 +1,7 @@
 package org.zwobble.couscous.ast.visitors;
 
 import org.zwobble.couscous.ast.ExpressionStatementNode;
+import org.zwobble.couscous.ast.IfStatementNode;
 import org.zwobble.couscous.ast.LocalVariableDeclarationNode;
 import org.zwobble.couscous.ast.ReturnNode;
 
@@ -8,4 +9,5 @@ public interface StatementNodeMapper<T> {
     T visit(ReturnNode returnNode);
     T visit(ExpressionStatementNode expressionStatement);
     T visit(LocalVariableDeclarationNode localVariableDeclaration);
+    T visit(IfStatementNode ifStatement);
 }