Implement visitors on OperationNode instead of desugarring

1 files changed, 6 insertions, 0 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/structure/NodeStructure.java b/src/main/java/org/zwobble/couscous/ast/structure/NodeStructure.java
index 2fd444d..a4cedca 100644
--- a/src/main/java/org/zwobble/couscous/ast/structure/NodeStructure.java
+++ b/src/main/java/org/zwobble/couscous/ast/structure/NodeStructure.java
@@ -28,6 +28,7 @@ public class NodeStructure {
     }
 
     public static Stream<? extends Node> childNodes(Node node) {
+        // TODO: is this just duplicating the logic already in node.transform?
         return node.accept(new NodeMapper<Stream<? extends Node>>() {
             @Override
             public Stream<Node> visit(LiteralNode literal) {
@@ -70,6 +71,11 @@ public class NodeStructure {
             }
 
             @Override
+            public Stream<? extends Node> visit(OperationNode operation) {
+                return operation.getArguments().stream();
+            }
+
+            @Override
             public Stream<Node> visit(FieldAccessNode fieldAccess) {
                 return Stream.of(fieldAccess.getLeft());
             }