Push generation of type info for constructed typos into TypeInfos

2 files changed, 46 insertions, 35 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 239e8d4d..78bf4638 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -444,20 +444,8 @@ public class TypeCheckerContext {
     }
 
     public StructTypeInfo structTypeInfo(ConstructedType<StructType> type, Source source) {
-        var typeInfo = structTypeInfo(type.genericType(), source);
-
-        var typeSubstitutionMap = type.typeSubstitutionMap();
-
-        var fieldsInfoBuilder = FieldsInfo.builder();
-
-        for (var field : typeInfo.fields()) {
-            fieldsInfoBuilder = fieldsInfoBuilder.addField(
-                field.name(),
-                TypeSubstitution.substitute(field.type(), typeSubstitutionMap)
-            );
-        }
-
-        return new StructTypeInfo(fieldsInfoBuilder.build());
+        return this.typeInfos.structTypeInfo(type)
+            .orElseThrow(() -> new MissingStructTypeInfoError(type.genericType(), source));
     }
 
     public StructTypeInfo structTypeInfo(StructType type, Source source) {
@@ -481,27 +469,8 @@ public class TypeCheckerContext {
     }
 
     public SumTypeInfo sumTypeInfo(ConstructedType<SumType> type, Source source) {
-        var typeInfo = sumTypeInfo(type.genericType(), source);
-
-        var typeSubstitutionMap = type.typeSubstitutionMap();
-
-        var variants = typeInfo.variants().stream()
-            .map(variant -> switch (variant) {
-                case SumVariant.ConstructedStruct constructedStruct ->
-                    SumVariant.constructedStruct(TypeSubstitution.substituteConstructedType(
-                        constructedStruct.type(),
-                        typeSubstitutionMap
-                    ));
-
-                case SumVariant.Struct struct ->
-                    SumVariant.struct(TypeSubstitution.substituteStructType(
-                        struct.type(),
-                        typeSubstitutionMap
-                    ));
-            })
-            .collect(PCollectors.toOrderedSet());
-
-        return new SumTypeInfo(variants);
+        return this.typeInfos.sumTypeInfo(type)
+            .orElseThrow(() -> new MissingSumTypeInfoError(type.genericType(), source));
     }
 
     public TypeCheckerContext addUserDefinedEffectInfo(
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
index 8de1f0de..075daf51 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
@@ -5,6 +5,7 @@ import org.pcollections.OrderedPMap;
 import org.pcollections.PMap;
 import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
+import org.shedlang.compiler.util.PCollectors;
 
 import java.util.Optional;
 import java.util.Set;
@@ -61,6 +62,23 @@ public class TypeInfos {
         return Optional.ofNullable(structTypeInfos.get(type));
     }
 
+    public Optional<StructTypeInfo> structTypeInfo(ConstructedType<StructType> type) {
+        return structTypeInfo(type.genericType()).map(typeInfo -> {
+            var typeSubstitutionMap = type.typeSubstitutionMap();
+
+            var fieldsInfoBuilder = FieldsInfo.builder();
+
+            for (var field : typeInfo.fields()) {
+                fieldsInfoBuilder = fieldsInfoBuilder.addField(
+                    field.name(),
+                    TypeSubstitution.substitute(field.type(), typeSubstitutionMap)
+                );
+            }
+
+            return new StructTypeInfo(fieldsInfoBuilder.build());
+        });
+    }
+
     public Set<StructType> structTypes() {
         return structTypeInfos.keySet();
     }
@@ -78,6 +96,30 @@ public class TypeInfos {
         return Optional.ofNullable(sumTypeInfos.get(type));
     }
 
+    public Optional<SumTypeInfo> sumTypeInfo(ConstructedType<SumType> type) {
+        return sumTypeInfo(type.genericType()).map(typeInfo -> {
+            var typeSubstitutionMap = type.typeSubstitutionMap();
+
+            var variants = typeInfo.variants().stream()
+                .map(variant -> switch (variant) {
+                    case SumVariant.ConstructedStruct constructedStruct ->
+                        SumVariant.constructedStruct(TypeSubstitution.substituteConstructedType(
+                            constructedStruct.type(),
+                            typeSubstitutionMap
+                        ));
+
+                    case SumVariant.Struct struct ->
+                        SumVariant.struct(TypeSubstitution.substituteStructType(
+                            struct.type(),
+                            typeSubstitutionMap
+                        ));
+                })
+                .collect(PCollectors.toOrderedSet());
+
+            return new SumTypeInfo(variants);
+        });
+    }
+
     public TypeInfos addUserDefinedEffectInfo(UserDefinedEffect effect, UserDefinedEffectInfo effectInfo) {
         return new TypeInfos(
             this.namespaceTypeInfos,