Add type info for sum types

6 files changed, 82 insertions, 10 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 e6900457..150c2c59 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedSumTypeDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedSumTypeDefinitionNode.java
@@ -22,8 +22,8 @@ public record UntypedSumTypeDefinitionNode(
     ) {
         public UntypedSumTypeDefinitionNode build() {
             return new UntypedSumTypeDefinitionNode(
-                name,
-                List.of(),
+                this.name,
+                this.variants,
                 Sources.NULL
             );
         }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index e184e26b..c9413be8 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -925,6 +925,21 @@ public class TypeChecker {
             }
 
             @Override
+            public TypeCheckerContext defineTypes(TypeCheckerContext context) {
+                var type = context.getLexicalState(SumType.class, untyped.name());
+
+                var variants = untyped.variants().stream()
+                    .map(untypedVariant -> typeCheckTypeLevelExpression(
+                        TypeSets.META_TYPE,
+                        untypedVariant,
+                        context
+                    ).value())
+                    .collect(PCollectors.toOrderedSet());
+
+                return context.addSumTypeInfo(type, new SumTypeInfo(variants));
+            }
+
+            @Override
             public TypeCheckerResult<TypedNamespaceStatementNode> finish(TypeCheckerContext context) {
                 var type = context.getLexicalState(SumType.class, untyped.name());
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 0f81a7fb..348380cb 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -352,6 +352,16 @@ public class TypeCheckerContext {
             .orElseThrow(() -> new MissingStructTypeInfoError(type, source));
     }
 
+    public TypeCheckerContext addSumTypeInfo(
+        SumType sumType,
+        SumTypeInfo sumTypeInfo
+    ) {
+        return new TypeCheckerContext(
+            this.typeInfos.addSumTypeInfo(sumType, sumTypeInfo),
+            this.stack
+        );
+    }
+
     public <T> TypeCheckerResult<T> toResult(T value) {
         return new TypeCheckerResult<>(value, this);
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
index 92a96fdd..5e6d6ee0 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
@@ -2,25 +2,25 @@ package org.shedlang.compiler.typechecker;
 
 import org.pcollections.Empty;
 import org.pcollections.PMap;
-import org.shedlang.compiler.types.NamespaceName;
-import org.shedlang.compiler.types.NamespaceTypeInfo;
-import org.shedlang.compiler.types.StructType;
-import org.shedlang.compiler.types.StructTypeInfo;
+import org.shedlang.compiler.types.*;
 
 import java.util.Optional;
 
 public class TypeInfos {
-    public static final TypeInfos EMPTY = new TypeInfos(Empty.map(), Empty.map());
+    public static final TypeInfos EMPTY = new TypeInfos(Empty.map(), Empty.map(), Empty.map());
 
     private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfos;
     private final PMap<StructType, StructTypeInfo> structTypeInfos;
+    private final PMap<SumType, SumTypeInfo> sumTypeInfos;
 
     public TypeInfos(
         PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfos,
-        PMap<StructType, StructTypeInfo> structTypeInfos
+        PMap<StructType, StructTypeInfo> structTypeInfos,
+        PMap<SumType, SumTypeInfo> sumTypeInfos
     ) {
         this.namespaceTypeInfos = namespaceTypeInfos;
         this.structTypeInfos = structTypeInfos;
+        this.sumTypeInfos = sumTypeInfos;
     }
 
     public TypeInfos addNamespaceTypeInfo(
@@ -29,7 +29,8 @@ public class TypeInfos {
     ) {
         return new TypeInfos(
             this.namespaceTypeInfos.plus(namespaceName, namespaceTypeInfo),
-            this.structTypeInfos
+            this.structTypeInfos,
+            this.sumTypeInfos
         );
     }
 
@@ -43,11 +44,24 @@ public class TypeInfos {
     ) {
         return new TypeInfos(
             this.namespaceTypeInfos,
-            this.structTypeInfos.plus(type, structTypeInfo)
+            this.structTypeInfos.plus(type, structTypeInfo),
+            this.sumTypeInfos
         );
     }
 
     public Optional<StructTypeInfo> structTypeInfo(StructType type) {
         return Optional.ofNullable(structTypeInfos.get(type));
     }
+
+    public TypeInfos addSumTypeInfo(SumType sumType, SumTypeInfo sumTypeInfo) {
+        return new TypeInfos(
+            this.namespaceTypeInfos,
+            this.structTypeInfos,
+            this.sumTypeInfos.plus(sumType, sumTypeInfo)
+        );
+    }
+
+    public Optional<SumTypeInfo> sumTypeInfo(SumType type) {
+        return Optional.ofNullable(sumTypeInfos.get(type));
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/types/SumTypeInfo.java b/src/main/java/org/shedlang/compiler/types/SumTypeInfo.java
new file mode 100644
index 00000000..916842cf
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/SumTypeInfo.java
@@ -0,0 +1,6 @@
+package org.shedlang.compiler.types;
+
+import org.pcollections.OrderedPSet;
+
+public record SumTypeInfo(OrderedPSet<Type> variants) {
+}
diff --git a/src/main/java/org/shedlang/compiler/util/PCollectors.java b/src/main/java/org/shedlang/compiler/util/PCollectors.java
index a0019eef..28e1fb28 100644
--- a/src/main/java/org/shedlang/compiler/util/PCollectors.java
+++ b/src/main/java/org/shedlang/compiler/util/PCollectors.java
@@ -1,6 +1,7 @@
 package org.shedlang.compiler.util;
 
 import org.pcollections.Empty;
+import org.pcollections.OrderedPSet;
 import org.pcollections.PMap;
 import org.pcollections.PVector;
 
@@ -63,6 +64,32 @@ public class PCollectors {
         }
     }
 
+    public static <T> Collector<T, OrderedPSetCollector<T>, OrderedPSet<T>> toOrderedSet() {
+        return Collector.of(
+            OrderedPSetCollector::new,
+            OrderedPSetCollector::add,
+            OrderedPSetCollector::addAll,
+            OrderedPSetCollector::toSet
+        );
+    }
+
+    private static class OrderedPSetCollector<T> {
+        private OrderedPSet<T> set = OrderedPSet.empty();
+
+        void add(T value) {
+            set = set.plus(value);
+        }
+
+        OrderedPSetCollector<T> addAll(OrderedPSetCollector<T> other) {
+            set = set.plusAll(other.set);
+            return this;
+        }
+
+        OrderedPSet<T> toSet() {
+            return set;
+        }
+    }
+
     public static <T> Collector<T, PVectorCollector<T>, PVector<T>> toVector() {
         return Collector.of(
             PVectorCollector::new,