Store variants bidirectionally

6 files changed, 63 insertions, 15 deletions

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 3e89209..7a4fb58 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
@@ -344,5 +344,9 @@ public class JavaTypesGenerator implements Generator {
         public List<SumType> variantOf(StructType type) {
             return this.typesInfo.variantOf(type);
         }
+
+        public List<Type> variants(SumType type) {
+            return this.typesInfo.variants(type);
+        }
     }
 }
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 e653321..9bc0efb 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerGlobalContext.java
@@ -1,6 +1,7 @@
 package org.zwobble.hobgoblin.compiler.typechecker;
 
 import org.zwobble.hobgoblin.compiler.types.*;
+import org.zwobble.hobgoblin.compiler.util.ManyToMany;
 
 import java.util.HashMap;
 import java.util.List;
@@ -9,7 +10,7 @@ import java.util.Map;
 public class TypeCheckerGlobalContext {
     private final Map<String, Type> nativeTypes = new HashMap<>();
     private final Map<StructType, List<Field>> fieldsOf = new HashMap<>();
-    private final Map<Type, List<SumType>> variantOf = new HashMap<>();
+    private final ManyToMany<SumType, Type> variants = new ManyToMany<>();
 
     public static TypeCheckerGlobalContext initial() {
         return new TypeCheckerGlobalContext();
@@ -19,7 +20,7 @@ public class TypeCheckerGlobalContext {
     }
 
     public TypesInfo toTypesInfo() {
-        return new TypesInfo(this.fieldsOf, this.variantOf);
+        return new TypesInfo(this.fieldsOf, this.variants);
     }
 
     public void addNativeType(SimpleNativeType type) {
@@ -35,7 +36,7 @@ public class TypeCheckerGlobalContext {
             namespaceName,
             this.nativeTypes,
             this.fieldsOf,
-            this.variantOf
+            this.variants
         );
     }
 }
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 5685a6e..6366a12 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,8 @@
 package org.zwobble.hobgoblin.compiler.typechecker;
 
 import org.zwobble.hobgoblin.compiler.types.*;
+import org.zwobble.hobgoblin.compiler.util.ManyToMany;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -13,7 +13,7 @@ public class TypeCheckerNamespaceContext {
         NamespaceName namespaceName,
         Map<String, Type> nativeTypes,
         Map<StructType, List<Field>> fieldsOf,
-        Map<Type, List<SumType>> variantOf
+        ManyToMany<SumType, Type> variants
     ) {
         var variables = nativeTypes.entrySet().stream()
             .collect(Collectors.toMap(
@@ -21,24 +21,24 @@ public class TypeCheckerNamespaceContext {
                 entry -> Variable.defined(entry.getValue())
             ));
 
-        return new TypeCheckerNamespaceContext(namespaceName, variables, fieldsOf, variantOf);
+        return new TypeCheckerNamespaceContext(namespaceName, variables, fieldsOf, variants);
     }
 
     private final NamespaceName namespaceName;
     private final Map<String, Variable> variables;
     private final Map<StructType, List<Field>> fieldsOf;
-    private final Map<Type, List<SumType>> variantOf;
+    private final ManyToMany<SumType, Type> variants;
 
     private TypeCheckerNamespaceContext(
         NamespaceName namespaceName,
         Map<String, Variable> variables,
         Map<StructType, List<Field>> fieldsOf,
-        Map<Type, List<SumType>> variantOf
+        ManyToMany<SumType, Type> variants
     ) {
         this.namespaceName = namespaceName;
         this.variables = variables;
         this.fieldsOf = fieldsOf;
-        this.variantOf = variantOf;
+        this.variants = variants;
     }
 
     public NamespaceName namespaceName() {
@@ -65,8 +65,7 @@ public class TypeCheckerNamespaceContext {
 
     public void defineSumType(SumType sumType, List<Type> variantTypes) {
         for (var variantType : variantTypes) {
-            this.variantOf.putIfAbsent(variantType, new ArrayList<>());
-            this.variantOf.get(variantType).add(sumType);
+            this.variants.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 b174be8..649ded0 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfo.java
@@ -4,6 +4,7 @@ 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.util.ManyToMany;
 
 import java.util.List;
 import java.util.Map;
@@ -11,11 +12,14 @@ import java.util.Optional;
 
 public class TypesInfo {
     private final Map<StructType, List<Field>> fieldsOf;
-    private final Map<Type, List<SumType>> variantOf;
+    private final ManyToMany<SumType, Type> variants;
 
-    public TypesInfo(Map<StructType, List<Field>> fieldsOf, Map<Type, List<SumType>> variantOf) {
+    public TypesInfo(
+        Map<StructType, List<Field>> fieldsOf,
+        ManyToMany<SumType, Type> variants
+    ) {
         this.fieldsOf = fieldsOf;
-        this.variantOf = variantOf;
+        this.variants = variants;
     }
 
     public List<Field> fieldsOf(StructType structType) {
@@ -24,6 +28,10 @@ public class TypesInfo {
     }
 
     public List<SumType> variantOf(StructType variantType) {
-        return this.variantOf.getOrDefault(variantType, List.of());
+        return this.variants.rightToLeft(variantType);
+    }
+
+    public List<Type> variants(SumType type) {
+        return this.variants.leftToRight(type);
     }
 }
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/util/ManyToMany.java b/src/main/java/org/zwobble/hobgoblin/compiler/util/ManyToMany.java
new file mode 100644
index 0000000..e5113d9
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/util/ManyToMany.java
@@ -0,0 +1,35 @@
+package org.zwobble.hobgoblin.compiler.util;
+
+import java.util.*;
+
+public class ManyToMany<TLeft, TRight> {
+    private final Map<TLeft, List<TRight>> leftToRight = new HashMap<>();
+    private final Map<TRight, List<TLeft>> rightToLeft = new HashMap<>();
+
+    public ManyToMany() {}
+
+    public List<TRight> leftToRight(TLeft left) {
+        var rights = this.leftToRight.get(left);
+        if (rights == null) {
+            return List.of();
+        } else {
+            return Collections.unmodifiableList(rights);
+        }
+    }
+
+    public List<TLeft> rightToLeft(TRight right) {
+        var lefts = this.rightToLeft.get(right);
+        if (lefts == null) {
+            return List.of();
+        } else {
+            return Collections.unmodifiableList(lefts);
+        }
+    }
+
+    public void add(TLeft left, TRight right) {
+        this.leftToRight.putIfAbsent(left, new ArrayList<>());
+        this.leftToRight.get(left).add(right);
+        this.rightToLeft.putIfAbsent(right, new ArrayList<>());
+        this.rightToLeft.get(right).add(left);
+    }
+}
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 fdf88ee..51de6a4 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
@@ -82,5 +82,6 @@ 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)));
     }
 }