Replace captured variable references in method bodies

16 files changed, 128 insertions, 2 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();
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/ConstructorCallNode.java b/src/main/java/org/zwobble/couscous/ast/ConstructorCallNode.java
index 6a4655a..ab4b1ea 100644
--- a/src/main/java/org/zwobble/couscous/ast/ConstructorCallNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ConstructorCallNode.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.ExpressionNodeMapper;
 
+import static org.zwobble.couscous.util.ExtraLists.eagerMap;
+
 public class ConstructorCallNode implements ExpressionNode {
     public static ConstructorCallNode constructorCall(
             TypeName type,
@@ -36,6 +39,11 @@ public class ConstructorCallNode implements ExpressionNode {
     }
 
     @Override
+    public ExpressionNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new ConstructorCallNode(type, eagerMap(arguments, replace::apply));
+    }
+
+    @Override
     public String toString() {
         return "ConstructorCallNode(type=" + type + ", arguments=" + arguments
                + ")";
diff --git a/src/main/java/org/zwobble/couscous/ast/ExpressionNode.java b/src/main/java/org/zwobble/couscous/ast/ExpressionNode.java
index 65757ad..0f1ab4f 100644
--- a/src/main/java/org/zwobble/couscous/ast/ExpressionNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ExpressionNode.java
@@ -3,6 +3,8 @@ package org.zwobble.couscous.ast;
 import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
 import org.zwobble.couscous.ast.visitors.NodeMapper;
 
+import java.util.function.Function;
+
 public interface ExpressionNode extends Node {
     TypeName getType();
     <T> T accept(ExpressionNodeMapper<T> visitor);
@@ -56,4 +58,6 @@ public interface ExpressionNode extends Node {
             }
         });
     }
+
+    ExpressionNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace);
 }
