Capture "this" references

4 files changed, 74 insertions, 48 deletions

diff --git a/src/main/java/org/zwobble/couscous/frontends/java/FreeVariables.java b/src/main/java/org/zwobble/couscous/frontends/java/FreeVariables.java
index 99cc9b4..00184b1 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/FreeVariables.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/FreeVariables.java
@@ -1,9 +1,9 @@
 package org.zwobble.couscous.frontends.java;
 
-import com.google.common.collect.Sets;
 import org.zwobble.couscous.ast.*;
 
 import java.util.List;
+import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -12,19 +12,23 @@ import static org.zwobble.couscous.util.Casts.tryCast;
 import static org.zwobble.couscous.util.ExtraStreams.toStream;
 
 public class FreeVariables {
-    public static List<VariableDeclaration> findFreeVariables(List<? extends Node> body) {
-        Stream<VariableDeclaration> referencedDeclarations = body.stream()
-            .flatMap(FreeVariables::findReferencedDeclarations);
-        Stream<VariableDeclaration> declarations = body.stream().flatMap(FreeVariables::findDeclarations);
-        return Sets.difference(
-            referencedDeclarations.collect(Collectors.toSet()),
-            declarations.collect(Collectors.toSet())).stream().collect(Collectors.toList());
+    public static List<ReferenceNode> findFreeVariables(List<? extends Node> body) {
+        Set<VariableDeclaration> declarations = body.stream()
+            .flatMap(FreeVariables::findDeclarations)
+            .collect(Collectors.toSet());
+
+        return body.stream()
+            .flatMap(FreeVariables::findReferences)
+            .filter(reference ->
+                tryCast(VariableReferenceNode.class, reference)
+                    .map(variableReference -> !declarations.contains(variableReference.getReferent()))
+                    .orElse(true))
+            .collect(Collectors.toList());
     }
 
-    private static Stream<VariableDeclaration> findReferencedDeclarations(Node root) {
-        Stream<VariableReferenceNode> references = descendantNodesAndSelf(root).flatMap(node ->
-            toStream(tryCast(VariableReferenceNode.class, node)));
-        return references.map(VariableReferenceNode::getReferent);
+    private static Stream<ReferenceNode> findReferences(Node root) {
+        return descendantNodesAndSelf(root).flatMap(node ->
+            toStream(tryCast(ReferenceNode.class, node)));
     }
 
     private static Stream<VariableDeclaration> findDeclarations(Node root) {
diff --git a/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java b/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
index 88d214c..9386ff6 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
@@ -1,6 +1,7 @@
 package org.zwobble.couscous.frontends.java;
 
 import org.zwobble.couscous.ast.ClassNode;
+import org.zwobble.couscous.ast.ReferenceNode;
 import org.zwobble.couscous.ast.TypeName;
 import org.zwobble.couscous.ast.VariableDeclaration;
 
@@ -8,9 +9,9 @@ import java.util.List;
 
 public class GeneratedClosure {
     private final ClassNode classNode;
-    private final List<VariableDeclaration> captures;
+    private final List<ReferenceNode> captures;
 
-    public GeneratedClosure(ClassNode classNode, List<VariableDeclaration> captures) {
+    public GeneratedClosure(ClassNode classNode, List<ReferenceNode> captures) {
         this.classNode = classNode;
         this.captures = captures;
     }
@@ -23,7 +24,7 @@ public class GeneratedClosure {
         return classNode.getName();
     }
 
-    public List<VariableDeclaration> getCaptures() {
+    public List<ReferenceNode> getCaptures() {
         return captures;
     }
 }
diff --git a/src/main/java/org/zwobble/couscous/frontends/java/JavaExpressionReader.java b/src/main/java/org/zwobble/couscous/frontends/java/JavaExpressionReader.java
index 13bc519..367a34f 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaExpressionReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaExpressionReader.java
@@ -17,7 +17,6 @@ import static org.zwobble.couscous.ast.MethodCallNode.not;
 import static org.zwobble.couscous.frontends.java.JavaOperators.readOperator;
 import static org.zwobble.couscous.frontends.java.JavaTypes.typeOf;
 import static org.zwobble.couscous.util.ExtraLists.concat;
-import static org.zwobble.couscous.util.ExtraLists.eagerMap;
 
 public class JavaExpressionReader {
     private final Scope scope;
@@ -232,8 +231,8 @@ public class JavaExpressionReader {
             captureArguments(closure));
     }
 
-    private List<ExpressionNode> captureArguments(GeneratedClosure closure) {
-        return eagerMap(closure.getCaptures(), VariableReferenceNode::reference);
+    private List<? extends ExpressionNode> captureArguments(GeneratedClosure closure) {
+        return closure.getCaptures();
     }
 
     private List<ExpressionNode> readArguments(IMethodBinding method, List<Expression> javaArguments) {
diff --git a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
index 1c53a44..5e81ec9 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
@@ -32,7 +32,6 @@ import static org.zwobble.couscous.ast.ThisReferenceNode.thisReference;
 import static org.zwobble.couscous.ast.VariableReferenceNode.reference;
 import static org.zwobble.couscous.frontends.java.FreeVariables.findFreeVariables;
 import static org.zwobble.couscous.frontends.java.JavaTypes.*;
-import static org.zwobble.couscous.util.Casts.tryCast;
 import static org.zwobble.couscous.util.ExtraLists.*;
 import static org.zwobble.couscous.util.ExtraMaps.toMap;
 
@@ -109,51 +108,47 @@ public class JavaReader {
     private List<StatementNode> replaceCaptureReferences(
         TypeName className,
         List<StatementNode> body,
-        List<VariableDeclaration> freeVariables)
+        List<CapturedVariable> freeVariables)
     {
-        Map<Identifier, ExpressionNode> freeVariablesById = Maps.transformValues(
-            Maps.uniqueIndex(freeVariables, VariableDeclaration::getId),
+        Map<ExpressionNode, ExpressionNode> replacements = Maps.transformValues(
+            Maps.uniqueIndex(freeVariables, variable -> variable.freeVariable),
             freeVariable -> captureAccess(className, freeVariable));
-        Function<ExpressionNode, ExpressionNode> replaceExpression = new CaptureReplacer(freeVariablesById);
+        Function<ExpressionNode, ExpressionNode> replaceExpression = new CaptureReplacer(replacements);
         return eagerMap(body, statement -> statement.replaceExpressions(replaceExpression));
     }
 
     private class CaptureReplacer implements Function<ExpressionNode, ExpressionNode> {
-        private final Map<Identifier, ExpressionNode> freeVariablesById;
+        private final Map<ExpressionNode, ExpressionNode> replacements;
 
-        public CaptureReplacer(Map<Identifier, ExpressionNode> freeVariablesById) {
-            this.freeVariablesById = freeVariablesById;
+        public CaptureReplacer(Map<ExpressionNode, ExpressionNode> replacements) {
+            this.replacements = replacements;
         }
 
         @Override
         public ExpressionNode apply(ExpressionNode expression) {
-            return tryCast(VariableReferenceNode.class, expression)
-                .flatMap(variableNode -> Optional.ofNullable(freeVariablesById.get(variableNode.getReferentId())))
+            return Optional.ofNullable(replacements.get(expression))
                 .orElseGet(() -> expression.replaceExpressions(this));
         }
     }
 
-    private ConstructorNode buildConstructor(Scope outerScope, TypeName type, List<VariableDeclaration> freeVariables) {
+    private ConstructorNode buildConstructor(Scope outerScope, TypeName type, List<CapturedVariable> freeVariables) {
         Scope scope = outerScope.enterConstructor();
-        Map<Identifier, VariableDeclaration> argumentDeclarationsById = toMap(
-            freeVariables,
-            freeVariable -> immutableEntry(
-                freeVariable.getId(),
-                scope.generateVariable(freeVariable.getName(), freeVariable.getType())));
 
-        List<FormalArgumentNode> arguments = eagerMap(
-            freeVariables,
-            freeVariable -> formalArg(argumentDeclarationsById.get(freeVariable.getId())));
+        ImmutableList.Builder<FormalArgumentNode> arguments = ImmutableList.builder();
+        ImmutableList.Builder<StatementNode> body = ImmutableList.builder();
 
-        List<StatementNode> body = eagerMap(freeVariables, freeVariable -> assignStatement(
-            captureAccess(type, freeVariable),
-            reference(argumentDeclarationsById.get(freeVariable.getId()))));
+        for (CapturedVariable variable : freeVariables) {
+            FormalArgumentNode argument = scope.formalArgument(variable.field.getName(), variable.field.getType());
+            arguments.add(argument);
+            body.add(assignStatement(captureAccess(type, variable), reference(argument)));
+        }
 
-        return constructor(arguments, body);
+        return constructor(arguments.build(), body.build());
     }
 
-    private FieldAccessNode captureAccess(TypeName type, VariableDeclaration freeVariable) {
-        return fieldAccess(thisReference(type), freeVariable.getName(), freeVariable.getType());
+    private FieldAccessNode captureAccess(TypeName type, CapturedVariable freeVariable) {
+        FieldDeclarationNode field = freeVariable.field;
+        return fieldAccess(thisReference(type), field.getName(), field.getType());
     }
 
     GeneratedClosure readAnonymousClass(Scope outerScope, AnonymousClassDeclaration declaration) {
@@ -170,6 +165,16 @@ public class JavaReader {
         return closure;
     }
 
+    private static class CapturedVariable {
+        private final ReferenceNode freeVariable;
+        private final FieldDeclarationNode field;
+
+        public CapturedVariable(ReferenceNode freeVariable, FieldDeclarationNode field) {
+            this.freeVariable = freeVariable;
+            this.field = field;
+        }
+    }
+
     private GeneratedClosure classWithCapture(
         Scope scope,
         TypeName className,
@@ -177,10 +182,11 @@ public class JavaReader {
         List<FieldDeclarationNode> declaredFields,
         List<MethodNode> methods
     ) {
-        List<VariableDeclaration> freeVariables = findFreeVariables(ExtraLists.concat(declaredFields, methods));
-        Iterable<FieldDeclarationNode> captureFields = transform(
+        List<ReferenceNode> freeVariables = findFreeVariables(ExtraLists.concat(declaredFields, methods));
+        List<CapturedVariable> capturedVariables = ImmutableList.copyOf(transform(
             freeVariables,
-            freeVariable -> field(freeVariable.getName(), freeVariable.getType()));
+            freeVariable -> new CapturedVariable(freeVariable, fieldForCapture(freeVariable))));
+        Iterable<FieldDeclarationNode> captureFields = transform(capturedVariables, capture -> capture.field);
 
         List<FieldDeclarationNode> fields = ImmutableList.copyOf(concat(declaredFields, captureFields));
 
@@ -188,12 +194,28 @@ public class JavaReader {
             className,
             superTypes,
             fields,
-            buildConstructor(scope, className, freeVariables),
+            buildConstructor(scope, className, capturedVariables),
             eagerMap(methods, method ->
-                method.mapBody(body -> replaceCaptureReferences(className, body, freeVariables))));
+                method.mapBody(body -> replaceCaptureReferences(className, body, capturedVariables))));
         return new GeneratedClosure(classNode, freeVariables);
     }
 
+    private FieldDeclarationNode fieldForCapture(ReferenceNode freeVariable) {
+        return freeVariable.accept(new ReferenceNode.Visitor<FieldDeclarationNode>() {
+            @Override
+            public FieldDeclarationNode visit(VariableReferenceNode reference) {
+                return field(reference.getReferent().getName(), reference.getType());
+            }
+
+            @Override
+            public FieldDeclarationNode visit(ThisReferenceNode thisReference) {
+                TypeName type = thisReference.getType();
+                String name = "this_" + type.getQualifiedName().replace(".", "__");
+                return field(name, type);
+            }
+        });
+    }
+
     private TypeName generateAnonymousName(ASTNode node) {
         ITypeBinding type = findDeclaringClass(node);
         while (type.isAnonymous()) {