Replace captured variable references in method bodies

1 files changed, 10 insertions, 0 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/WhileNode.java b/src/main/java/org/zwobble/couscous/ast/WhileNode.java
index 0a67aef..6462821 100644
--- a/src/main/java/org/zwobble/couscous/ast/WhileNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/WhileNode.java
@@ -4,6 +4,9 @@ import com.google.common.base.Objects;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
 import java.util.List;
+import java.util.function.Function;
+
+import static org.zwobble.couscous.util.ExtraLists.eagerMap;
 
 public class WhileNode implements StatementNode {
     public static WhileNode whileLoop(ExpressionNode condition, List<StatementNode> body) {
@@ -32,6 +35,13 @@ public class WhileNode implements StatementNode {
     }
 
     @Override
+    public StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new WhileNode(
+            replace.apply(condition),
+            eagerMap(body, statement -> statement.replaceExpressions(replace)));
+    }
+
+    @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;