Introduce FieldInfo

5 files changed, 57 insertions, 40 deletions

diff --git a/src/main/java/org/shedlang/compiler/ShedCompiler.java b/src/main/java/org/shedlang/compiler/ShedCompiler.java
index 6e7f187a..1c2a471f 100644
--- a/src/main/java/org/shedlang/compiler/ShedCompiler.java
+++ b/src/main/java/org/shedlang/compiler/ShedCompiler.java
@@ -10,10 +10,7 @@ import org.shedlang.compiler.stackir.StackCompiler;
 import org.shedlang.compiler.stackir.StackIrNamespace;
 import org.shedlang.compiler.typechecker.TypeChecker;
 import org.shedlang.compiler.typechecker.TypeCheckerContext;
-import org.shedlang.compiler.typechecker.TypeInfos;
-import org.shedlang.compiler.types.Effects;
-import org.shedlang.compiler.types.NamespaceName;
-import org.shedlang.compiler.types.Types;
+import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
 
 import java.io.IOException;
@@ -42,9 +39,12 @@ public class ShedCompiler {
         var namespaces = new ArrayList<StackIrNamespace>();
 
         var typeCheckerContext = TypeCheckerContext.EMPTY
-            .addNamespaceFields(ShedBuiltins.NAMESPACE_NAME, P.orderedMap(
-                Map.entry(ShedBuiltins.PRINT_NAME, ShedBuiltins.PRINT_TYPE)
-            ))
+            .addNamespaceTypeInfo(
+                ShedBuiltins.NAMESPACE_NAME,
+                new NamespaceTypeInfo(FieldsInfo.of(
+                    new FieldInfo<>(ShedBuiltins.PRINT_NAME, ShedBuiltins.PRINT_TYPE)
+                ))
+            )
             .enterGlobal(P.map(
                 Map.entry("Bool", Types.metaType(Types.BOOL)),
                 Map.entry("Int", Types.metaType(Types.INT)),
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 222ddc79..a1139a98 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -332,12 +332,7 @@ public class TypeChecker {
             var structTypeInfo = context.structTypeInfo(structType, typedReceiver.source());
 
             var namedParams = structTypeInfo.fields().stream()
-                .map(field -> {
-                    var name = field.getKey();
-                    var type = field.getValue();
-
-                    return new NamedParam(name, type);
-                })
+                .map(field -> new NamedParam(field.name(), field.type()))
                 .toList();
 
             return new FunctionType(
@@ -415,12 +410,10 @@ public class TypeChecker {
                 var fields = untyped.fields().stream()
                     .map(fieldNode -> {
                         var typedFunctionTypeNode = typeCheckFunctionType(fieldNode.type(), context);
+                        return new FieldInfo<>(fieldNode.name(), typedFunctionTypeNode.value());
+                    });
 
-                        return Map.entry(fieldNode.name(), typedFunctionTypeNode.value());
-                    })
-                    .collect(PCollectors.toOrderedMap());
-
-                var effectInfo = new UserDefinedEffectInfo(new FieldsInfo<>(fields));
+                var effectInfo = new UserDefinedEffectInfo(FieldsInfo.from(fields));
 
                 return context.addUserDefinedEffectInfo(effect, effectInfo);
             }
@@ -1176,7 +1169,7 @@ public class TypeChecker {
         declaredVariables = DeclaredVariables.forNamespaceStatements(untyped.body(), declaredVariables);
 
         var namespaceFinalContext = context;
-        var fieldTypes = declaredVariables.stream()
+        var fieldInfos = declaredVariables.stream()
             .filter(variable -> switch (variable.variableType()) {
                 case CONSTANT -> true;
                 case IMPORT -> false;
@@ -1188,16 +1181,16 @@ public class TypeChecker {
                 case PUBLIC -> true;
                 case PRIVATE -> false;
             })
-            .map(variable -> Map.entry(
+            .map(variable -> new FieldInfo<>(
                 variable.name(),
                 namespaceFinalContext.compileTimeLookup(variable.name())
                     .orElseThrow() // TODO: better error
                     .type()
                     .orElseThrow() // TODO: better error
-            ))
-            .collect(PCollectors.toOrderedMap());
+            ));
 
-        var fieldNames = fieldTypes.keys();
+        var fieldsInfo = FieldsInfo.from(fieldInfos);
+        var fieldNames = fieldsInfo.keys();
 
         var typed = new TypedNamespaceNode(
             untyped.namespaceName(),
@@ -1212,7 +1205,7 @@ public class TypeChecker {
 
         return context
             .exit()
-            .addNamespaceFields(typed.namespaceName(), fieldTypes)
+            .addNamespaceTypeInfo(typed.namespaceName(), new NamespaceTypeInfo(fieldsInfo))
             .toResult(typed);
     }
 
@@ -1305,7 +1298,7 @@ public class TypeChecker {
             @Override
             public TypeCheckerContext defineTypes(TypeCheckerContext context) {
                 var typedFieldNodes = new ArrayList<TypedStructFieldDefinitionNode>();
-                var fieldTypes = P.<String, Type>orderedMap();
+                var fieldInfos = new ArrayList<FieldInfo<Type>>();
 
                 for (var untypedField : untyped.fields()) {
                     var typedFieldNode = new TypedStructFieldDefinitionNode(
@@ -1314,12 +1307,12 @@ public class TypeChecker {
                         untypedField.source()
                     );
                     typedFieldNodes.add(typedFieldNode);
-                    fieldTypes = fieldTypes.plus(untypedField.name(), typedFieldNode.type().value());
+                    fieldInfos.add(new FieldInfo<>(untypedField.name(), typedFieldNode.type().value()));
                 }
 
                 var type = context.getLexicalState(StructType.class, untyped.name());
 
-                context = context.addStructTypeInfo(type, new StructTypeInfo(new FieldsInfo<>(fieldTypes)));
+                context = context.addStructTypeInfo(type, new StructTypeInfo(FieldsInfo.from(fieldInfos.stream())));
 
                 var typedNode = new TypedStructDefinitionNode(
                     type,
@@ -1626,7 +1619,7 @@ public class TypeChecker {
         var effectInfo = context.userDefinedEffectInfo(handledEffect, typedHandledEffect.source());
         // TODO: Don't create unnecessary list
         for (var effectField : effectInfo.fields().stream().toList()) {
-            throw new MissingEffectHandlerFieldError(effectField.getKey(), untyped.source());
+            throw new MissingEffectHandlerFieldError(effectField.name(), untyped.source());
         }
 
         for (var handlerField : untyped.handlers()) {
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 8209cd16..3f3147f0 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -9,7 +9,6 @@ import org.shedlang.compiler.typechecker.errors.*;
 import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
 import org.shedlang.compiler.util.PCollectors;
-import org.shedlang.compiler.util.POrderedMap;
 
 import java.util.*;
 import java.util.function.BiFunction;
@@ -408,12 +407,10 @@ public class TypeCheckerContext {
         return currentLexicalContext().freeVariables.build();
     }
 
-    public TypeCheckerContext addNamespaceFields(
+    public TypeCheckerContext addNamespaceTypeInfo(
         NamespaceName namespaceName,
-        POrderedMap<String, Type> fields
+        NamespaceTypeInfo namespaceTypeInfo
     ) {
-        var namespaceTypeInfo = new NamespaceTypeInfo(new FieldsInfo<>(fields));
-
         return new TypeCheckerContext(
             this.typeInfos.addNamespaceTypeInfo(namespaceName, namespaceTypeInfo),
             this.stack
diff --git a/src/main/java/org/shedlang/compiler/types/FieldInfo.java b/src/main/java/org/shedlang/compiler/types/FieldInfo.java
new file mode 100644
index 00000000..62a3e75b
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/FieldInfo.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.types;
+
+public record FieldInfo<T extends Type>(String name, T type) {
+}
diff --git a/src/main/java/org/shedlang/compiler/types/FieldsInfo.java b/src/main/java/org/shedlang/compiler/types/FieldsInfo.java
index afe76757..d42a3de5 100644
--- a/src/main/java/org/shedlang/compiler/types/FieldsInfo.java
+++ b/src/main/java/org/shedlang/compiler/types/FieldsInfo.java
@@ -1,13 +1,33 @@
 package org.shedlang.compiler.types;
 
+import org.pcollections.PVector;
 import org.shedlang.compiler.util.P;
+import org.shedlang.compiler.util.PCollectors;
 import org.shedlang.compiler.util.POrderedMap;
 
+import java.util.Arrays;
 import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Stream;
 
-public record FieldsInfo<T extends Type>(POrderedMap<String, T> fields) {
+public record FieldsInfo<T extends Type>(POrderedMap<String, FieldInfo<T>> fields) {
+    @SafeVarargs
+    public static <T extends Type> FieldsInfo<T> of(FieldInfo<T>... fields) {
+        return new FieldsInfo<>(
+            Arrays.stream(fields)
+                .map(x -> Map.entry(x.name(), x))
+                .collect(PCollectors.toOrderedMap())
+        );
+    }
+
+    public static <T extends Type> FieldsInfo<T> from(Stream<FieldInfo<T>> fields) {
+        return new FieldsInfo<>(
+            fields
+                .map(x -> Map.entry(x.name(), x))
+                .collect(PCollectors.toOrderedMap())
+        );
+    }
+
     public int fieldCount() {
         return fields.size();
     }
@@ -23,25 +43,28 @@ public record FieldsInfo<T extends Type>(POrderedMap<String, T> fields) {
     }
 
     public Optional<T> fieldType(String fieldName) {
-        return fields.get(fieldName);
+        return fields.get(fieldName).map(x -> x.type());
+    }
+
+    public PVector<String> keys() {
+        return fields.keys();
     }
 
-    // TODO: better element type
-    public Stream<Map.Entry<String, T>> stream() {
-        return fields.entryStream();
+    public Stream<FieldInfo<T>> stream() {
+        return fields.entryStream().map(x -> x.getValue());
     }
 
     public static <T extends Type> Builder<T> builder() {
         return new Builder<T>(P.orderedMap());
     }
 
-    public record Builder<T extends Type>(POrderedMap<String, T> fields) {
+    public record Builder<T extends Type>(POrderedMap<String, FieldInfo<T>> fields) {
         public FieldsInfo<T> build() {
             return new FieldsInfo<>(fields);
         }
 
         public Builder<T> addField(String name, T type) {
-            return new Builder<>(fields.plus(name, type));
+            return new Builder<>(fields.plus(name, new FieldInfo<>(name, type)));
         }
     }
 }