Replace captured variable references in method bodies

1 files changed, 11 insertions, 0 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java b/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java
index ae3e29b..cb5eee8 100644
--- a/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java
@@ -1,9 +1,12 @@
 package org.zwobble.couscous.ast;
 
 import java.util.List;
+import java.util.function.Function;
 
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
+import static org.zwobble.couscous.util.ExtraLists.eagerMap;
+
 public class IfStatementNode implements StatementNode {
     public static IfStatementNode ifStatement(
             ExpressionNode condition,
@@ -43,6 +46,14 @@ public class IfStatementNode implements StatementNode {
     }
 
     @Override
+    public StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new IfStatementNode(
+            replace.apply(condition),
+            eagerMap(trueBranch, statement -> statement.replaceExpressions(replace)),
+            eagerMap(falseBranch, statement -> statement.replaceExpressions(replace)));
+    }
+
+    @Override
     public String toString() {
         return "IfStatementNode(condition=" + condition + ", trueBranch="
                + trueBranch + ", falseBranch=" + falseBranch + ")";