Replace captured variable references in method bodies

1 files changed, 9 insertions, 0 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/AssignmentNode.java b/src/main/java/org/zwobble/couscous/ast/AssignmentNode.java
index b6c8f9c..15ecf8f 100644
--- a/src/main/java/org/zwobble/couscous/ast/AssignmentNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/AssignmentNode.java
@@ -2,6 +2,8 @@ package org.zwobble.couscous.ast;
 
 import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
 
+import java.util.function.Function;
+
 import static org.zwobble.couscous.ast.ExpressionStatementNode.expressionStatement;
 import static org.zwobble.couscous.ast.VariableReferenceNode.reference;
 
@@ -44,6 +46,13 @@ public class AssignmentNode implements ExpressionNode {
     }
 
     @Override
+    public ExpressionNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new AssignmentNode(
+            (AssignableExpressionNode) replace.apply(target),
+            replace.apply(value));
+    }
+
+    @Override
     public TypeName getType() {
         return value.getType();
     }