Fix generation of hoisted type constructor

1 files changed, 7 insertions, 1 deletions

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(