Extract inner type hoisting to separate step

6 files changed, 74 insertions, 185 deletions

diff --git a/src/main/java/org/zwobble/couscous/frontends/java/AnonymousType.java b/src/main/java/org/zwobble/couscous/frontends/java/AnonymousType.java
index 845c6f9..5d5915d 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/AnonymousType.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/AnonymousType.java
@@ -2,6 +2,8 @@ package org.zwobble.couscous.frontends.java;
 
 import org.zwobble.couscous.types.Type;
 
+import java.util.function.Function;
+
 public class AnonymousType implements Type {
     public static Type anonymousType(String key) {
         return new AnonymousType(key);
@@ -19,6 +21,11 @@ public class AnonymousType implements Type {
     }
 
     @Override
+    public Type transformSubTypes(Function<Type, Type> transform) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
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 b904f0d..7679431 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/FreeVariables.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/FreeVariables.java
@@ -14,8 +14,6 @@ import java.util.stream.Stream;
 import static org.zwobble.couscous.ast.structure.NodeStructure.descendantNodesAndSelf;
 import static org.zwobble.couscous.util.Casts.tryCast;
 import static org.zwobble.couscous.util.ExtraIterables.iterable;
-import static org.zwobble.couscous.util.ExtraIterables.lazyFlatMap;
-import static org.zwobble.couscous.util.ExtraLists.list;
 import static org.zwobble.couscous.util.ExtraStreams.toStream;
 
 public class FreeVariables {
@@ -50,32 +48,52 @@ public class FreeVariables {
         }
     }
 
+    public static class FindDeclaredTypes {
+        private static final Function<Node, Stream<TypeParameter>> VISITOR =
+            DynamicNodeMapper.instantiate(new FindDeclaredTypes(), "visit");
+
+        public Stream<TypeParameter> visit(Node node) {
+            return Stream.empty();
+        }
+
+        public Stream<TypeParameter> visit(MethodNode methodNode) {
+            return methodNode.getTypeParameters().stream().map(parameter -> parameter.getType());
+        }
+    }
+
     public static Set<TypeParameter> findFreeTypeParameters(Node root) {
+        // TODO: test removal of declared types
+        Set<TypeParameter> declaredTypes = descendantNodesAndSelf(root)
+            .flatMap(FindDeclaredTypes.VISITOR)
+            .collect(Collectors.toSet());
+
         // TODO: this assumes that *all* type parameters are free, which is often but not always true in inner types
-        Iterable<Type> types = iterable(() -> descendantNodesAndSelf(root)
-            .flatMap(FindReferencedTypes.VISITOR));
-
-        return ImmutableSet.copyOf(lazyFlatMap(types, type -> type.accept(new Type.Visitor<Iterable<TypeParameter>>() {
-            @Override
-            public Iterable<TypeParameter> visit(ScalarType type) {
-                return list();
-            }
-
-            @Override
-            public Iterable<TypeParameter> visit(TypeParameter parameter) {
-                return list(parameter);
-            }
-
-            @Override
-            public Iterable<TypeParameter> visit(ParameterizedType type) {
-                return lazyFlatMap(type.getParameters(), parameter -> parameter.accept(this));
-            }
-
-            @Override
-            public Iterable<TypeParameter> visit(BoundTypeParameter type) {
-                return type.getValue().accept(this);
-            }
-        })));
+        Stream<TypeParameter> types = descendantNodesAndSelf(root)
+            .flatMap(FindReferencedTypes.VISITOR)
+            .flatMap(type -> type.accept(new Type.Visitor<Stream<TypeParameter>>() {
+                @Override
+                public Stream<TypeParameter> visit(ScalarType type) {
+                    return Stream.empty();
+                }
+
+                @Override
+                public Stream<TypeParameter> visit(TypeParameter parameter) {
+                    return Stream.of(parameter);
+                }
+
+                @Override
+                public Stream<TypeParameter> visit(ParameterizedType type) {
+                    return type.getParameters().stream().flatMap(parameter -> parameter.accept(this));
+                }
+
+                @Override
+                public Stream<TypeParameter> visit(BoundTypeParameter type) {
+                    return type.getValue().accept(this);
+                }
+            }))
+            .filter(type -> !declaredTypes.contains(type));
+
+        return ImmutableSet.copyOf(iterable(() -> types));
     }
 
     public static List<ReferenceNode> findFreeVariables(List<? extends Node> body) {
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 9bad1d2..c019cb8 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
@@ -40,6 +40,10 @@ public class GeneratedClosure {
             concat(capturedVariables, arguments));
     }
 
+    public boolean hasCaptures() {
+        return !capturedTypes.isEmpty() || !capturedVariables.isEmpty();
+    }
+
     private ScalarType getType() {
         return classNode.getName();
     }
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 00198a6..aa64543 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaExpressionReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaExpressionReader.java
@@ -290,6 +290,7 @@ public class JavaExpressionReader {
                 return methodCall(
                     thisReference(typeOf(declaringClass.getDeclaringClass())),
                     // TODO: remove duplication with JavaReader in name generation
+                    // TODO: handle this when hoisting (i.e. in HoistNestedTypes)
                     "create_" + declaringClass.getName(),
                     arguments,
                     type);
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 5202127..a047e9e 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
@@ -1,8 +1,6 @@
 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.*;
 import org.zwobble.couscous.ast.*;
@@ -11,21 +9,19 @@ 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.*;
+import org.zwobble.couscous.transforms.ClosureGenerator;
+import org.zwobble.couscous.types.ScalarType;
 import org.zwobble.couscous.types.Type;
 import org.zwobble.couscous.util.ExtraLists;
-import org.zwobble.couscous.util.InsertionOrderSet;
 
 import java.io.IOException;
 import java.nio.file.Path;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
-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;
@@ -37,17 +33,12 @@ import static org.zwobble.couscous.ast.FieldDeclarationNode.field;
 import static org.zwobble.couscous.ast.FormalArgumentNode.formalArg;
 import static org.zwobble.couscous.ast.FormalTypeParameterNode.formalTypeParameter;
 import static org.zwobble.couscous.ast.InstanceReceiver.instanceReceiver;
-import static org.zwobble.couscous.ast.ReturnNode.returns;
 import static org.zwobble.couscous.ast.StaticReceiver.staticReceiver;
 import static org.zwobble.couscous.ast.ThisReferenceNode.thisReference;
-import static org.zwobble.couscous.ast.VariableReferenceNode.reference;
-import static org.zwobble.couscous.frontends.java.FreeVariables.findFreeTypeParameters;
-import static org.zwobble.couscous.frontends.java.FreeVariables.findFreeVariables;
 import static org.zwobble.couscous.frontends.java.JavaMethods.signature;
 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.ExtraIterables.lazyMap;
 import static org.zwobble.couscous.util.ExtraLists.*;
 import static org.zwobble.couscous.util.ExtraMaps.map;
 
@@ -125,7 +116,8 @@ public class JavaReader {
                 superTypes,
                 body.getFields(),
                 body.getStaticConstructor(),
-                body.getMethods()
+                body.getMethods(),
+                body.getInnerTypes()
             );
         } else {
             return ClassNode.declareClass(
@@ -135,7 +127,9 @@ public class JavaReader {
                 body.getFields(),
                 body.getStaticConstructor(),
                 body.getConstructor(),
-                body.getMethods());
+                body.getMethods(),
+                body.getInnerTypes()
+            );
         }
     }
 
@@ -197,46 +191,6 @@ public class JavaReader {
             list(method));
     }
 
