1 files changed, 16 insertions, 32 deletions
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) {