Add if-statement node

6 files changed, 127 insertions, 20 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java b/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java
new file mode 100644
index 0000000..ae3e29b
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java
@@ -0,0 +1,90 @@
+package org.zwobble.couscous.ast;
+
+import java.util.List;
+
+import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
+
+public class IfStatementNode implements StatementNode {
+    public static IfStatementNode ifStatement(
+            ExpressionNode condition,
+            List<StatementNode> trueBranch,
+            List<StatementNode> falseBranch) {
+        return new IfStatementNode(condition, trueBranch, falseBranch);
+    }
+    
+    private final ExpressionNode condition;
+    private final List<StatementNode> trueBranch;
+    private final List<StatementNode> falseBranch;
+
+    private IfStatementNode(
+            ExpressionNode condition,
+            List<StatementNode> trueBranch,
+            List<StatementNode> falseBranch) {
+        this.condition = condition;
+        this.trueBranch = trueBranch;
+        this.falseBranch = falseBranch;
+    }
+
+    public ExpressionNode getCondition() {
+        return condition;
+    }
+    
+    public List<StatementNode> getTrueBranch() {
+        return trueBranch;
+    }
+    
+    public List<StatementNode> getFalseBranch() {
+        return falseBranch;
+    }
+    
+    @Override
+    public <T> T accept(StatementNodeMapper<T> visitor) {
+        return visitor.visit(this);
+    }
+
+    @Override
+    public String toString() {
+        return "IfStatementNode(condition=" + condition + ", trueBranch="
+               + trueBranch + ", falseBranch=" + falseBranch + ")";
+    }
+
+    @Override
+    public int hashCode() {
+        final int prime = 31;
+        int result = 1;
+        result = prime * result
+                 + ((condition == null) ? 0 : condition.hashCode());
+        result = prime * result
+                 + ((falseBranch == null) ? 0 : falseBranch.hashCode());
+        result = prime * result
+                 + ((trueBranch == null) ? 0 : trueBranch.hashCode());
+        return result;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj)
+            return true;
+        if (obj == null)
+            return false;
+        if (getClass() != obj.getClass())
+            return false;
+        IfStatementNode other = (IfStatementNode) obj;
+        if (condition == null) {
+            if (other.condition != null)
+                return false;
+        } else if (!condition.equals(other.condition))
+            return false;
+        if (falseBranch == null) {
+            if (other.falseBranch != null)
+                return false;
+        } else if (!falseBranch.equals(other.falseBranch))
+            return false;
+        if (trueBranch == null) {
+            if (other.trueBranch != null)
+                return false;
+        } else if (!trueBranch.equals(other.trueBranch))
+            return false;
+        return true;
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/ast/StatementNode.java b/src/main/java/org/zwobble/couscous/ast/StatementNode.java
index 7a1b25b..1186619 100644
--- a/src/main/java/org/zwobble/couscous/ast/StatementNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/StatementNode.java
@@ -27,6 +27,12 @@ public interface StatementNode extends Node {
                 visitor.visit(localVariableDeclaration);
                 return null;
             }
+
+            @Override
+            public Void visit(IfStatementNode ifStatement) {
+                visitor.visit(ifStatement);
+                return null;
+            }
         });
     }
 }
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);
 }