-    private List<StatementNode> replaceCaptureReferences(
-        ScalarType className,
-        List<StatementNode> body,
-        InsertionOrderSet<CapturedVariable> freeVariables)
-    {
-        Map<ExpressionNode, ExpressionNode> replacements = Maps.transformValues(
-            Maps.uniqueIndex(freeVariables, variable -> variable.freeVariable),
-            freeVariable -> captureAccess(className, freeVariable));
-        NodeTransformer transformer = NodeTransformer.replaceExpressions(replacements);
-        return eagerFlatMap(body, transformer::transformStatement);
-    }
-
-    private ConstructorNode buildConstructor(
-        Scope outerScope,
-        ScalarType type,
-        InsertionOrderSet<CapturedVariable> freeVariables,
-        ConstructorNode existing)
-    {
-        Scope scope = outerScope.enterConstructor();
-
-        ImmutableList.Builder<FormalArgumentNode> arguments = ImmutableList.builder();
-        ImmutableList.Builder<StatementNode> body = ImmutableList.builder();
-
-        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)));
-        }
-
-        arguments.addAll(existing.getArguments());
-        body.addAll(existing.getBody());
-
-        return new ConstructorNode(arguments.build(), body.build());
-    }
-
-    private FieldAccessNode captureAccess(ScalarType type, CapturedVariable freeVariable) {
-        FieldDeclarationNode field = freeVariable.field;
-        return fieldAccess(thisReference(type), field.getName(), field.getType());
-    }
-
     GeneratedClosure readAnonymousClass(Scope outerScope, AnonymousClassDeclaration declaration) {
         ScalarType className = generateAnonymousName(declaration);
         Scope scope = outerScope.enterClass(className);
@@ -255,61 +209,6 @@ 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,
-        ClassNode classNode
-    ) {
-        InsertionOrderSet<org.zwobble.couscous.types.TypeParameter> freeTypeParameters = InsertionOrderSet.copyOf(
-            findFreeTypeParameters(classNode));
-
-        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)));
-        InsertionOrderSet<CapturedVariable> capturedVariables = InsertionOrderSet.copyOf(transform(
-            freeVariables,
-            freeVariable -> new CapturedVariable(freeVariable, fieldForCapture(freeVariable))));
-        Iterable<FieldDeclarationNode> captureFields = transform(capturedVariables, capture -> capture.field);
-
-        List<FieldDeclarationNode> fields = ImmutableList.copyOf(concat(
-            classNode.getFields(),
-            captureFields));
-
-        if (!classNode.getStaticConstructor().isEmpty()) {
-            throw new RuntimeException("Class has unexpected static constructor");
-        }
-
-        ClassNode generatedClass = new ClassNode(
-            classNode.getName(),
-            // TODO: generate fresh type parameter and replace in body
-            ExtraLists.concat(lazyMap(freeTypeParameters, parameter -> formalTypeParameter(parameter)), classNode.getTypeParameters()),
-            classNode.getSuperTypes(),
-            fields,
-            list(),
-            buildConstructor(scope, classNode.getName(), capturedVariables, classNode.getConstructor()),
-            eagerMap(classNode.getMethods(), method ->
-                method.mapBody(body -> replaceCaptureReferences(classNode.getName(), body, capturedVariables))));
-        return new GeneratedClosure(generatedClass, freeTypeParameters, freeVariables);
-    }
-
-    private boolean isThisReference(ScalarType name, ReferenceNode reference) {
-        return tryCast(ThisReferenceNode.class, reference)
-            .map(node -> node.getType().equals(name))
-            .orElse(false);
-    }
-
     private GeneratedClosure classWithCapture(
         Scope scope,
         ScalarType className,
@@ -322,7 +221,9 @@ public class JavaReader {
             anonymousClass.getFields(),
             list(),
             ConstructorNode.DEFAULT,
-            anonymousClass.getMethods());
+            anonymousClass.getMethods(),
+            list()
+        );
         if (anonymousClass.getType().isPresent()) {
             NodeTransformer nodeTransformer = NodeTransformer.replaceExpressions(map(
                 thisReference(anonymousClass.getType().get()),
@@ -330,23 +231,7 @@ public class JavaReader {
             ));
             classNode = nodeTransformer.transformClass(classNode);
         }
-        return classWithCapture(scope, classNode);
-    }
-
-    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) {
-                Type type = thisReference.getType();
-                String name = "this_" + erasure(type).getQualifiedName().replace(".", "__");
-                return field(name, type);
-            }
-        });
+        return ClosureGenerator.classWithCapture(scope, classNode);
     }
 
     private ScalarType generateAnonymousName(ASTNode node) {
@@ -380,39 +265,11 @@ public class JavaReader {
                 .ifPresent(field -> readField(body, scope, type, field, isInterface));
 
             tryCast(TypeDeclaration.class, declaration)
-                .ifPresent(typeDeclaration -> classes.add(readNestedTypeDeclaration(body, typeDeclaration)));
+                .ifPresent(typeDeclaration -> body.addInnerType(readTypeDeclaration(typeDeclaration)));
         }
         return body.build();
     }
 
