Remove duplicates when finding captured variables

4 files changed, 45 insertions, 17 deletions

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 f9e6415..4629fa5 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
@@ -3,14 +3,13 @@ package org.zwobble.couscous.frontends.java;
 import org.zwobble.couscous.ast.ClassNode;
 import org.zwobble.couscous.ast.ReferenceNode;
 import org.zwobble.couscous.types.ScalarType;
-
-import java.util.List;
+import org.zwobble.couscous.util.InsertionOrderSet;
 
 public class GeneratedClosure {
     private final ClassNode classNode;
-    private final List<ReferenceNode> captures;
+    private final InsertionOrderSet<ReferenceNode> captures;
 
-    public GeneratedClosure(ClassNode classNode, List<ReferenceNode> captures) {
+    public GeneratedClosure(ClassNode classNode, InsertionOrderSet<ReferenceNode> captures) {
         this.classNode = classNode;
         this.captures = captures;
     }
@@ -23,7 +22,7 @@ public class GeneratedClosure {
         return classNode.getName();
     }
 
-    public List<ReferenceNode> getCaptures() {
+    public InsertionOrderSet<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 4adb92c..71826fe 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaExpressionReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaExpressionReader.java
@@ -5,6 +5,7 @@ import org.eclipse.jdt.core.dom.*;
 import org.zwobble.couscous.ast.*;
 import org.zwobble.couscous.types.Type;
 import org.zwobble.couscous.types.Types;
+import org.zwobble.couscous.util.InsertionOrderSet;
 
 import java.util.List;
 import java.util.stream.Collectors;
@@ -229,17 +230,17 @@ public class JavaExpressionReader {
         GeneratedClosure closure = javaReader.readLambda(scope, expression);
         return constructorCall(
             closure.getType(),
-            captureArguments(closure));
+            captureArguments(closure).asList());
     }
 
     private ExpressionNode readExpressionMethodReference(ExpressionMethodReference expression) {
         GeneratedClosure closure = javaReader.readExpressionMethodReference(scope, expression);
         return constructorCall(
             closure.getType(),
-            captureArguments(closure));
+            captureArguments(closure).asList());
     }
 
-    private List<? extends ExpressionNode> captureArguments(GeneratedClosure closure) {
+    private InsertionOrderSet<? extends ExpressionNode> captureArguments(GeneratedClosure closure) {
         return closure.getCaptures();
     }
 
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 a6da3c3..68a70ed 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
@@ -1,6 +1,7 @@
 package org.zwobble.couscous.frontends.java;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import org.eclipse.jdt.core.compiler.IProblem;
 import org.eclipse.jdt.core.dom.*;
@@ -8,11 +9,12 @@ import org.zwobble.couscous.ast.*;
 import org.zwobble.couscous.ast.sugar.AnonymousClass;
 import org.zwobble.couscous.ast.sugar.Lambda;
 import org.zwobble.couscous.ast.sugar.TypeDeclarationBody;
+import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.types.ScalarType;
 import org.zwobble.couscous.types.Type;
-import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.types.Types;
 import org.zwobble.couscous.util.ExtraLists;
+import org.zwobble.couscous.util.InsertionOrderSet;
 
 import java.io.IOException;
 import java.nio.file.Path;
@@ -37,15 +39,14 @@ import static org.zwobble.couscous.ast.StaticReceiver.staticReceiver;
 import static org.zwobble.couscous.ast.ThisReferenceNode.thisReference;
 import static org.zwobble.couscous.ast.TypeCoercionNode.typeCoercion;
 import static org.zwobble.couscous.ast.VariableReferenceNode.reference;
-import static org.zwobble.couscous.types.Types.erasure;
 import static org.zwobble.couscous.frontends.java.FreeVariables.findFreeVariables;
 import static org.zwobble.couscous.frontends.java.JavaTypes.*;
+import static org.zwobble.couscous.types.Types.erasure;
 import static org.zwobble.couscous.util.Casts.tryCast;
 import static org.zwobble.couscous.util.ExtraLists.*;
 import static org.zwobble.couscous.util.ExtraMaps.lookup;
 
 public class JavaReader {
-    // TODO: de-dupe captures (multiple reference nodes may reference the same free variable)
     public static List<TypeNode> readClassFromFile(List<Path> sourcePaths, Path sourcePath) throws IOException {
         CompilationUnit ast = new JavaParser().parseCompilationUnit(sourcePaths, sourcePath);
 
@@ -70,7 +71,7 @@ public class JavaReader {
     private final ImmutableList.Builder<TypeNode> classes;
     private int anonymousClassCount = 0;
     private final Scope topScope = Scope.create();
-    private final Map<ScalarType, List<ReferenceNode>> nestedClasses = new HashMap<>();
+    private final Map<ScalarType, InsertionOrderSet<ReferenceNode>> nestedClasses = new HashMap<>();
 
     private JavaReader() {
         classes = ImmutableList.builder();
@@ -160,7 +161,7 @@ public class JavaReader {
     private List<StatementNode> replaceCaptureReferences(
         ScalarType className,
         List<StatementNode> body,
-        List<CapturedVariable> freeVariables)
+        InsertionOrderSet<CapturedVariable> freeVariables)
     {
         Map<ExpressionNode, ExpressionNode> replacements = Maps.transformValues(
             Maps.uniqueIndex(freeVariables, variable -> variable.freeVariable),
@@ -172,7 +173,7 @@ public class JavaReader {
     private ConstructorNode buildConstructor(
         Scope outerScope,
         ScalarType type,
-        List<CapturedVariable> freeVariables,
+        InsertionOrderSet<CapturedVariable> freeVariables,
         ConstructorNode existing)
     {
         Scope scope = outerScope.enterConstructor();
@@ -228,13 +229,13 @@ public class JavaReader {
         Scope scope,
         ClassNode classNode
     ) {
-        List<ReferenceNode> freeVariables = eagerFilter(
+        InsertionOrderSet<ReferenceNode> freeVariables = InsertionOrderSet.copyOf(Iterables.filter(
             findFreeVariables(ExtraLists.concat(
                 classNode.getFields(),
                 classNode.getMethods())),
             // TODO: check this references work correctly in anonymous classes
-            variable -> !isThisReference(classNode.getName(), variable));
-        List<CapturedVariable> capturedVariables = ImmutableList.copyOf(transform(
+            variable -> !isThisReference(classNode.getName(), variable)));
+        InsertionOrderSet<CapturedVariable> capturedVariables = InsertionOrderSet.copyOf(transform(
             freeVariables,
             freeVariable -> new CapturedVariable(freeVariable, fieldForCapture(freeVariable))));
         Iterable<FieldDeclarationNode> captureFields = transform(capturedVariables, capture -> capture.field);
diff --git a/src/main/java/org/zwobble/couscous/util/InsertionOrderSet.java b/src/main/java/org/zwobble/couscous/util/InsertionOrderSet.java
new file mode 100644
index 0000000..4272add
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/util/InsertionOrderSet.java
@@ -0,0 +1,27 @@
+package org.zwobble.couscous.util;
+
+import java.util.Iterator;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+public interface InsertionOrderSet<T> extends Iterable<T> {
+    static <T> InsertionOrderSet<T> copyOf(Iterable<T> iterable) {
+        List<T> list = StreamSupport.stream(iterable.spliterator(), false)
+            .distinct()
+            .collect(Collectors.toList());
+        return new InsertionOrderSet<T>() {
+            @Override
+            public List<T> asList() {
+                return list;
+            }
+
+            @Override
+            public Iterator<T> iterator() {
+                return list.iterator();
+            }
+        };
+    }
+
+    List<T> asList();
+}