Store enum variants in TypesInfo

9 files changed, 55 insertions, 20 deletions

diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javapreciselymatchers/JavaPreciselyMatchersGenerator.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javapreciselymatchers/JavaPreciselyMatchersGenerator.java
index bc29ab9..8cd8b40 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javapreciselymatchers/JavaPreciselyMatchersGenerator.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javapreciselymatchers/JavaPreciselyMatchersGenerator.java
@@ -275,7 +275,7 @@ public class JavaPreciselyMatchersGenerator implements Generator {
         }
 
         public List<Type> variants(SumType type) {
-            return this.typesInfo.variants(type);
+            return this.typesInfo.sumVariants(type);
         }
     }
 
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java
index bb073b6..bd56880 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java
@@ -452,7 +452,7 @@ public class JavaTypesGenerator implements Generator {
         }
 
         public List<Type> variants(SumType type) {
-            return this.typesInfo.variants(type);
+            return this.typesInfo.sumVariants(type);
         }
     }
 }
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
index ed4ad54..5003e3a 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
@@ -125,6 +125,12 @@ public class TypeChecker {
             ))
             .toList();
 
+        var enumVariants = untyped.variants().stream()
+            .map(untypedVariantNode -> new EnumVariant(untypedVariantNode.name()))
+            .toList();
+
+        context.defineEnumType(enumType, enumVariants);
+
         return new TypedEnumDefinitionNode(
             enumType,
             typedVariantNodes,
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java
index ff6cd5e..3e7665c 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java
@@ -7,9 +7,10 @@ import java.util.*;
 
 public class TypeCheckerGlobalContext {
     private final Map<String, Type> nativeTypes = new HashMap<>();
+    private final Map<EnumType, List<EnumVariant>> enumVariants = new HashMap<>();
     private final Map<StructType, List<Field>> structFieldsOf = new HashMap<>();
     private final Map<NamespaceName, Fields> namespaceFieldsOf = new HashMap<>();
-    private final ManyToMany<SumType, Type> variants = new ManyToMany<>();
+    private final ManyToMany<SumType, Type> sumVariants = new ManyToMany<>();
 
     public static TypeCheckerGlobalContext initial() {
         return new TypeCheckerGlobalContext();
@@ -19,7 +20,11 @@ public class TypeCheckerGlobalContext {
     }
 
     public TypesInfo toTypesInfo() {
-        return new TypesInfo(this.structFieldsOf, this.variants);
+        return new TypesInfo(
+            this.enumVariants,
+            this.structFieldsOf,
+            this.sumVariants
+        );
     }
 
     public void addNativeType(SimpleNativeType type) {
@@ -34,8 +39,9 @@ public class TypeCheckerGlobalContext {
         return TypeCheckerNamespaceContext.initial(
             namespaceName,
             this.nativeTypes,
+            this.enumVariants,
             this.structFieldsOf,
-            this.variants
+            this.sumVariants
         );
     }
 
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerNamespaceContext.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerNamespaceContext.java
index 69827a8..906cdfa 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerNamespaceContext.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerNamespaceContext.java
@@ -12,12 +12,14 @@ public class TypeCheckerNamespaceContext {
     public static TypeCheckerNamespaceContext initial(
         NamespaceName namespaceName,
         Map<String, Type> nativeTypes,
+        Map<EnumType, List<EnumVariant>> enumVariants,
         Map<StructType, List<Field>> fieldsOf,
         ManyToMany<SumType, Type> variants
     ) {
         return new TypeCheckerNamespaceContext(
             namespaceName,
             new HashMap<>(nativeTypes),
+            enumVariants,
             fieldsOf,
             variants
         );
@@ -25,19 +27,22 @@ public class TypeCheckerNamespaceContext {
 
     private final NamespaceName namespaceName;
     private final Map<String, Type> variables;
+    private final Map<EnumType, List<EnumVariant>> enumVariants;
     private final Map<StructType, List<Field>> fieldsOf;
-    private final ManyToMany<SumType, Type> variants;
+    private final ManyToMany<SumType, Type> sumVariants;
 
     private TypeCheckerNamespaceContext(
         NamespaceName namespaceName,
         Map<String, Type> variables,
+        Map<EnumType, List<EnumVariant>> enumVariants,
         Map<StructType, List<Field>> fieldsOf,
-        ManyToMany<SumType, Type> variants
+        ManyToMany<SumType, Type> sumVariants
     ) {
         this.namespaceName = namespaceName;
         this.variables = variables;
+        this.enumVariants = enumVariants;
         this.fieldsOf = fieldsOf;
-        this.variants = variants;
+        this.sumVariants = sumVariants;
     }
 
     public NamespaceName namespaceName() {
@@ -53,13 +58,17 @@ public class TypeCheckerNamespaceContext {
         return Optional.ofNullable(this.variables.get(name));
     }
 
+    public void defineEnumType(EnumType enumType, List<EnumVariant> enumVariants) {
+        this.enumVariants.put(enumType, enumVariants);
+    }
+
     public void defineStructType(StructType structType, List<Field> fields) {
         this.fieldsOf.put(structType, fields);
     }
 
     public void defineSumType(SumType sumType, List<Type> variantTypes) {
         for (var variantType : variantTypes) {
-            this.variants.add(sumType, variantType);
+            this.sumVariants.add(sumType, variantType);
         }
     }
 }
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java
index e64a4b8..0fb4f81 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java
@@ -1,9 +1,6 @@
 package org.zwobble.hobgoblin.compiler.typechecker;
 
-import org.zwobble.hobgoblin.compiler.types.Field;
-import org.zwobble.hobgoblin.compiler.types.StructType;
-import org.zwobble.hobgoblin.compiler.types.SumType;
-import org.zwobble.hobgoblin.compiler.types.Type;
+import org.zwobble.hobgoblin.compiler.types.*;
 import org.zwobble.hobgoblin.compiler.util.ManyToMany;
 
 import java.util.List;
@@ -11,15 +8,22 @@ import java.util.Map;
 import java.util.Optional;
 
 public class TypesInfo {
+    private final Map<EnumType, List<EnumVariant>> enumVariants;
     private final Map<StructType, List<Field>> structFieldsOf;
-    private final ManyToMany<SumType, Type> variants;
+    private final ManyToMany<SumType, Type> sumVariants;
 
     public TypesInfo(
+        Map<EnumType, List<EnumVariant>> enumVariants,
         Map<StructType, List<Field>> structFieldsOf,
-        ManyToMany<SumType, Type> variants
+        ManyToMany<SumType, Type> sumVariants
     ) {
+        this.enumVariants = enumVariants;
         this.structFieldsOf = structFieldsOf;
-        this.variants = variants;
+        this.sumVariants = sumVariants;
+    }
+
+    public List<EnumVariant> enumVariants(EnumType enumType) {
+        return this.enumVariants.get(enumType);
     }
 
     public List<Field> fieldsOf(StructType structType) {
@@ -28,10 +32,10 @@ public class TypesInfo {
     }
 
     public List<SumType> variantOf(StructType variantType) {
-        return this.variants.rightToLeft(variantType);
+        return this.sumVariants.rightToLeft(variantType);
     }
 
-    public List<Type> variants(SumType type) {
-        return this.variants.leftToRight(type);
+    public List<Type> sumVariants(SumType type) {
+        return this.sumVariants.leftToRight(type);
     }
 }
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/types/EnumVariant.java b/src/main/java/org/zwobble/hobgoblin/compiler/types/EnumVariant.java
new file mode 100644
index 0000000..2297a64
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/types/EnumVariant.java
@@ -0,0 +1,4 @@
+package org.zwobble.hobgoblin.compiler.types;
+
+public record EnumVariant(String name) {
+}
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerEnumDefinitionTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerEnumDefinitionTests.java
index f8a8fb4..ac85b7d 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerEnumDefinitionTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerEnumDefinitionTests.java
@@ -57,5 +57,11 @@ public class TypeCheckerEnumDefinitionTests {
                 )
             )
         ));
+        var typesInfo = globalContext.toTypesInfo();
+        var enumType = new EnumType(NamespaceName.of("a", "b"), "X");
+        assertThat(typesInfo.enumVariants(enumType), isSequence(
+            equalTo(new EnumVariant("a")),
+            equalTo(new EnumVariant("b"))
+        ));
     }
 }
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
index 60c2b9a..57c40f5 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
@@ -76,7 +76,7 @@ public class TypeCheckerSumDefinitionTests {
         var typesInfo = globalContext.toTypesInfo();
         assertThat(typesInfo.variantOf(rectangleType), isSequence(equalTo(sumType)));
         assertThat(typesInfo.variantOf(triangleType), isSequence(equalTo(sumType)));
-        assertThat(typesInfo.variants(sumType), isSequence(equalTo(rectangleType), equalTo(triangleType)));
+        assertThat(typesInfo.sumVariants(sumType), isSequence(equalTo(rectangleType), equalTo(triangleType)));
     }
 
     @Test