Allow descendants of nodes to be traversed

18 files changed, 531 insertions, 335 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/AnnotationNode.java b/src/main/java/org/zwobble/couscous/ast/AnnotationNode.java
index f7a646b..0d135ba 100644
--- a/src/main/java/org/zwobble/couscous/ast/AnnotationNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/AnnotationNode.java
@@ -1,6 +1,8 @@
 package org.zwobble.couscous.ast;
 
-public class AnnotationNode {
+import org.zwobble.couscous.ast.visitors.NodeMapper;
+
+public class AnnotationNode implements Node {
     public static AnnotationNode annotation(TypeName type) {
         return new AnnotationNode(type);
     }
@@ -16,6 +18,11 @@ public class AnnotationNode {
     }
 
     @Override
+    public <T> T accept(NodeMapper<T> visitor) {
+        return visitor.visit(this);
+    }
+
+    @Override
     public String toString() {
         return "AnnotationNode(type=" + type + ")";
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/ClassNode.java b/src/main/java/org/zwobble/couscous/ast/ClassNode.java
index d2c64f5..c252bb3 100644
--- a/src/main/java/org/zwobble/couscous/ast/ClassNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ClassNode.java
@@ -1,8 +1,8 @@
 package org.zwobble.couscous.ast;
 
-import java.util.List;
+import org.zwobble.couscous.ast.visitors.NodeMapper;
 
-import org.zwobble.couscous.ast.visitors.NodeVisitor;
+import java.util.List;
 
 public class ClassNode implements Node {
     public static ClassNodeBuilder builder(String name) {
@@ -54,8 +54,8 @@ public class ClassNode implements Node {
     }
 
     @Override
-    public void accept(NodeVisitor visitor) {
-        visitor.visit(this);
+    public <T> T accept(NodeMapper<T> visitor) {
+        return visitor.visit(this);
     }
 
     @Override
diff --git a/src/main/java/org/zwobble/couscous/ast/ConstructorNode.java b/src/main/java/org/zwobble/couscous/ast/ConstructorNode.java
index c94554b..2e8b85a 100644
--- a/src/main/java/org/zwobble/couscous/ast/ConstructorNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ConstructorNode.java
@@ -2,6 +2,7 @@ package org.zwobble.couscous.ast;
 
 import java.util.List;
 
+import org.zwobble.couscous.ast.visitors.NodeMapper;
 import org.zwobble.couscous.ast.visitors.NodeVisitor;
 
 public class ConstructorNode implements CallableNode {
@@ -28,10 +29,10 @@ public class ConstructorNode implements CallableNode {
     public List<StatementNode> getBody() {
         return body;
     }
-    
+
     @Override
-    public void accept(NodeVisitor visitor) {
-        visitor.visit(this);
+    public <T> T accept(NodeMapper<T> visitor) {
+        return visitor.visit(this);
     }
 
     @Override
diff --git a/src/main/java/org/zwobble/couscous/ast/ExpressionNode.java b/src/main/java/org/zwobble/couscous/ast/ExpressionNode.java
index 8210ae1..65757ad 100644
--- a/src/main/java/org/zwobble/couscous/ast/ExpressionNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ExpressionNode.java
@@ -1,67 +1,58 @@
 package org.zwobble.couscous.ast;
 
 import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
-import org.zwobble.couscous.ast.visitors.NodeVisitor;
+import org.zwobble.couscous.ast.visitors.NodeMapper;
 
 public interface ExpressionNode extends Node {
     TypeName getType();
     <T> T accept(ExpressionNodeMapper<T> visitor);
     
     @Override
-    default void accept(NodeVisitor visitor) {
-        accept(new ExpressionNodeMapper<Void>() {
+    default <T> T accept(NodeMapper<T> visitor) {
+        return accept(new ExpressionNodeMapper<T>() {
             @Override
-            public Void visit(LiteralNode literal) {
-                visitor.visit(literal);
-                return null;
+            public T visit(LiteralNode literal) {
+                return visitor.visit(literal);
             }
 
             @Override
-            public Void visit(VariableReferenceNode variableReference) {
-                visitor.visit(variableReference);
-                return null;
+            public T visit(VariableReferenceNode variableReference) {
+                return visitor.visit(variableReference);
             }
 
             @Override
-            public Void visit(ThisReferenceNode reference) {
-                visitor.visit(reference);
-                return null;
+            public T visit(ThisReferenceNode reference) {
+                return visitor.visit(reference);
             }
 
             @Override
-            public Void visit(AssignmentNode assignment) {
-                visitor.visit(assignment);
-                return null;
+            public T visit(AssignmentNode assignment) {
+                return visitor.visit(assignment);
             }
 
             @Override
-            public Void visit(TernaryConditionalNode ternaryConditional) {
-                visitor.visit(ternaryConditional);
-                return null;
+            public T visit(TernaryConditionalNode ternaryConditional) {
+                return visitor.visit(ternaryConditional);
             }
 
             @Override
-            public Void visit(MethodCallNode methodCall) {
-                visitor.visit(methodCall);
-                return null;
+            public T visit(MethodCallNode methodCall) {
+                return visitor.visit(methodCall);
             }
 
             @Override
-            public Void visit(StaticMethodCallNode staticMethodCall) {
-                visitor.visit(staticMethodCall);
-                return null;
+            public T visit(StaticMethodCallNode staticMethodCall) {
+                return visitor.visit(staticMethodCall);
             }
 
             @Override
-            public Void visit(ConstructorCallNode call) {
-                visitor.visit(call);
-                return null;
+            public T visit(ConstructorCallNode call) {
+                return visitor.visit(call);
             }
 
             @Override
-            public Void visit(FieldAccessNode fieldAccess) {
-                visitor.visit(fieldAccess);
-                return null;
+            public T visit(FieldAccessNode fieldAccess) {
+                return visitor.visit(fieldAccess);
             }
         });
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/FieldDeclarationNode.java b/src/main/java/org/zwobble/couscous/ast/FieldDeclarationNode.java
index 58b05f2..7fea43c 100644
--- a/src/main/java/org/zwobble/couscous/ast/FieldDeclarationNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/FieldDeclarationNode.java
@@ -1,6 +1,9 @@
 package org.zwobble.couscous.ast;
 
-public class FieldDeclarationNode {
+import org.zwobble.couscous.ast.visitors.NodeMapper;
+import org.zwobble.couscous.ast.visitors.NodeVisitor;
+
+public class FieldDeclarationNode implements Node {
     public static FieldDeclarationNode field(String name, TypeName type) {
         return new FieldDeclarationNode(name, type);
     }
@@ -22,6 +25,11 @@ public class FieldDeclarationNode {
     }
 
     @Override
+    public <T> T accept(NodeMapper<T> visitor) {
+        return visitor.visit(this);
+    }
+
+    @Override
     public String toString() {
         return "FieldDeclarationNode(name=" + name + ", type=" + type + ")";
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/FormalArgumentNode.java b/src/main/java/org/zwobble/couscous/ast/FormalArgumentNode.java
index 80c3273..41fb982 100644
--- a/src/main/java/org/zwobble/couscous/ast/FormalArgumentNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/FormalArgumentNode.java
@@ -1,6 +1,8 @@
 package org.zwobble.couscous.ast;
 
-public class FormalArgumentNode implements VariableNode {
+import org.zwobble.couscous.ast.visitors.NodeMapper;
+
+public class FormalArgumentNode implements VariableNode, Node {
     public static FormalArgumentNode formalArg(VariableDeclaration declaration) {
         return new FormalArgumentNode(declaration);
     }
@@ -24,6 +26,11 @@ public class FormalArgumentNode implements VariableNode {
     }
 
     @Override
+    public <T> T accept(NodeMapper<T> visitor) {
+        return visitor.visit(this);
+    }
+
+    @Override
     public String toString() {
         return "FormalArgumentNode(declaration=" + declaration + ")";
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/MethodNode.java b/src/main/java/org/zwobble/couscous/ast/MethodNode.java
index d37f7d7..3e97a36 100644
--- a/src/main/java/org/zwobble/couscous/ast/MethodNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/MethodNode.java
@@ -1,10 +1,9 @@
 package org.zwobble.couscous.ast;
 
-import java.util.List;
-
-import org.zwobble.couscous.ast.visitors.NodeVisitor;
-
 import com.google.common.collect.ImmutableList;
+import org.zwobble.couscous.ast.visitors.NodeMapper;
+
+import java.util.List;
 
 public class MethodNode implements CallableNode {
     public static Builder staticMethod(String name) {
@@ -115,10 +114,10 @@ public class MethodNode implements CallableNode {
     public List<StatementNode> getBody() {
         return body;
     }
-    
+
     @Override
-    public void accept(NodeVisitor visitor) {
-        visitor.visit(this);
+    public <T> T accept(NodeMapper<T> visitor) {
+        return visitor.visit(this);
     }
     
     @Override
diff --git a/src/main/java/org/zwobble/couscous/ast/Node.java b/src/main/java/org/zwobble/couscous/ast/Node.java
index acde931..7c4d347 100644
--- a/src/main/java/org/zwobble/couscous/ast/Node.java
+++ b/src/main/java/org/zwobble/couscous/ast/Node.java
@@ -1,7 +1,131 @@
 package org.zwobble.couscous.ast;
 
+import org.zwobble.couscous.ast.visitors.NodeMapper;
 import org.zwobble.couscous.ast.visitors.NodeVisitor;
 
 public interface Node {
-    void accept(NodeVisitor visitor);
+    default void accept(NodeVisitor visitor) {
+        accept(new NodeMapper<Void>() {
+            @Override
+            public Void visit(LiteralNode literal) {
+                visitor.visit(literal);
+                return null;
+            }
+
+            @Override
+            public Void visit(VariableReferenceNode variableReference) {
+                visitor.visit(variableReference);
+                return null;
+            }
+
+            @Override
+            public Void visit(ThisReferenceNode reference) {
+                visitor.visit(reference);
+                return null;
+            }
+
+            @Override
+            public Void visit(AssignmentNode assignment) {
+                visitor.visit(assignment);
+                return null;
+            }
+
+            @Override
+            public Void visit(TernaryConditionalNode ternaryConditional) {
+                visitor.visit(ternaryConditional);
+                return null;
+            }
+
+            @Override
+            public Void visit(MethodCallNode methodCall) {
+                visitor.visit(methodCall);
+                return null;
+            }
+
+            @Override
+            public Void visit(StaticMethodCallNode staticMethodCall) {
+                visitor.visit(staticMethodCall);
+                return null;
+            }
+
+            @Override
+            public Void visit(ConstructorCallNode call) {
+                visitor.visit(call);
+                return null;
+            }
+
+            @Override
+            public Void visit(FieldAccessNode fieldAccess) {
+                visitor.visit(fieldAccess);
+                return null;
+            }
+
+            @Override
+            public Void visit(ReturnNode returnNode) {
+                visitor.visit(returnNode);
+                return null;
+            }
+
+            @Override
+            public Void visit(ExpressionStatementNode expressionStatement) {
+                visitor.visit(expressionStatement);
+                return null;
+            }
+
+            @Override
+            public Void visit(LocalVariableDeclarationNode localVariableDeclaration) {
+                visitor.visit(localVariableDeclaration);
+                return null;
+            }
+
+            @Override
+            public Void visit(IfStatementNode ifStatement) {
+                visitor.visit(ifStatement);
+                return null;
+            }
+
+            @Override
+            public Void visit(WhileNode whileLoop) {
+                visitor.visit(whileLoop);
+                return null;
+            }
+
+            @Override
+            public Void visit(FormalArgumentNode formalArgumentNode) {
+                visitor.visit(formalArgumentNode);
+                return null;
+            }
+
+            @Override
+            public Void visit(AnnotationNode annotation) {
+                visitor.visit(annotation);
+                return null;
+            }
+
+            @Override
+            public Void visit(MethodNode methodNode) {
+                visitor.visit(methodNode);
+                return null;
+            }
+
+            @Override
+            public Void visit(ConstructorNode constructorNode) {
+                visitor.visit(constructorNode);
+                return null;
+            }
+
+            @Override
+            public Void visit(FieldDeclarationNode declaration) {
+                visitor.visit(declaration);
+                return null;
+            }
+
+            @Override
+            public Void visit(ClassNode classNode) {
+                visitor.visit(classNode);
+                return null;
+            }
+        });
+    }
+    <T> T accept(NodeMapper<T> visitor);
 }
diff --git a/src/main/java/org/zwobble/couscous/ast/StatementNode.java b/src/main/java/org/zwobble/couscous/ast/StatementNode.java
index 6ba198f..a052917 100644
--- a/src/main/java/org/zwobble/couscous/ast/StatementNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/StatementNode.java
@@ -1,43 +1,37 @@
 package org.zwobble.couscous.ast;
 
-import org.zwobble.couscous.ast.visitors.NodeVisitor;
+import org.zwobble.couscous.ast.visitors.NodeMapper;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
 public interface StatementNode extends Node {
     <T> T accept(StatementNodeMapper<T> visitor);
     
     @Override
-    default void accept(NodeVisitor visitor) {
-        accept(new StatementNodeMapper<Void>() {
+    default <T> T accept(NodeMapper<T> visitor) {
+        return accept(new StatementNodeMapper<T>() {
             @Override
-            public Void visit(ReturnNode returnNode) {
-                visitor.visit(returnNode);
-                return null;
+            public T visit(ReturnNode returnNode) {
+                return visitor.visit(returnNode);
             }
 
             @Override
-            public Void visit(ExpressionStatementNode expressionStatement) {
-                visitor.visit(expressionStatement);
-                return null;
+            public T visit(ExpressionStatementNode expressionStatement) {
+                return visitor.visit(expressionStatement);
             }
 
             @Override
-            public Void visit(
-                LocalVariableDeclarationNode localVariableDeclaration) {
-                visitor.visit(localVariableDeclaration);
-                return null;
+            public T visit(LocalVariableDeclarationNode localVariableDeclaration) {
+                return visitor.visit(localVariableDeclaration);
             }
 
             @Override
-            public Void visit(IfStatementNode ifStatement) {
-                visitor.visit(ifStatement);
-                return null;
+            public T visit(IfStatementNode ifStatement) {
+                return visitor.visit(ifStatement);
             }
 
             @Override
-            public Void visit(WhileNode whileLoop) {
-                visitor.visit(whileLoop);
-                return null;
+            public T visit(WhileNode whileLoop) {
+                return visitor.visit(whileLoop);
             }
         });
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/structure/NodeStructure.java b/src/main/java/org/zwobble/couscous/ast/structure/NodeStructure.java
new file mode 100644
index 0000000..5fdd505
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/ast/structure/NodeStructure.java
@@ -0,0 +1,138 @@
+package org.zwobble.couscous.ast.structure;
+
+import org.zwobble.couscous.ast.*;
+import org.zwobble.couscous.ast.visitors.NodeMapper;
+
+import java.util.stream.Stream;
+
+import static org.zwobble.couscous.util.ExtraStreams.concatStreams;
+
+public class NodeStructure {
+    public static Stream<? extends Node> descendantNodesAndSelf(Node node) {
+        return Stream.concat(Stream.of(node), descendantNodes(node));
+    }
+
+    public static Stream<? extends Node> descendantNodes(Node node) {
+        return childNodes(node).flatMap(NodeStructure::descendantNodesAndSelf);
+    }
+
+    public static Stream<? extends Node> childNodes(Node node) {
+        return node.accept(new NodeMapper<Stream<? extends Node>>() {
+            @Override
+            public Stream<Node> visit(LiteralNode literal) {
+                return Stream.empty();
+            }
+
+            @Override
+            public Stream<Node> visit(VariableReferenceNode variableReference) {
+                return Stream.empty();
+            }
+
+            @Override
+            public Stream<Node> visit(ThisReferenceNode reference) {
+                return Stream.empty();
+            }
+
+            @Override
+            public Stream<Node> visit(AssignmentNode assignment) {
+                return Stream.of(assignment.getValue(), assignment.getTarget());
+            }
+
+            @Override
+            public Stream<Node> visit(TernaryConditionalNode ternaryConditional) {
+                return Stream.of(
+                    ternaryConditional.getCondition(),
+                    ternaryConditional.getIfTrue(),
+                    ternaryConditional.getIfFalse());
+            }
+
+            @Override
+            public Stream<Node> visit(MethodCallNode methodCall) {
+                return Stream.concat(
+                    Stream.of(methodCall.getReceiver()),
+                    methodCall.getArguments().stream());
+            }
+
+            @Override
+            public Stream<ExpressionNode> visit(StaticMethodCallNode staticMethodCall) {
+                return staticMethodCall.getArguments().stream();
+            }
+
+            @Override
+            public Stream<ExpressionNode> visit(ConstructorCallNode call) {
+                return call.getArguments().stream();
+            }
+
+            @Override
+            public Stream<Node> visit(FieldAccessNode fieldAccess) {
+                return Stream.of(fieldAccess.getLeft());
+            }
+
+            @Override
+            public Stream<Node> visit(ReturnNode returnNode) {
+                return Stream.of(returnNode.getValue());
+            }
+
+            @Override
+            public Stream<Node> visit(ExpressionStatementNode expressionStatement) {
+                return Stream.of(expressionStatement.getExpression());
+            }
+
+            @Override
+            public Stream<Node> visit(LocalVariableDeclarationNode localVariableDeclaration) {
+                return Stream.of(localVariableDeclaration.getInitialValue());
+            }
+
+            @Override
+            public Stream<Node> visit(IfStatementNode ifStatement) {
+                return concatStreams(
+                    Stream.of(ifStatement.getCondition()),
+                    ifStatement.getTrueBranch().stream(),
+                    ifStatement.getFalseBranch().stream());
+            }
+
+            @Override
+            public Stream<Node> visit(WhileNode whileLoop) {
+                return Stream.concat(
+                    Stream.of(whileLoop.getCondition()),
+                    whileLoop.getBody().stream());
+            }
+
+            @Override
+            public Stream<? extends Node> visit(FormalArgumentNode formalArgumentNode) {
+                return Stream.empty();
+            }
+
+            @Override
+            public Stream<? extends Node> visit(AnnotationNode annotation) {
+                return Stream.empty();
+            }
+
+            @Override
+            public Stream<Node> visit(ClassNode classNode) {
+                return concatStreams(
+                    Stream.of(classNode.getConstructor()),
+                    classNode.getFields().stream(),
+                    classNode.getMethods().stream());
+            }
+
+            @Override
+            public Stream<Node> visit(MethodNode methodNode) {
+                return concatStreams(
+                    methodNode.getAnnotations().stream(),
+                    methodNode.getArguments().stream(),
+                    methodNode.getBody().stream());
+            }
+
+            @Override
+            public Stream<Node> visit(ConstructorNode constructorNode) {
+                return Stream.concat(constructorNode.getArguments().stream(), constructorNode.getBody().stream());
+            }
+
+            @Override
+            public Stream<? extends Node> visit(FieldDeclarationNode declaration) {
+                return Stream.empty();
+            }
+        });
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/NodeMapper.java b/src/main/java/org/zwobble/couscous/ast/visitors/NodeMapper.java
new file mode 100644
index 0000000..d38d7fa
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeMapper.java
@@ -0,0 +1,28 @@
+package org.zwobble.couscous.ast.visitors;
+
+import org.zwobble.couscous.ast.*;
+
+public interface NodeMapper<T> {
+    T visit(LiteralNode literal);
+    T visit(VariableReferenceNode variableReference);
+    T visit(ThisReferenceNode reference);
+    T visit(AssignmentNode assignment);
+    T visit(TernaryConditionalNode ternaryConditional);
+    T visit(MethodCallNode methodCall);
+    T visit(StaticMethodCallNode staticMethodCall);
+    T visit(ConstructorCallNode call);
+    T visit(FieldAccessNode fieldAccess);
+
+    T visit(ReturnNode returnNode);
+    T visit(ExpressionStatementNode expressionStatement);
+    T visit(LocalVariableDeclarationNode localVariableDeclaration);
+    T visit(IfStatementNode ifStatement);
+    T visit(WhileNode whileLoop);
+
+    T visit(FormalArgumentNode formalArgumentNode);
+    T visit(AnnotationNode annotation);
+    T visit(MethodNode methodNode);
+    T visit(ConstructorNode constructorNode);
+    T visit(FieldDeclarationNode declaration);
+    T visit(ClassNode classNode);
+}
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/NodeMapperWithDefault.java b/src/main/java/org/zwobble/couscous/ast/visitors/NodeMapperWithDefault.java
new file mode 100644
index 0000000..e4203e0
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeMapperWithDefault.java
@@ -0,0 +1,111 @@
+package org.zwobble.couscous.ast.visitors;
+
+import org.zwobble.couscous.ast.*;
+
+public class NodeMapperWithDefault<T> implements NodeMapper<T> {
+    private final T defaultValue;
+
+    public NodeMapperWithDefault(T defaultValue) {
+        this.defaultValue = defaultValue;
+    }
+
+    @Override
+    public T visit(LiteralNode literal) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(VariableReferenceNode variableReference) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(ThisReferenceNode reference) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(AssignmentNode assignment) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(TernaryConditionalNode ternaryConditional) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(MethodCallNode methodCall) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(StaticMethodCallNode staticMethodCall) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(ConstructorCallNode call) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(FieldAccessNode fieldAccess) {
+        return defaultValue;
+    }
+
+    @Override
+    public T visit(ReturnNode returnNode) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(ExpressionStatementNode expressionStatement) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(LocalVariableDeclarationNode localVariableDeclaration) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(IfStatementNode ifStatement) {
+        return defaultValue;
+    }
+
+    @Override
+    public T visit(WhileNode whileLoop) {
+        return defaultValue;
+    }
+
+    @Override
+    public T visit(AnnotationNode annotation) {
+        return defaultValue;
+    }
+
+    @Override
+    public T visit(FieldDeclarationNode declaration) {
+        return defaultValue;
+    }
+
+    @Override
+    public T visit(FormalArgumentNode formalArgumentNode) {
+        return defaultValue;
+    }
+
+    @Override
+    public T visit(ClassNode classNode) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(MethodNode methodNode) {
+        return defaultValue;
+    }
+    
+    @Override
+    public T visit(ConstructorNode constructorNode) {
+        return defaultValue;
+    }
+}
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 8d42f5c..75e6da5 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitor.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitor.java
@@ -18,8 +18,11 @@ public interface NodeVisitor {
     void visit(LocalVariableDeclarationNode localVariableDeclaration);
     void visit(IfStatementNode ifStatement);
     void visit(WhileNode whileLoop);
-    
-    void visit(ClassNode classNode);
+
+    void visit(FormalArgumentNode formalArgumentNode);
+    void visit(AnnotationNode annotation);
     void visit(MethodNode methodNode);
     void visit(ConstructorNode constructorNode);
+    void visit(FieldDeclarationNode declaration);
+    void visit(ClassNode classNode);
 }
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitorWithEmptyDefaults.java b/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitorWithEmptyDefaults.java
deleted file mode 100644
index c8dd2cf..0000000
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitorWithEmptyDefaults.java
+++ /dev/null
@@ -1,73 +0,0 @@
-package org.zwobble.couscous.ast.visitors;
-
-import org.zwobble.couscous.ast.*;
-
-public interface NodeVisitorWithEmptyDefaults extends NodeVisitor {
-    @Override
-    default void visit(LiteralNode literal) {
-    }
-    
-    @Override
-    default void visit(VariableReferenceNode variableReference) {
-    }
-    
-    @Override
-    default void visit(ThisReferenceNode reference) {
-    }
-    
-    @Override
-    default void visit(AssignmentNode assignment) {
-    }
-    
-    @Override
-    default void visit(TernaryConditionalNode ternaryConditional) {
-    }
-    
-    @Override
-    default void visit(MethodCallNode methodCall) {
-    }
-    
-    @Override
-    default void visit(StaticMethodCallNode staticMethodCall) {
-    }
-    
-    @Override
-    default void visit(ConstructorCallNode call) {
-    }
-    
-    @Override
-    default void visit(FieldAccessNode fieldAccess) {
-    }
-
-    @Override
-    default void visit(ReturnNode returnNode) {
-    }
-    
-    @Override
-    default void visit(ExpressionStatementNode expressionStatement) {
-    }
-    
-    @Override
-    default void visit(LocalVariableDeclarationNode localVariableDeclaration) {
-    }
-    
-    @Override
-    default void visit(IfStatementNode ifStatement) {
-    }
-
-    @Override
-    default void visit(WhileNode whileLoop) {
-    }
-    
-    @Override
-    default void visit(ClassNode classNode) {
-    }
-    
-    @Override
-    default void visit(MethodNode methodNode) {
-    }
-    
-    @Override
-    default void visit(ConstructorNode constructorNode) {
-    }
-}
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitors.java b/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitors.java
deleted file mode 100644
index 78a8254..0000000
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeVisitors.java
+++ /dev/null
@@ -1,119 +0,0 @@
-package org.zwobble.couscous.ast.visitors;
-
-import java.util.List;
-
-import org.zwobble.couscous.ast.*;
-
-public class NodeVisitors {
-    public static void visitAll(Node node, NodeVisitor visitor) {
-        node.accept(new NodeVisitor(){
-            @Override
-            public void visit(LiteralNode literal) {
-                visitor.visit(literal);
-            }
-            
-            @Override
-            public void visit(VariableReferenceNode variableReference) {
-                visitor.visit(variableReference);
-            }
-            
-            @Override
-            public void visit(ThisReferenceNode reference) {
-                visitor.visit(reference);
-            }
-            
-            @Override
-            public void visit(AssignmentNode assignment) {
-                visitor.visit(assignment);
-                assignment.getValue().accept(this);
-                assignment.getTarget().accept(this);
-            }
-            
-            @Override
-            public void visit(TernaryConditionalNode ternaryConditional) {
-                visitor.visit(ternaryConditional);
-                ternaryConditional.getCondition().accept(this);
-                ternaryConditional.getIfTrue().accept(this);
-                ternaryConditional.getIfFalse().accept(this);
-            }
-            
-            @Override
-            public void visit(MethodCallNode methodCall) {
-                visitor.visit(methodCall);
-                methodCall.getReceiver().accept(this);
-                visitAll(methodCall.getArguments());
-            }
-            
-            @Override
-            public void visit(StaticMethodCallNode staticMethodCall) {
-                visitor.visit(staticMethodCall);
-                visitAll(staticMethodCall.getArguments());
-            }
-            
-            @Override
-            public void visit(ConstructorCallNode call) {
-                visitor.visit(call);
-                visitAll(call.getArguments());
-            }
-            
-            @Override
-            public void visit(FieldAccessNode fieldAccess) {
-                visitor.visit(fieldAccess);
-            }
-            
-            @Override
-            public void visit(ReturnNode returnNode) {
-                visitor.visit(returnNode);
-                returnNode.getValue().accept(this);
-            }
-            
-            @Override
-            public void visit(ExpressionStatementNode expressionStatement) {
-                visitor.visit(expressionStatement);
-                expressionStatement.getExpression().accept(this);
-            }
-            
-            @Override
-            public void visit(LocalVariableDeclarationNode localVariableDeclaration) {
-                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(WhileNode whileLoop) {
-                whileLoop.getCondition().accept(visitor);
-                visitAll(whileLoop.getBody());
-            }
-
-            @Override
-            public void visit(ClassNode classNode) {
-                visitor.visit(classNode);
-                classNode.getConstructor().accept(this);
-                visitAll(classNode.getMethods());
-            }
-            
-            @Override
-            public void visit(MethodNode methodNode) {
-                visitAll(methodNode.getBody());
-            }
-            
-            @Override
-            public void visit(ConstructorNode constructorNode) {
-                visitAll(constructorNode.getBody());
-            }
-            
-            private <T extends Node> void visitAll(List<T> nodes) {
-                for (Node node : nodes) {
-                    node.accept(this);
-                }
-            }
-        });
-    }
-}
\ No newline at end of file
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 5f6fefb..574854d 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -1,5 +1,18 @@
 package org.zwobble.couscous.backends.python;
 
+import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Iterators;
+import org.zwobble.couscous.ast.*;
+import org.zwobble.couscous.ast.structure.NodeStructure;
+import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
+import org.zwobble.couscous.ast.visitors.NodeMapperWithDefault;
+import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
+import org.zwobble.couscous.backends.python.PrimitiveMethods.PrimitiveMethodGenerator;
+import org.zwobble.couscous.backends.python.ast.*;
+import org.zwobble.couscous.values.*;
+
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
@@ -7,35 +20,6 @@ import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-import org.zwobble.couscous.ast.*;
-import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
-import org.zwobble.couscous.ast.visitors.NodeVisitorWithEmptyDefaults;
-import org.zwobble.couscous.ast.visitors.NodeVisitors;
-import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
-import org.zwobble.couscous.backends.python.PrimitiveMethods.PrimitiveMethodGenerator;
-import org.zwobble.couscous.backends.python.ast.PythonAttributeAccessNode;
-import org.zwobble.couscous.backends.python.ast.PythonBlock;
-import org.zwobble.couscous.backends.python.ast.PythonClassNode;
-import org.zwobble.couscous.backends.python.ast.PythonExpressionNode;
-import org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode;
-import org.zwobble.couscous.backends.python.ast.PythonImportNode;
-import org.zwobble.couscous.backends.python.ast.PythonModuleNode;
-import org.zwobble.couscous.backends.python.ast.PythonStatementNode;
-import org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNode;
-import org.zwobble.couscous.values.BooleanValue;
-import org.zwobble.couscous.values.IntegerValue;
-import org.zwobble.couscous.values.InternalCouscousValue;
-import org.zwobble.couscous.values.PrimitiveValue;
-import org.zwobble.couscous.values.PrimitiveValueVisitor;
-import org.zwobble.couscous.values.StringValue;
-import org.zwobble.couscous.values.UnitValue;
-
-import com.google.common.base.Strings;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Iterators;
-
 import static com.google.common.collect.Iterators.singletonIterator;
 import static java.util.Arrays.asList;
 import static org.zwobble.couscous.backends.python.ast.PythonAssignmentNode.pythonAssignment;
@@ -103,19 +87,19 @@ public class PythonCodeGenerator {
     }
     
     private static Set<TypeName> findReferencedClasses(ClassNode classNode) {
-        ImmutableSet.Builder<TypeName> imports = ImmutableSet.builder();
-        NodeVisitors.visitAll(classNode, new NodeVisitorWithEmptyDefaults(){
-            @Override
-            public void visit(StaticMethodCallNode staticMethodCall) {
-                imports.add(staticMethodCall.getClassName());
-            }
-            
-            @Override
-            public void visit(ConstructorCallNode call) {
-                imports.add(call.getType());
-            }
-        });
-        return imports.build();
+        return NodeStructure.descendantNodes(classNode)
+            .flatMap(node -> node.accept(new NodeMapperWithDefault<Stream<TypeName>>(Stream.empty()) {
+                @Override
+                public Stream<TypeName> visit(StaticMethodCallNode staticMethodCall) {
+                    return Stream.of(staticMethodCall.getClassName());
+                }
+
+                @Override
+                public Stream<TypeName> visit(ConstructorCallNode call) {
+                    return Stream.of(call.getType());
+                }
+            }))
+            .collect(Collectors.toSet());
     }
     
     public static PythonExpressionNode generateCode(PrimitiveValue value) {
diff --git a/src/main/java/org/zwobble/couscous/interpreter/Executor.java b/src/main/java/org/zwobble/couscous/interpreter/Executor.java
index 8909acc..c9528be 100644
--- a/src/main/java/org/zwobble/couscous/interpreter/Executor.java
+++ b/src/main/java/org/zwobble/couscous/interpreter/Executor.java
@@ -1,15 +1,18 @@
 package org.zwobble.couscous.interpreter;
 
-import java.util.List;
-import java.util.Optional;
-import java.util.stream.Stream;
-
 import org.zwobble.couscous.ast.*;
+import org.zwobble.couscous.ast.structure.NodeStructure;
+import org.zwobble.couscous.ast.visitors.NodeMapperWithDefault;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 import org.zwobble.couscous.interpreter.values.BooleanInterpreterValue;
 import org.zwobble.couscous.interpreter.values.InterpreterValue;
 import org.zwobble.couscous.interpreter.values.UnitInterpreterValue;
 
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Stream;
+
+import static org.zwobble.couscous.ast.structure.NodeStructure.descendantNodesAndSelf;
 import static org.zwobble.couscous.interpreter.Evaluator.eval;
 
 public class Executor implements StatementNodeMapper<Optional<InterpreterValue>> {
@@ -34,34 +37,15 @@ public class Executor implements StatementNodeMapper<Optional<InterpreterValue>>
     }
 
     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);
-            }
-
-            @Override
-            public Stream<VariableNode> visit(IfStatementNode ifStatement) {
-                // TODO: Handle declarations in branches
-                return Stream.empty();
-            }
-
-            @Override
-            public Stream<VariableNode> visit(WhileNode whileLoop) {
-                // TODO: Handle declarations in branches
-                return Stream.empty();
-            }
-        }));
+        // TODO: add filter argument to descendantNodesAndSelf to allow recursion to be stopped early
+        return body.stream()
+            .flatMap(NodeStructure::descendantNodesAndSelf)
+            .flatMap(node -> node.accept(new NodeMapperWithDefault<Stream<VariableNode>>(Stream.empty()) {
+                @Override
+                public Stream<VariableNode> visit(LocalVariableDeclarationNode localVariableDeclaration) {
+                    return Stream.of(localVariableDeclaration);
+                }
+            }));
     }
 
     public static Optional<InterpreterValue> exec(Environment environment, StatementNode statement) {
diff --git a/src/main/java/org/zwobble/couscous/util/ExtraStreams.java b/src/main/java/org/zwobble/couscous/util/ExtraStreams.java
new file mode 100644
index 0000000..6bbf1c8
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/util/ExtraStreams.java
@@ -0,0 +1,9 @@
+package org.zwobble.couscous.util;
+
+import java.util.stream.Stream;
+
+public class ExtraStreams {
+    public static <T> Stream<T> concatStreams(Stream<? extends T> first, Stream<? extends T> second, Stream<? extends T> third) {
+        return Stream.concat(first, Stream.concat(second, third));
+    }
+}