Fix generation of hoisted type constructor

3 files changed, 16 insertions, 1 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/MethodNode.java b/src/main/java/org/zwobble/couscous/ast/MethodNode.java
index ccef0da..63a6c46 100644
--- a/src/main/java/org/zwobble/couscous/ast/MethodNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/MethodNode.java
@@ -55,6 +55,11 @@ public class MethodNode implements Node {
             return this;
         }
 
+        public Builder typeParameters(Iterable<FormalTypeParameterNode> additionalParameters) {
+            parameters.addAll(additionalParameters);
+            return this;
+        }
+
         public Builder typeParameter(FormalTypeParameterNode parameter) {
             parameters.add(parameter);
             return this;
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 c019cb8..0bf7649 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/GeneratedClosure.java
@@ -47,4 +47,8 @@ public class GeneratedClosure {
     private ScalarType getType() {
         return classNode.getName();
     }
+
+    public InsertionOrderSet<TypeParameter> getCapturedTypes() {
+        return capturedTypes;
+    }
 }
diff --git a/src/main/java/org/zwobble/couscous/transforms/HoistNestedTypes.java b/src/main/java/org/zwobble/couscous/transforms/HoistNestedTypes.java
index 1cb74c8..9d252d2 100644
--- a/src/main/java/org/zwobble/couscous/transforms/HoistNestedTypes.java
+++ b/src/main/java/org/zwobble/couscous/transforms/HoistNestedTypes.java
@@ -7,16 +7,19 @@ import org.zwobble.couscous.frontends.java.Scope;
 import org.zwobble.couscous.types.ScalarType;
 import org.zwobble.couscous.types.Type;
 import org.zwobble.couscous.types.Types;
+import org.zwobble.couscous.util.ExtraLists;
 import org.zwobble.couscous.util.ExtraSets;
 
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 
+import static org.zwobble.couscous.ast.FormalTypeParameterNode.formalTypeParameter;
 import static org.zwobble.couscous.ast.MethodCallNode.methodCall;
 import static org.zwobble.couscous.ast.ReturnNode.returns;
 import static org.zwobble.couscous.ast.ThisReferenceNode.thisReference;
 import static org.zwobble.couscous.ast.VariableReferenceNode.reference;
+import static org.zwobble.couscous.types.Types.addTypeParameters;
 import static org.zwobble.couscous.util.Casts.tryCast;
 import static org.zwobble.couscous.util.ExtraIterables.*;
 import static org.zwobble.couscous.util.ExtraLists.*;
@@ -152,9 +155,12 @@ public class HoistNestedTypes {
                         node.getConstructor().getArguments(),
                         // TODO: use scope of outer class
                         argument -> scope.formalArgument(argument.getName(), argument.getType()));
+                    // TODO: test addition of type parameters (both captured and on type)
+                    List<FormalTypeParameterNode> typeParameters = ExtraLists.concat(lazyMap(closure.getCapturedTypes(), parameter -> formalTypeParameter(parameter)), node.getTypeParameters());
                     MethodNode method = MethodNode.builder(createMethodName(type))
                         .arguments(methodArguments)
-                        .returns(type)
+                        .typeParameters(typeParameters)
+                        .returns(addTypeParameters(type, eagerMap(typeParameters, parameter -> parameter.getType())))
                         .statement(returns(closure.generateConstructor(lazyMap(methodArguments, argument -> reference(argument)))))
                         .build();
                     return Optional.of(new HoistedType(