1 files changed, 23 insertions, 22 deletions
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfoInMemory.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfoInMemory.java
index 7da4d2c..64c285b 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfoInMemory.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypesInfoInMemory.java
@@ -5,12 +5,16 @@ import org.zwobble.hobgoblin.compiler.types.*;
 import java.util.*;
 
 public class TypesInfoInMemory implements TypesInfo {
+    private record SumTypeInfo(
+        List<SumVariant> variants,
+        List<Field> fields
+    ) {}
+
     public static TypesInfoInMemory empty() {
         return new TypesInfoInMemory(
             new HashMap<>(),
             new HashMap<>(),
             new HashMap<>(),
-            new HashMap<>(),
             new HashMap<>()
         );
     }
@@ -19,21 +23,18 @@ public class TypesInfoInMemory implements TypesInfo {
 
     private final Map<Type, Optional<List<Field>>> structFields;
 
-    private final Map<Type, List<Field>> sumFields;
-    private final Map<SumType, List<SumVariant>> sumToVariants;
+    private final Map<SumType, SumTypeInfo> sumTypeInfos;
     private final Map<StructType, List<SumType>> variantToSums;
 
-    public TypesInfoInMemory(
+    private TypesInfoInMemory(
         Map<EnumType, List<EnumVariant>> enumVariants,
         Map<Type, Optional<List<Field>>> structFields,
-        Map<Type, List<Field>> sumFields,
-        Map<SumType, List<SumVariant>> sumToVariants,
+        Map<SumType, SumTypeInfo> sumTypeInfos,
         Map<StructType, List<SumType>> variantToSums
     ) {
         this.enumVariants = enumVariants;
         this.structFields = structFields;
-        this.sumFields = sumFields;
-        this.sumToVariants = sumToVariants;
+        this.sumTypeInfos = sumTypeInfos;
         this.variantToSums = variantToSums;
     }
 
@@ -70,12 +71,13 @@ public class TypesInfoInMemory implements TypesInfo {
         List<SumVariant> variants,
         List<Field> fields
     ) {
-        this.sumToVariants.put(sumType, variants);
+        var sumTypeInfo = new SumTypeInfo(variants, fields);
+        this.sumTypeInfos.put(sumType, sumTypeInfo);
+
         for (var variant : variants) {
             this.variantToSums.putIfAbsent(variant.valueType(), new ArrayList<>());
             this.variantToSums.get(variant.valueType()).add(sumType);
         }
-        this.sumFields.put(sumType, fields);
     }
 
     public List<SumType> variantOf(StructType variantType) {
@@ -83,21 +85,20 @@ public class TypesInfoInMemory implements TypesInfo {
     }
 
     public List<SumVariant> sumVariants(SumType type) {
-        var variants = this.sumToVariants.get(type);
-        if (variants == null) {
-            // TODO: better error
-            throw new RuntimeException("sum not defined");
-        } else {
-            return variants;
-        }
+        return sumTypeInfo(type).variants;
     }
 
-    public List<Field> fieldsOf(SumType sumType) {
-        var fields = this.sumFields.get(sumType);
-        if (fields == null) {
-            return List.of();
+    public List<Field> fieldsOf(SumType type) {
+        return sumTypeInfo(type).fields;
+    }
+
+    private SumTypeInfo sumTypeInfo(SumType type) {
+        var typeInfo = this.sumTypeInfos.get(type);
+        if (typeInfo == null) {
+            // TODO: better error
+            throw new RuntimeException("sum not defined");
         } else {
-            return fields;
+            return typeInfo;
         }
     }
 }