Collect mapping of variants to sum types

4 files changed, 51 insertions, 5 deletions

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 012fd2b..6fd902b 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
@@ -111,14 +111,20 @@ public class TypeChecker {
         var sumType = (SumType) lookupMetaType(untyped.name(), untyped.source(), context);
 
         var typedVariants = new ArrayList<TypedSumVariantDefinitionNode>();
+        // TODO: should be StructType?
+        var variantTypes = new ArrayList<Type>();
         for (var untypedVariant : untyped.variants()) {
+            var variantType = typeCheckMetaType(untypedVariant.type(), context);
             var typedVariant = new TypedSumVariantDefinitionNode(
-                typeCheckMetaType(untypedVariant.type(), context),
+                variantType,
                 untyped.source()
             );
             typedVariants.add(typedVariant);
+            variantTypes.add(variantType.value());
         }
 
+        context.defineSumType(sumType, variantTypes);
+
         return new TypedSumDefinitionNode(
             sumType,
             typedVariants,
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 82a7e96..d3d4c43 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java
@@ -3,10 +3,12 @@ package org.zwobble.hobgoblin.compiler.typechecker;
 import org.zwobble.hobgoblin.compiler.types.*;
 
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 public class TypeCheckerGlobalContext {
     private final Map<String, Type> nativeTypes = new HashMap<>();
+    private final Map<Type, List<SumType>> variantOf = new HashMap<>();
 
     public static TypeCheckerGlobalContext initial() {
         return new TypeCheckerGlobalContext();
@@ -15,6 +17,10 @@ public class TypeCheckerGlobalContext {
     private TypeCheckerGlobalContext() {
     }
 
+    public TypesInfo toTypesInfo() {
+        return new TypesInfo(this.variantOf);
+    }
+
     public void addNativeType(SimpleNativeType type) {
         this.nativeTypes.put(type.name(), new TypeLevelValueType(type));
     }
@@ -24,6 +30,6 @@ public class TypeCheckerGlobalContext {
     }
 
     public TypeCheckerNamespaceContext enterNamespace(NamespaceName namespaceName) {
-        return TypeCheckerNamespaceContext.initial(namespaceName, this.nativeTypes);
+        return TypeCheckerNamespaceContext.initial(namespaceName, this.nativeTypes, this.variantOf);
     }
 }
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 b18267b..795bd6a 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerNamespaceContext.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerNamespaceContext.java
@@ -1,8 +1,11 @@
 package org.zwobble.hobgoblin.compiler.typechecker;
 
 import org.zwobble.hobgoblin.compiler.types.NamespaceName;
+import org.zwobble.hobgoblin.compiler.types.SumType;
 import org.zwobble.hobgoblin.compiler.types.Type;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -10,7 +13,8 @@ import java.util.stream.Collectors;
 public class TypeCheckerNamespaceContext {
     public static TypeCheckerNamespaceContext initial(
         NamespaceName namespaceName,
-        Map<String, Type> nativeTypes
+        Map<String, Type> nativeTypes,
+        Map<Type, List<SumType>> variantOf
     ) {
         var variables = nativeTypes.entrySet().stream()
             .collect(Collectors.toMap(
@@ -18,18 +22,21 @@ public class TypeCheckerNamespaceContext {
                 entry -> Variable.defined(entry.getValue())
             ));
 
-        return new TypeCheckerNamespaceContext(namespaceName, variables);
+        return new TypeCheckerNamespaceContext(namespaceName, variables, variantOf);
     }
 
     private final NamespaceName namespaceName;
     private final Map<String, Variable> variables;
+    private final Map<Type, List<SumType>> variantOf;
 
     private TypeCheckerNamespaceContext(
         NamespaceName namespaceName,
-        Map<String, Variable> variables
+        Map<String, Variable> variables,
+        Map<Type, List<SumType>> variantOf
     ) {
         this.namespaceName = namespaceName;
         this.variables = variables;
+        this.variantOf = variantOf;
     }
 
     public NamespaceName namespaceName() {
@@ -49,4 +56,11 @@ public class TypeCheckerNamespaceContext {
     public Optional<Variable> lookup(String name) {
         return Optional.ofNullable(this.variables.get(name));
     }
+
+    public void defineSumType(SumType sumType, ArrayList<Type> variantTypes) {
+        for (var variantType : variantTypes) {
+            this.variantOf.putIfAbsent(variantType, new ArrayList<>());
+            this.variantOf.get(variantType).add(sumType);
+        }
+    }
 }
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java
new file mode 100644
index 0000000..1c15553
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java
@@ -0,0 +1,20 @@
+package org.zwobble.hobgoblin.compiler.typechecker;
+
+import org.zwobble.hobgoblin.compiler.types.StructType;
+import org.zwobble.hobgoblin.compiler.types.SumType;
+import org.zwobble.hobgoblin.compiler.types.Type;
+
+import java.util.List;
+import java.util.Map;
+
+public class TypesInfo {
+    private final Map<Type, List<SumType>> variantOf;
+
+    public TypesInfo(Map<Type, List<SumType>> variantOf) {
+        this.variantOf = variantOf;
+    }
+
+    public List<SumType> variantOf(StructType variantType) {
+        return this.variantOf.getOrDefault(variantType, List.of());
+    }
+}