Add type level params to environment for sum variants

8 files changed, 95 insertions, 26 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index f7cf63b1..89a2f926 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1,5 +1,6 @@
 package org.shedlang.compiler.typechecker;
 
+import org.pcollections.OrderedPSet;
 import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.ast.untyped.*;
 import org.shedlang.compiler.sources.Source;
@@ -1513,18 +1514,7 @@ public class TypeChecker {
                     untyped.name()
                 );
 
-                var environment = EnvironmentBuilder.EMPTY;
-
-                if (typeOrConstructor instanceof TypeOrConstructor.Constructor<StructType> constructor) {
-                    for (var param : constructor.value().params()) {
-                        // TODO: proper source
-                        environment = environment.add(
-                            param.name(),
-                            Variable.defined(Types.metaType(param), VariableType.CONSTANT),
-                            Sources.NULL
-                        );
-                    }
-                }
+                var environment = environmentForTypeOrConstructor(typeOrConstructor);
 
                 context = context.enterBlock(environment.build());
 
@@ -1601,10 +1591,23 @@ public class TypeChecker {
                     untyped.name()
                 );
 
+                var environment = environmentForTypeOrConstructor(typeOrConstructor);
+
+                context = context.enterBlock(environment.build());
+
+                var variants = typeCheckVariants(context);
+
+                context = context.exit();
+
+                return context.addSumTypeInfo(typeOrConstructor.toType(), new SumTypeInfo(variants));
+            }
+
+            private OrderedPSet<SumVariant> typeCheckVariants(TypeCheckerContext context) {
                 // TODO: restrict variants to structs. Do we allow other sum types?
                 // TODO: prevent constructed types with same underlying type from appearing multiple times
                 // (e.g. Option[String], Option[Int])
-                var variants = untyped.variants().stream()
+
+                return untyped.variants().stream()
                     .map(untypedVariant -> {
                         var variantType = typeCheckTypeExpression(
                             untypedVariant,
@@ -1612,7 +1615,14 @@ public class TypeChecker {
                         ).value();
 
                         if (variantType instanceof StructType variantStructType) {
-                            return variantStructType;
+                            return SumVariant.struct(variantStructType);
+                        } else if (
+                            variantType instanceof ConstructedType<?> variantConstructedType &&
+                            variantConstructedType.genericType() instanceof StructType
+                        ) {
+                            return SumVariant.constructedStruct(
+                                (ConstructedType<StructType>) variantConstructedType
+                            );
                         } else {
                             throw new VariantTypeIsNotStructError(
                                 variantType,
@@ -1621,8 +1631,6 @@ public class TypeChecker {
                         }
                     })
                     .collect(PCollectors.toOrderedSet());
-
-                return context.addSumTypeInfo(typeOrConstructor.toType(), new SumTypeInfo(variants));
             }
 
             @Override
@@ -2072,7 +2080,7 @@ public class TypeChecker {
 
         for (var typedCase : typedCases) {
             var caseType = typedCase.caseType().value();
-            if (!variantTypes.contains(caseType)) {
+            if (!variantTypes.contains(SumVariant.struct(caseType))) {
                 throw new UnreachableWhenCaseError(caseType, typedCase.source());
             }
             if (handledVariantTypes.contains(caseType)) {
@@ -2082,7 +2090,7 @@ public class TypeChecker {
         }
 
         for (var variantType : variantTypes) {
-            if (!handledVariantTypes.contains(variantType)) {
+            if (!handledVariantTypes.contains(variantType.toType())) {
                 throw new UnhandledVariantError(variantType, untyped.source());
             }
         }
@@ -2189,4 +2197,22 @@ public class TypeChecker {
             )))
             .orElseGet(() -> TypeOrConstructor.type(type));
     }
+
+    private static EnvironmentBuilder environmentForTypeOrConstructor(
+        TypeOrConstructor<?> typeOrConstructor
+    ) {
+        var environment = EnvironmentBuilder.EMPTY;
+
+        if (typeOrConstructor instanceof TypeOrConstructor.Constructor<?> constructor) {
+            for (var param : constructor.value().params()) {
+                // TODO: proper source
+                environment = environment.add(
+                    param.name(),
+                    Variable.defined(Types.metaType(param), VariableType.CONSTANT),
+                    Sources.NULL
+                );
+            }
+        }
+        return environment;
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
index d1857276..2c69ac14 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
@@ -25,7 +25,7 @@ public class TypeCheckerSubTyping {
         if (supertype instanceof SumType supertypeSum) {
             var superTypeInfo = context.sumTypeInfo(supertypeSum, Sources.NULL);
             return superTypeInfo.variants().stream()
-                .anyMatch(supertypeVariant -> isSubType(subtype, supertypeVariant, context));
+                .anyMatch(supertypeVariant -> isSubType(subtype, supertypeVariant.toType(), context));
         }
 
         if (
diff --git a/src/main/java/org/shedlang/compiler/typechecker/errors/UnhandledVariantError.java b/src/main/java/org/shedlang/compiler/typechecker/errors/UnhandledVariantError.java
index de30b900..377573b6 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/errors/UnhandledVariantError.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/errors/UnhandledVariantError.java
@@ -2,17 +2,17 @@ package org.shedlang.compiler.typechecker.errors;
 
 import org.shedlang.compiler.errors.SourceError;
 import org.shedlang.compiler.sources.Source;
-import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.types.SumVariant;
 
 public class UnhandledVariantError extends SourceError {
-    private final Type unhandledVariant;
+    private final SumVariant unhandledVariant;
 
-    public UnhandledVariantError(Type unhandledVariant, Source source) {
+    public UnhandledVariantError(SumVariant unhandledVariant, Source source) {
         super("The variant " + unhandledVariant.describe() + " is unhandled", source);
         this.unhandledVariant = unhandledVariant;
     }
 
-    public Type unhandledVariant() {
+    public SumVariant unhandledVariant() {
         return unhandledVariant;
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/types/SumTypeInfo.java b/src/main/java/org/shedlang/compiler/types/SumTypeInfo.java
index a9027512..7dd359e3 100644
--- a/src/main/java/org/shedlang/compiler/types/SumTypeInfo.java
+++ b/src/main/java/org/shedlang/compiler/types/SumTypeInfo.java
@@ -2,5 +2,5 @@ package org.shedlang.compiler.types;
 
 import org.pcollections.OrderedPSet;
 
-public record SumTypeInfo(OrderedPSet<StructType> variants) {
+public record SumTypeInfo(OrderedPSet<SumVariant> variants) {
 }
diff --git a/src/main/java/org/shedlang/compiler/types/SumVariant.java b/src/main/java/org/shedlang/compiler/types/SumVariant.java
new file mode 100644
index 00000000..0270a93c
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/SumVariant.java
@@ -0,0 +1,39 @@
+package org.shedlang.compiler.types;
+
+public sealed interface SumVariant {
+    static SumVariant constructedStruct(ConstructedType<StructType> type) {
+        return new ConstructedStruct(type);
+    }
+
+    static SumVariant struct(StructType type) {
+        return new Struct(type);
+    }
+
+    String describe();
+
+    Type toType();
+
+    record ConstructedStruct(ConstructedType<StructType> type) implements SumVariant {
+        @Override
+        public String describe() {
+            return type.describe();
+        }
+
+        @Override
+        public Type toType() {
+            return type;
+        }
+    }
+
+    record Struct(StructType type) implements SumVariant {
+        @Override
+        public String describe() {
+            return type.describe();
+        }
+
+        @Override
+        public Type toType() {
+            return type;
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/Type.java b/src/main/java/org/shedlang/compiler/types/Type.java
index 9a062156..f6f51b1e 100644
--- a/src/main/java/org/shedlang/compiler/types/Type.java
+++ b/src/main/java/org/shedlang/compiler/types/Type.java
@@ -6,5 +6,5 @@ package org.shedlang.compiler.types;
  * found through auxiliary data.
  */
 public sealed interface Type extends TypeLevelValue
-    permits BoolType, BottomType, ConstructedType, FunctionType, IntType, NamespaceType, StringType, StructuredType, SumType, TopType, TupleType, TypeLevelValueType, TypeParam, UnitType {
+    permits BoolType, BottomType, ConstructedType, FunctionType, IntType, NamespaceType, StringType, StructuredType, TopType, TupleType, TypeLevelValueType, TypeParam, UnitType {
 }
diff --git a/src/main/java/org/shedlang/compiler/types/TypeParam.java b/src/main/java/org/shedlang/compiler/types/TypeParam.java
index 702d6a5f..4788276a 100644
--- a/src/main/java/org/shedlang/compiler/types/TypeParam.java
+++ b/src/main/java/org/shedlang/compiler/types/TypeParam.java
@@ -5,6 +5,10 @@ public record TypeParam(Scope scope, String name) implements Type {
         return new TypeParam(new Scope.Struct(structType), name);
     }
 
+    public static TypeParam forSumType(SumType sumType, String name) {
+        return new TypeParam(new Scope.Sum(sumType), name);
+    }
+
     @Override
     public String describe() {
         return name;
diff --git a/src/main/java/org/shedlang/compiler/types/Types.java b/src/main/java/org/shedlang/compiler/types/Types.java
index 0e502249..20c8e0cc 100644
--- a/src/main/java/org/shedlang/compiler/types/Types.java
+++ b/src/main/java/org/shedlang/compiler/types/Types.java
@@ -49,7 +49,7 @@ public class Types {
         return new TypeConstructor<>(params, genericType);
     }
 
-    public static <T extends StructuredType> Type constructedType(
+    public static <T extends StructuredType> ConstructedType<T> constructedType(
         TypeConstructor<T> constructor,
         List<Type> args
     ) {