-    private TypeNode readNestedTypeDeclaration(TypeDeclarationBody.Builder body, TypeDeclaration typeDeclaration) {
-        TypeNode typeNode = readTypeDeclaration(typeDeclaration);
-        // TODO: can we remove duplication of scope creation with readTypeDeclaration()?
-        Scope scope = topScope.enterClass(typeNode.getName());
-        return tryCast(ClassNode.class, typeNode)
-            .filter(node -> !Modifier.isStatic(typeDeclaration.getModifiers()))
-            .<TypeNode>map(node -> {
-                GeneratedClosure closure = classWithCapture(scope, node);
-//                Type type = node.getTypeParameters().isEmpty()
-//                    ? typeNode.getName()
-//                    : parameterizedType(typeNode.getName(), eagerMap(node.getTypeParameters()));
-                Type type = typeNode.getName();
-                // TODO: method type parameters
-                List<FormalArgumentNode> methodArguments = eagerMap(
-                    node.getConstructor().getArguments(),
-                    // TODO: use scope of outer class
-                    argument -> scope.formalArgument(argument.getName(), argument.getType()));
-                MethodNode method = MethodNode.builder("create_" + typeDeclaration.getName().getIdentifier())
-                    .arguments(methodArguments)
-                    .returns(type)
-                    .statement(returns(closure.generateConstructor(lazyMap(methodArguments, argument -> reference(argument)))))
-                    .build();
-                body.addMethod(method);
-                return closure.getClassNode();
-            })
-            .orElse(typeNode);
-    }
-
     private void readField(TypeDeclarationBody.Builder builder, Scope scope, ScalarType declaringType, FieldDeclaration field, boolean isInterface) {
         @SuppressWarnings("unchecked")
         List<VariableDeclarationFragment> fragments = field.fragments();
diff --git a/src/main/java/org/zwobble/couscous/frontends/java/JavaTypes.java b/src/main/java/org/zwobble/couscous/frontends/java/JavaTypes.java
index aa20fc9..5d9eddd 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaTypes.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaTypes.java
@@ -47,10 +47,12 @@ public class JavaTypes {
             Identifier declaringScope = getDeclaringScope(typeBinding);
             return new TypeParameter(declaringScope, typeBinding.getName());
         }
-        ScalarType rawType = outerClass == null
-            ? ScalarType.of(typeBinding.getErasure().getQualifiedName())
-            // TODO: test for erasure of inner type name
-            : ScalarType.of(erasure(typeOf(outerClass)).getQualifiedName() + "__" + typeBinding.getErasure().getName());
+        // TODO: tidy-up
+//        ScalarType rawType = outerClass == null
+//            ? ScalarType.of(typeBinding.getErasure().getQualifiedName())
+//            // TODO: test for erasure of inner type name
+//            : ScalarType.of(erasure(typeOf(outerClass)).getQualifiedName() + "__" + typeBinding.getErasure().getName());
+        ScalarType rawType = ScalarType.of(typeBinding.getErasure().getQualifiedName());
         if (typeBinding.isParameterizedType()) {
             List<Type> typeParameters = eagerMap(
                 asList(typeBinding.getTypeArguments()),