Move childNodes onto Node interface

1 files changed, 11 insertions, 0 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/TryNode.java b/src/main/java/org/zwobble/couscous/ast/TryNode.java
index 800931f..3c18e1a 100644
--- a/src/main/java/org/zwobble/couscous/ast/TryNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/TryNode.java
@@ -1,10 +1,12 @@
 package org.zwobble.couscous.ast;
 
+import com.google.common.collect.Iterables;
 import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
 import java.util.List;
 
+import static org.zwobble.couscous.util.ExtraIterables.lazyFlatMap;
 import static org.zwobble.couscous.util.ExtraLists.eagerMap;
 
 public class TryNode implements StatementNode {
@@ -48,6 +50,15 @@ public class TryNode implements StatementNode {
     }
 
     @Override
+    public Iterable<? extends Node> childNodes() {
+        return Iterables.concat(
+            body,
+            lazyFlatMap(exceptionHandlers, handler -> handler.getBody()),
+            finallyBody
+        );
+    }
+
+    @Override
     public StatementNode transform(NodeTransformer transformer) {
         return new TryNode(
             transformer.transformStatements(body),