Pass in namespace field types in one go

4 files changed, 46 insertions, 21 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 7e3a5153..259e517d 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -12,6 +12,7 @@ import org.shedlang.compiler.util.PCollectors;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Stream;
 
 public class TypeChecker {
     private TypeChecker() {
@@ -261,18 +262,15 @@ public class TypeChecker {
             untyped.source()
         );
 
-        var context2 = context;
-        for (var bodyStatement : typed.body()) {
-            if (bodyStatement instanceof TypedFunctionDefinitionNode functionDefinition) {
-                context2 = context2.addNamespaceField(
-                    untyped.namespaceName(),
-                    functionDefinition.name(),
-                    functionDefinition.type()
-                );
-            }
-        }
+        var fieldTypes = typed.body().stream()
+            .flatMap(bodyStatement ->
+                bodyStatement instanceof TypedFunctionDefinitionNode functionDefinition
+                    ? Stream.of(Map.<String, Type>entry(functionDefinition.name(), functionDefinition.type()))
+                    : Stream.of()
+            )
+            .collect(PCollectors.toOrderedMap());
 
-        return context2.toResult(typed);
+        return context.addNamespaceFields(typed.namespaceName(), fieldTypes).toResult(typed);
     }
 
     static TypedNamespaceStatementNode typeCheckNamespaceStatement(
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 93207cd3..b16169a7 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -71,22 +71,15 @@ public class TypeCheckerContext {
         return freeVariables;
     }
 
-    public TypeCheckerContext addNamespaceField(
+    public TypeCheckerContext addNamespaceFields(
         NamespaceName namespaceName,
-        String fieldName,
-        Type type
+        POrderedMap<String, Type> fields
     ) {
-        var namespaceFields = this.namespaceFields.plus(
-            namespaceName,
-            this.namespaceFields.getOrDefault(namespaceName, POrderedMap.empty())
-                .plus(fieldName, type)
-        );
-
         return new TypeCheckerContext(
             this.environment,
             this.boundVariables,
             this.freeVariables,
-            namespaceFields
+            namespaceFields.plus(namespaceName, fields)
         );
     }
 
diff --git a/src/main/java/org/shedlang/compiler/util/PCollectors.java b/src/main/java/org/shedlang/compiler/util/PCollectors.java
index 6c2d8208..26033670 100644
--- a/src/main/java/org/shedlang/compiler/util/PCollectors.java
+++ b/src/main/java/org/shedlang/compiler/util/PCollectors.java
@@ -35,4 +35,30 @@ public class PCollectors {
             return map;
         }
     }
+
+    public static <K, V> Collector<Map.Entry<K, V>, POrderedMapCollector<K, V>, POrderedMap<K, V>> toOrderedMap() {
+        return Collector.of(
+            POrderedMapCollector::new,
+            POrderedMapCollector::add,
+            POrderedMapCollector::addAll,
+            POrderedMapCollector::toOrderedMap
+        );
+    }
+
+    private static class POrderedMapCollector<K, V> {
+        private POrderedMap<K, V> map = POrderedMap.empty();
+
+        void add(Map.Entry<K, V> entry) {
+            map = map.plus(entry.getKey(), entry.getValue());
+        }
+
+        POrderedMapCollector<K, V> addAll(POrderedMapCollector<K, V> other) {
+            map = map.plusAll(other.map);
+            return this;
+        }
+
+        POrderedMap<K, V> toOrderedMap() {
+            return map;
+        }
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/util/POrderedMap.java b/src/main/java/org/shedlang/compiler/util/POrderedMap.java
index 6d273cdc..60bc6209 100644
--- a/src/main/java/org/shedlang/compiler/util/POrderedMap.java
+++ b/src/main/java/org/shedlang/compiler/util/POrderedMap.java
@@ -26,6 +26,14 @@ public class POrderedMap<K, V> {
         return new POrderedMap<>(map.plus(key, value), keys.plus(key));
     }
 
+    public POrderedMap<K, V> plusAll(POrderedMap<K, V> other) {
+        var newMap = this;
+        for (var otherKey : other.keys) {
+            newMap = newMap.plus(otherKey, other.map.get(otherKey));
+        }
+        return newMap;
+    }
+
     public boolean containsKey(K key) {
         return map.containsKey(key);
     }