Replace captured variable references in method bodies

1 files changed, 10 insertions, 0 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/TernaryConditionalNode.java b/src/main/java/org/zwobble/couscous/ast/TernaryConditionalNode.java
index e9d3f6c..f370b0a 100644
--- a/src/main/java/org/zwobble/couscous/ast/TernaryConditionalNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/TernaryConditionalNode.java
@@ -2,6 +2,8 @@ package org.zwobble.couscous.ast;
 
 import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
 
+import java.util.function.Function;
+
 public class TernaryConditionalNode implements ExpressionNode {
     public static TernaryConditionalNode ternaryConditional(
             ExpressionNode condition,
@@ -41,6 +43,14 @@ public class TernaryConditionalNode implements ExpressionNode {
     }
 
     @Override
+    public ExpressionNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new TernaryConditionalNode(
+            replace.apply(condition),
+            replace.apply(ifTrue),
+            replace.apply(ifFalse));
+    }
+
+    @Override
     public TypeName getType() {
         return ifTrue.getType();
     }