Ensure variants are structs

3 files changed, 35 insertions, 6 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 14992468..4cb52e54 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -930,11 +930,22 @@ public class TypeChecker {
 
                 // TODO: restrict variants to structs. Do we allow other sum types?
                 var variants = untyped.variants().stream()
-                    .map(untypedVariant -> typeCheckTypeLevelExpression(
-                        TypeSets.META_TYPE,
-                        untypedVariant,
-                        context
-                    ).value())
+                    .map(untypedVariant -> {
+                        var variantType = typeCheckTypeLevelExpression(
+                            TypeSets.META_TYPE,
+                            untypedVariant,
+                            context
+                        ).value();
+
+                        if (variantType instanceof StructType variantStructType) {
+                            return variantStructType;
+                        } else {
+                            throw new VariantTypeIsNotStructError(
+                                variantType,
+                                untypedVariant.source()
+                            );
+                        }
+                    })
                     .collect(PCollectors.toOrderedSet());
 
                 return context.addSumTypeInfo(type, new SumTypeInfo(variants));
diff --git a/src/main/java/org/shedlang/compiler/typechecker/errors/VariantTypeIsNotStructError.java b/src/main/java/org/shedlang/compiler/typechecker/errors/VariantTypeIsNotStructError.java
new file mode 100644
index 00000000..dd7f4bed
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/errors/VariantTypeIsNotStructError.java
@@ -0,0 +1,18 @@
+package org.shedlang.compiler.typechecker.errors;
+
+import org.shedlang.compiler.errors.SourceError;
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.Type;
+
+public class VariantTypeIsNotStructError extends SourceError {
+    private final Type variantType;
+
+    public VariantTypeIsNotStructError(Type variantType, Source source) {
+        super("Variant types must be structs, but variant is " + variantType.describe(), source);
+        this.variantType = variantType;
+    }
+
+    public Type variantType() {
+        return variantType;
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/SumTypeInfo.java b/src/main/java/org/shedlang/compiler/types/SumTypeInfo.java
index 916842cf..a9027512 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<Type> variants) {
+public record SumTypeInfo(OrderedPSet<StructType> variants) {
 }