2 files changed, 23 insertions, 20 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 224283f..88d214c 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
@@ -1,21 +1,26 @@
 package org.zwobble.couscous.frontends.java;
 
+import org.zwobble.couscous.ast.ClassNode;
 import org.zwobble.couscous.ast.TypeName;
 import org.zwobble.couscous.ast.VariableDeclaration;
 
 import java.util.List;
 
 public class GeneratedClosure {
-    private final TypeName type;
+    private final ClassNode classNode;
     private final List<VariableDeclaration> captures;
 
-    public GeneratedClosure(TypeName type, List<VariableDeclaration> captures) {
-        this.type = type;
+    public GeneratedClosure(ClassNode classNode, List<VariableDeclaration> captures) {
+        this.classNode = classNode;
         this.captures = captures;
     }
 
+    public ClassNode getClassNode() {
+        return classNode;
+    }
+
     public TypeName getType() {
-        return type;
+        return classNode.getName();
     }
 
     public List<VariableDeclaration> 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 9b5f561..76dcf22 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
@@ -1,11 +1,11 @@
 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.dom.*;
 import org.zwobble.couscous.ast.*;
 import org.zwobble.couscous.ast.VariableDeclaration;
+import org.zwobble.couscous.util.ExtraLists;
 
 import java.io.IOException;
 import java.nio.file.Path;
@@ -14,6 +14,7 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.function.Function;
 
+import static com.google.common.collect.Iterables.concat;
 import static com.google.common.collect.Iterables.transform;
 import static java.util.Arrays.asList;
 import static java.util.Collections.emptyList;
@@ -66,8 +67,6 @@ public class JavaReader {
         IMethodBinding functionalInterfaceMethod = expression.resolveTypeBinding().getFunctionalInterfaceMethod();
         TypeName className = generateClassName(expression);
         FunctionDeclaration lambda = functionDeclaration(expression);
-        List<VariableDeclaration> freeVariables = findFreeVariables(
-            concat(lambda.getFormalArguments(), lambda.getBody()));
 
         MethodNode method = MethodNode.method(
             emptyList(),
@@ -76,10 +75,10 @@ public class JavaReader {
             lambda.getFormalArguments(),
             lambda.getBody());
 
-        ClassNode classNode = classWithCapture(className, freeVariables, emptyList(), ImmutableList.of(method));
+        GeneratedClosure closure = classWithCapture(className, emptyList(), ImmutableList.of(method));
 
-        classes.add(classNode);
-        return new GeneratedClosure(classNode.getName(), freeVariables);
+        classes.add(closure.getClassNode());
+        return closure;
     }
 
     private List<StatementNode> replaceCaptureReferences(
@@ -133,31 +132,30 @@ public class JavaReader {
     GeneratedClosure readAnonymousClass(AnonymousClassDeclaration declaration) {
         TypeName className = generateClassName(declaration);
         TypeDeclarationBody bodyDeclarations = readTypeDeclarationBody(declaration.bodyDeclarations());
-        List<VariableDeclaration> freeVariables = findFreeVariables(bodyDeclarations.getNodes());
-
-        ClassNode classNode = classWithCapture(className, freeVariables, bodyDeclarations.getFields(), bodyDeclarations.getMethods());
-        classes.add(classNode);
-        return new GeneratedClosure(classNode.getName(), freeVariables);
+        GeneratedClosure closure = classWithCapture(className, bodyDeclarations.getFields(), bodyDeclarations.getMethods());
+        classes.add(closure.getClassNode());
+        return closure;
     }
 
-    private ClassNode classWithCapture(
+    private GeneratedClosure classWithCapture(
         TypeName className,
-        List<VariableDeclaration> freeVariables,
         List<FieldDeclarationNode> declaredFields,
         List<MethodNode> methods
     ) {
+        List<VariableDeclaration> freeVariables = findFreeVariables(ExtraLists.concat(declaredFields, methods));
         Iterable<FieldDeclarationNode> captureFields = transform(
             freeVariables,
             freeVariable -> field(freeVariable.getName(), freeVariable.getType()));
 
-        List<FieldDeclarationNode> fields = ImmutableList.copyOf(Iterables.concat(declaredFields, captureFields));
+        List<FieldDeclarationNode> fields = ImmutableList.copyOf(concat(declaredFields, captureFields));
 
-        return ClassNode.declareClass(
+        ClassNode classNode = ClassNode.declareClass(
             className,
             fields,
             buildConstructor(className, freeVariables),
             eagerMap(methods, method ->
                 method.mapBody(body -> replaceCaptureReferences(className, body, freeVariables))));
+        return new GeneratedClosure(classNode, freeVariables);
     }
 
     private TypeName generateClassName(AnonymousClassDeclaration declaration) {
@@ -196,7 +194,7 @@ public class JavaReader {
         }
 
         public List<Node> getNodes() {
-            return concat(fields, asList(constructor), methods);
+            return ExtraLists.concat(fields, asList(constructor), methods);
         }
     }