Use TypeInfos as builder in TypeCheckerContext

3 files changed, 25 insertions, 50 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index a0822d7e..5871c510 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -35,11 +35,15 @@ public record PatioCodeGenContext(
         OrderedPMap<String, PatioCompilationUnitElement> elementsByName
     ) {
         public NamespaceTypeInfo namespaceTypeInfo(NamespaceName namespaceName) {
-            return typeInfos.namespaceTypeInfo(namespaceName);
+            return typeInfos.namespaceTypeInfo(namespaceName)
+                // TODO: throw better error
+                .orElseThrow(() -> new RuntimeException("No type info for: " + namespaceName));
         }
 
         public StructTypeInfo structTypeInfo(StructType type) {
-            return typeInfos.structTypeInfo(type);
+            return typeInfos.structTypeInfo(type)
+                // TODO: throw better error
+                .orElseThrow(() -> new RuntimeException("No type info for: " + type));
         }
 
         public static final CompilationUnitContext EMPTY = forTypeInfos(TypeInfos.EMPTY);
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 39968737..0f81a7fb 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -20,8 +20,7 @@ import java.util.function.Function;
 
 public class TypeCheckerContext {
     public static final TypeCheckerContext EMPTY = new TypeCheckerContext(
-        Empty.map(),
-        Empty.map(),
+        TypeInfos.EMPTY,
         P.stack(new LexicalContext(
             Optional.empty(),
             Empty.map(),
@@ -104,17 +103,14 @@ public class TypeCheckerContext {
         }
     }
 
-    private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfos;
-    private final PMap<StructType, StructTypeInfo> structTypeInfos;
+    private final TypeInfos typeInfos;
     private final PStack<StackElement> stack;
 
     public TypeCheckerContext(
-        PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfos,
-        PMap<StructType, StructTypeInfo> structTypeInfos,
+        TypeInfos typeInfos,
         PStack<StackElement> stack
     ) {
-        this.namespaceTypeInfos = namespaceTypeInfos;
-        this.structTypeInfos = structTypeInfos;
+        this.typeInfos = typeInfos;
         this.stack = stack;
     }
 
@@ -142,8 +138,7 @@ public class TypeCheckerContext {
         }
 
         return new TypeCheckerContext(
-            this.namespaceTypeInfos,
-            this.structTypeInfos,
+            this.typeInfos,
             newStack
         );
     }
@@ -179,8 +174,7 @@ public class TypeCheckerContext {
         Optional<Type> newReturnType
     ) {
         return new TypeCheckerContext(
-            this.namespaceTypeInfos,
-            this.structTypeInfos,
+            this.typeInfos,
             this.stack.plus(new LexicalContext(
                 newNamespaceName,
                 this.environment().plusAll(newEnvironment),
@@ -195,8 +189,7 @@ public class TypeCheckerContext {
 
     public TypeCheckerContext exit() {
         return new TypeCheckerContext(
-            this.namespaceTypeInfos,
-            this.structTypeInfos,
+            this.typeInfos,
             this.stack.minus(0)
         );
     }
@@ -204,8 +197,7 @@ public class TypeCheckerContext {
     public TypeCheckerContext enterFreeVariableCapture() {
         // TODO: do this properly (should not affect environment, possibly other bits of lexical context)
         return new TypeCheckerContext(
-            this.namespaceTypeInfos,
-            this.structTypeInfos,
+            this.typeInfos,
             this.stack.plus(FreeVariableCapture.EMPTY)
         );
     }
@@ -242,8 +234,7 @@ public class TypeCheckerContext {
         }
 
         return new TypeCheckerContext(
-            this.namespaceTypeInfos,
-            this.structTypeInfos,
+            this.typeInfos,
             newLexicalContextStack
         );
     }
@@ -341,36 +332,26 @@ public class TypeCheckerContext {
         var namespaceTypeInfo = new NamespaceTypeInfo(new FieldsInfo(fields));
 
         return new TypeCheckerContext(
-            this.namespaceTypeInfos.plus(namespaceName, namespaceTypeInfo),
-            this.structTypeInfos,
+            this.typeInfos.addNamespaceTypeInfo(namespaceName, namespaceTypeInfo),
             this.stack
         );
     }
 
-    public PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo() {
-        return namespaceTypeInfos;
-    }
-
     public TypeCheckerContext addStructTypeInfo(
         StructType structType,
         StructTypeInfo structTypeInfo
     ) {
         return new TypeCheckerContext(
-            this.namespaceTypeInfos,
-            this.structTypeInfos.plus(structType, structTypeInfo),
+            this.typeInfos.addStructTypeInfo(structType, structTypeInfo),
             this.stack
         );
     }
 
     public StructTypeInfo structTypeInfo(StructType type, Source source) {
-        return Optional.ofNullable(this.structTypeInfos.get(type))
+        return this.typeInfos.structTypeInfo(type)
             .orElseThrow(() -> new MissingStructTypeInfoError(type, source));
     }
 
-    public PMap<StructType, StructTypeInfo> structTypeInfos() {
-        return structTypeInfos;
-    }
-
     public <T> TypeCheckerResult<T> toResult(T value) {
         return new TypeCheckerResult<>(value, this);
     }
@@ -392,6 +373,6 @@ public class TypeCheckerContext {
     }
 
     public TypeInfos toTypeInfos() {
-        return new TypeInfos(namespaceTypeInfos, structTypeInfos);
+        return typeInfos;
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
index 6aa6ac59..92a96fdd 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
@@ -7,6 +7,8 @@ import org.shedlang.compiler.types.NamespaceTypeInfo;
 import org.shedlang.compiler.types.StructType;
 import org.shedlang.compiler.types.StructTypeInfo;
 
+import java.util.Optional;
+
 public class TypeInfos {
     public static final TypeInfos EMPTY = new TypeInfos(Empty.map(), Empty.map());
 
@@ -31,14 +33,8 @@ public class TypeInfos {
         );
     }
 
-    public NamespaceTypeInfo namespaceTypeInfo(NamespaceName namespaceName) {
-        // TODO: throw better error
-        var typeInfo = namespaceTypeInfos.get(namespaceName);
-        if (typeInfo == null) {
-            throw new RuntimeException("No type info for: " + namespaceName);
-        } else {
-            return typeInfo;
-        }
+    public Optional<NamespaceTypeInfo> namespaceTypeInfo(NamespaceName namespaceName) {
+        return Optional.ofNullable(namespaceTypeInfos.get(namespaceName));
     }
 
     public TypeInfos addStructTypeInfo(
@@ -51,13 +47,7 @@ public class TypeInfos {
         );
     }
 
-    public StructTypeInfo structTypeInfo(StructType type) {
-        // TODO: throw better error
-        var typeInfo = structTypeInfos.get(type);
-        if (typeInfo == null) {
-            throw new RuntimeException("No type info for: " + type);
-        } else {
-            return typeInfo;
-        }
+    public Optional<StructTypeInfo> structTypeInfo(StructType type) {
+        return Optional.ofNullable(structTypeInfos.get(type));
     }
 }