diff --git a/src/main/java/org/zwobble/couscous/ast/ExpressionStatementNode.java b/src/main/java/org/zwobble/couscous/ast/ExpressionStatementNode.java
index b814076..0ac7381 100644
--- a/src/main/java/org/zwobble/couscous/ast/ExpressionStatementNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ExpressionStatementNode.java
@@ -2,6 +2,8 @@ package org.zwobble.couscous.ast;
 
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
+import java.util.function.Function;
+
 public class ExpressionStatementNode implements StatementNode {
     public static ExpressionStatementNode expressionStatement(ExpressionNode expression) {
         return new ExpressionStatementNode(expression);
@@ -23,6 +25,11 @@ public class ExpressionStatementNode implements StatementNode {
     }
 
     @Override
+    public StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new ExpressionStatementNode(replace.apply(expression));
+    }
+
+    @Override
     public String toString() {
         return "ExpressionStatementNode(expression=" + expression + ")";
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/FieldAccessNode.java b/src/main/java/org/zwobble/couscous/ast/FieldAccessNode.java
index f3c92f9..9f881b4 100644
--- a/src/main/java/org/zwobble/couscous/ast/FieldAccessNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/FieldAccessNode.java
@@ -3,6 +3,8 @@ package org.zwobble.couscous.ast;
 import org.zwobble.couscous.ast.visitors.AssignableExpressionNodeVisitor;
 import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
 
+import java.util.function.Function;
+
 public class FieldAccessNode implements AssignableExpressionNode {
     public static FieldAccessNode fieldAccess(
             ExpressionNode left,
@@ -42,6 +44,11 @@ public class FieldAccessNode implements AssignableExpressionNode {
     }
 
     @Override
+    public ExpressionNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new FieldAccessNode(replace.apply(left), fieldName, type);
+    }
+
+    @Override
     public void accept(AssignableExpressionNodeVisitor mapper) {
         mapper.visit(this);
     }
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 + ")";
diff --git a/src/main/java/org/zwobble/couscous/ast/LiteralNode.java b/src/main/java/org/zwobble/couscous/ast/LiteralNode.java
index ca005de..1aa3293 100644
--- a/src/main/java/org/zwobble/couscous/ast/LiteralNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/LiteralNode.java
@@ -6,6 +6,8 @@ import org.zwobble.couscous.values.IntegerValue;
 import org.zwobble.couscous.values.PrimitiveValue;
 import org.zwobble.couscous.values.StringValue;
 
+import java.util.function.Function;
+
 import static org.zwobble.couscous.values.PrimitiveValues.value;
 
 public class LiteralNode implements ExpressionNode {
@@ -51,6 +53,11 @@ public class LiteralNode implements ExpressionNode {
     }
 
     @Override
+    public ExpressionNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return this;
+    }
+
+    @Override
     public String toString() {
         return "LiteralNode(value=" + value + ", type=" + type + ")";
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/LocalVariableDeclarationNode.java b/src/main/java/org/zwobble/couscous/ast/LocalVariableDeclarationNode.java
index bb272fc..3499e26 100644
--- a/src/main/java/org/zwobble/couscous/ast/LocalVariableDeclarationNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/LocalVariableDeclarationNode.java
@@ -2,6 +2,8 @@ package org.zwobble.couscous.ast;
 
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
+import java.util.function.Function;
+
 import static org.zwobble.couscous.ast.VariableDeclaration.var;
 
 public class LocalVariableDeclarationNode implements VariableNode, StatementNode {
@@ -51,6 +53,11 @@ public class LocalVariableDeclarationNode implements VariableNode, StatementNode
     }
 
     @Override
+    public StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new LocalVariableDeclarationNode(declaration, replace.apply(initialValue));
+    }
+
+    @Override
     public String toString() {
         return "LocalVariableDeclarationNode(declaration=" + declaration
                + ", initialValue=" + initialValue + ")";
diff --git a/src/main/java/org/zwobble/couscous/ast/MethodCallNode.java b/src/main/java/org/zwobble/couscous/ast/MethodCallNode.java
index cce140f..f65a9ae 100644
--- a/src/main/java/org/zwobble/couscous/ast/MethodCallNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/MethodCallNode.java
@@ -1,14 +1,15 @@
 package org.zwobble.couscous.ast;
 
-import com.sun.org.apache.xpath.internal.operations.Bool;
 import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
 import org.zwobble.couscous.values.BooleanValue;
 import org.zwobble.couscous.values.IntegerValue;
 
 import java.util.Collections;
 import java.util.List;
+import java.util.function.Function;
 
 import static java.util.Arrays.asList;
+import static org.zwobble.couscous.util.ExtraLists.eagerMap;
 
 public class MethodCallNode implements ExpressionNode {
     public static ExpressionNode not(ExpressionNode value) {
@@ -117,6 +118,15 @@ public class MethodCallNode implements ExpressionNode {
     }
 
     @Override
+    public ExpressionNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new MethodCallNode(
+            replace.apply(receiver),
+            methodName,
+            eagerMap(arguments, replace::apply),
+            type);
+    }
+
+    @Override
     public String toString() {
         return "MethodCallNode(receiver=" + receiver + ", methodName="
                + methodName + ", arguments=" + arguments + ", type=" + type
diff --git a/src/main/java/org/zwobble/couscous/ast/ReturnNode.java b/src/main/java/org/zwobble/couscous/ast/ReturnNode.java
index 008b985..68330af 100644
--- a/src/main/java/org/zwobble/couscous/ast/ReturnNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ReturnNode.java
@@ -2,6 +2,8 @@ package org.zwobble.couscous.ast;
 
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
+import java.util.function.Function;
+
 public class ReturnNode implements StatementNode {
     public static ReturnNode returns(ExpressionNode value) {
         return new ReturnNode(value);
@@ -23,6 +25,11 @@ public class ReturnNode implements StatementNode {
     }
 
     @Override
+    public StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new ReturnNode(replace.apply(value));
+    }
+
+    @Override
     public String toString() {
         return "ReturnNode(value=" + value + ")";
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/StatementNode.java b/src/main/java/org/zwobble/couscous/ast/StatementNode.java
index a052917..4b9c9de 100644
--- a/src/main/java/org/zwobble/couscous/ast/StatementNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/StatementNode.java
@@ -3,6 +3,8 @@ package org.zwobble.couscous.ast;
 import org.zwobble.couscous.ast.visitors.NodeMapper;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
+import java.util.function.Function;
+
 public interface StatementNode extends Node {
     <T> T accept(StatementNodeMapper<T> visitor);
     
@@ -35,4 +37,6 @@ public interface StatementNode extends Node {
             }
         });
     }
+
+    StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace);
 }
diff --git a/src/main/java/org/zwobble/couscous/ast/StaticMethodCallNode.java b/src/main/java/org/zwobble/couscous/ast/StaticMethodCallNode.java
index b33e94c..ccb9854 100644
--- a/src/main/java/org/zwobble/couscous/ast/StaticMethodCallNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/StaticMethodCallNode.java
@@ -1,6 +1,7 @@
 package org.zwobble.couscous.ast;
 
 import java.util.List;
+import java.util.function.Function;
 
 import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
 import org.zwobble.couscous.values.BooleanValue;
@@ -9,6 +10,7 @@ import org.zwobble.couscous.values.InternalCouscousValue;
 import org.zwobble.couscous.values.ObjectValues;
 
 import static java.util.Arrays.asList;
+import static org.zwobble.couscous.util.ExtraLists.eagerMap;
 
 public class StaticMethodCallNode implements ExpressionNode {
     public static ExpressionNode same(ExpressionNode left, ExpressionNode right) {
@@ -101,6 +103,15 @@ public class StaticMethodCallNode implements ExpressionNode {
     }
 
     @Override
+    public ExpressionNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return new StaticMethodCallNode(
+            className,
+            methodName,
+            eagerMap(arguments, replace::apply),
+            type);
+    }
+
+    @Override
     public String toString() {
         return "StaticMethodCallNode(className=" + className + ", methodName="
                + methodName + ", arguments=" + arguments + ", type=" + type
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();
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/ThisReferenceNode.java b/src/main/java/org/zwobble/couscous/ast/ThisReferenceNode.java
index 4e7adc6..3ef5527 100644
--- a/src/main/java/org/zwobble/couscous/ast/ThisReferenceNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ThisReferenceNode.java
@@ -2,6 +2,8 @@ package org.zwobble.couscous.ast;
 
 import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
 
+import java.util.function.Function;
+
 public class ThisReferenceNode implements ExpressionNode {
     public static ThisReferenceNode thisReference(TypeName type) {
         return new ThisReferenceNode(type);
@@ -23,6 +25,11 @@ public class ThisReferenceNode implements ExpressionNode {
     }
 
     @Override
+    public ExpressionNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return this;
+    }
+
+    @Override
     public String toString() {
         return "ThisReferenceNode(type=" + type + ")";
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/VariableReferenceNode.java b/src/main/java/org/zwobble/couscous/ast/VariableReferenceNode.java
index e95510f..b576389 100644
--- a/src/main/java/org/zwobble/couscous/ast/VariableReferenceNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/VariableReferenceNode.java
@@ -3,6 +3,8 @@ package org.zwobble.couscous.ast;
 import org.zwobble.couscous.ast.visitors.AssignableExpressionNodeVisitor;
 import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
 
+import java.util.function.Function;
+
 public class VariableReferenceNode implements AssignableExpressionNode {
     public static VariableReferenceNode reference(VariableNode variable) {
         return new VariableReferenceNode(variable.getDeclaration());
@@ -26,7 +28,12 @@ public class VariableReferenceNode implements AssignableExpressionNode {
     public <T> T accept(ExpressionNodeMapper<T> visitor) {
         return visitor.visit(this);
     }
-    
+
+    @Override
+    public ExpressionNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return this;
+    }
+
     public String getReferentId() {
         return referent.getId();
     }
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;