Create type constructor for generic sum type definitions

3 files changed, 57 insertions, 17 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedSumTypeDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedSumTypeDefinitionNode.java
index d727b98e..3f64c4ee 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedSumTypeDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedSumTypeDefinitionNode.java
@@ -50,6 +50,17 @@ public record UntypedSumTypeDefinitionNode(
             );
         }
 
+        public Builder withTypeLevelParams(
+            List<UntypedTypeLevelParamNode> typeLevelParams
+        ) {
+            return new Builder(
+                this.name,
+                this.visibility,
+                Optional.of(typeLevelParams),
+                this.variants
+            );
+        }
+
         public Builder addVariant(UntypedTypeLevelExpressionNode variant) {
             return new Builder(
                 this.name,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 4f428c1f..15e740b5 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1491,17 +1491,11 @@ public class TypeChecker {
 
                 var type = new StructType(namespaceName, untyped.name());
 
-                var typeOrConstructor = untyped.typeLevelParams()
-                    .map(typeLevelParamNodes -> TypeOrConstructor.constructor(new TypeConstructor<>(
-                        typeLevelParamNodes.stream()
-                            .map(typeLevelParamNode -> new TypeParam(
-                                new Scope.Struct(type),
-                                typeLevelParamNode.name()
-                            ))
-                            .toList(),
-                        type
-                    )))
-                    .orElseGet(() -> TypeOrConstructor.type(type));
+                var typeOrConstructor = createTypeOrConstructor(
+                    untyped.typeLevelParams(),
+                    type,
+                    new Scope.Struct(type)
+                );
 
                 context = context.addLexicalState(untyped.name(), typeOrConstructor);
 
@@ -1585,14 +1579,27 @@ public class TypeChecker {
 
                 var type = new SumType(namespaceName, untyped.name());
 
-                context = context.addLexicalState(untyped.name(), type);
+                var typeOrConstructor = createTypeOrConstructor(
+                    untyped.typeLevelParams(),
+                    type,
+                    new Scope.Sum(type)
+                );
 
-                return context.declareVariableType(untyped.name(), Types.metaType(type), untyped.source());
+                context = context.addLexicalState(untyped.name(), typeOrConstructor);
+
+                return context.declareVariableType(
+                    untyped.name(),
+                    typeOrConstructor.toMetaType(),
+                    untyped.source()
+                );
             }
 
             @Override
             public TypeCheckerContext defineTypes(TypeCheckerContext context) {
-                var type = context.getLexicalState(SumType.class, untyped.name());
+                var typeOrConstructor = (TypeOrConstructor<SumType>) context.getLexicalState(
+                    TypeOrConstructor.class,
+                    untyped.name()
+                );
 
                 // TODO: restrict variants to structs. Do we allow other sum types?
                 var variants = untyped.variants().stream()
@@ -1613,15 +1620,18 @@ public class TypeChecker {
                     })
                     .collect(PCollectors.toOrderedSet());
 
-                return context.addSumTypeInfo(type, new SumTypeInfo(variants));
+                return context.addSumTypeInfo(typeOrConstructor.toType(), new SumTypeInfo(variants));
             }
 
             @Override
             public TypeCheckerResult<TypedNamespaceStatementNode> finish(TypeCheckerContext context) {
-                var type = context.getLexicalState(SumType.class, untyped.name());
+                var typeOrConstructor = (TypeOrConstructor<SumType>) context.getLexicalState(
+                    TypeOrConstructor.class,
+                    untyped.name()
+                );
 
                 return context.toResult(new TypedSumTypeDefinitionNode(
-                    TypeOrConstructor.type(type),
+                    typeOrConstructor,
                     untyped.source()
                 ));
             }
@@ -2159,4 +2169,22 @@ public class TypeChecker {
             untyped.source()
         ));
     }
+
+    private static <T extends StructuredType> TypeOrConstructor<T> createTypeOrConstructor(
+        Optional<List<UntypedTypeLevelParamNode>> untypedTypeLevelParamNodes,
+        T type,
+        Scope scope
+    ) {
+        return untypedTypeLevelParamNodes
+            .map(typeLevelParamNodes -> TypeOrConstructor.constructor(new TypeConstructor<>(
+                typeLevelParamNodes.stream()
+                    .map(typeLevelParamNode -> new TypeParam(
+                        scope,
+                        typeLevelParamNode.name()
+                    ))
+                    .toList(),
+                type
+            )))
+            .orElseGet(() -> TypeOrConstructor.type(type));
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/types/Scope.java b/src/main/java/org/shedlang/compiler/types/Scope.java
index 0ba3c776..ed8e8c5a 100644
--- a/src/main/java/org/shedlang/compiler/types/Scope.java
+++ b/src/main/java/org/shedlang/compiler/types/Scope.java
@@ -2,4 +2,5 @@ package org.shedlang.compiler.types;
 
 public sealed interface Scope {
     record Struct(StructType type) implements Scope {}
+    record Sum(SumType type) implements Scope {}
 }