Store types of fields on namespace node

5 files changed, 89 insertions, 12 deletions

diff --git a/src/main/java/org/shedlang/compiler/ShedCompiler.java b/src/main/java/org/shedlang/compiler/ShedCompiler.java
index dd553f45..6f252776 100644
--- a/src/main/java/org/shedlang/compiler/ShedCompiler.java
+++ b/src/main/java/org/shedlang/compiler/ShedCompiler.java
@@ -65,7 +65,7 @@ public class ShedCompiler {
                 ));
             var typedNode = TypeChecker.typeCheckNamespace(untypedNode, context);
 
-            var instructions = StackCompiler.compileNamespace(typedNode);
+            var instructions = StackCompiler.compileNamespace(typedNode.value());
             namespaces.put(namespaceName, instructions);
         }
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 4e89e97f..7e3a5153 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -249,17 +249,30 @@ public class TypeChecker {
         ));
     }
 
-    public static TypedNamespaceNode typeCheckNamespace(
+    public static TypeCheckerResult<TypedNamespaceNode> typeCheckNamespace(
         UntypedNamespaceNode untyped,
         TypeCheckerContext context
     ) {
-        return new TypedNamespaceNode(
+        var typed = new TypedNamespaceNode(
             untyped.namespaceName(),
             untyped.body().stream()
                 .map(statement -> typeCheckNamespaceStatement(statement, context))
                 .toList(),
             untyped.source()
         );
+
+        var context2 = context;
+        for (var bodyStatement : typed.body()) {
+            if (bodyStatement instanceof TypedFunctionDefinitionNode functionDefinition) {
+                context2 = context2.addNamespaceField(
+                    untyped.namespaceName(),
+                    functionDefinition.name(),
+                    functionDefinition.type()
+                );
+            }
+        }
+
+        return context2.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 da51b959..3429be63 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -3,7 +3,9 @@ package org.shedlang.compiler.typechecker;
 import org.pcollections.Empty;
 import org.pcollections.PMap;
 import org.pcollections.POrderedSet;
+import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.util.POrderedMap;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -15,28 +17,33 @@ public class TypeCheckerContext {
     public static final TypeCheckerContext EMPTY = new TypeCheckerContext(
         Empty.map(),
         Set.of(),
-        Empty.orderedSet()
+        Empty.orderedSet(),
+        Empty.map()
     );
 
     private final PMap<String, Type> environment;
     private final Set<String> boundVariables;
     private final POrderedSet<String> freeVariables;
+    private final PMap<NamespaceName, POrderedMap<String, Type>> namespaceFields;
 
     public TypeCheckerContext(
         PMap<String, Type> environment,
         Set<String> boundVariables,
-        POrderedSet<String> freeVariables
+        POrderedSet<String> freeVariables,
+        PMap<NamespaceName, POrderedMap<String, Type>> namespaceFields
     ) {
         this.environment = environment;
         this.boundVariables = boundVariables;
         this.freeVariables = freeVariables;
+        this.namespaceFields = namespaceFields;
     }
 
     public TypeCheckerContext enter(PMap<String, Type> environment) {
         return new TypeCheckerContext(
             this.environment.plusAll(environment),
             environment.keySet(),
-            Empty.orderedSet()
+            Empty.orderedSet(),
+            this.namespaceFields
         );
     }
 
@@ -52,7 +59,39 @@ public class TypeCheckerContext {
     }
     
     private TypeCheckerContext withFreeVariable(String name) {
-        return new TypeCheckerContext(environment, boundVariables, freeVariables.plus(name));
+        return new TypeCheckerContext(
+            environment,
+            boundVariables,
+            freeVariables.plus(name),
+            this.namespaceFields
+        );
+    }
+
+    public Set<String> freeVariables() {
+        return freeVariables;
+    }
+
+    public TypeCheckerContext addNamespaceField(
+        NamespaceName namespaceName,
+        String fieldName,
+        Type type
+    ) {
+        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
+        );
+    }
+
+    public PMap<NamespaceName, POrderedMap<String, Type>> namespaceFields() {
+        return namespaceFields;
     }
 
     public <T> TypeCheckerResult<T> toResult(T value) {
@@ -74,8 +113,4 @@ public class TypeCheckerContext {
 
         return new TypeCheckerResult<>(mappedValues, context);
     }
-
-    public Set<String> freeVariables() {
-        return freeVariables;
-    }
 }
diff --git a/src/main/java/org/shedlang/compiler/util/P.java b/src/main/java/org/shedlang/compiler/util/P.java
index a09d7940..fba3c512 100644
--- a/src/main/java/org/shedlang/compiler/util/P.java
+++ b/src/main/java/org/shedlang/compiler/util/P.java
@@ -22,6 +22,15 @@ public class P {
     }
 
     @SafeVarargs
+    public static <K, V> POrderedMap<K, V> orderedMap(Map.Entry<K, V>... entries) {
+        var map = POrderedMap.<K, V>empty();
+        for (var entry : entries) {
+            map = map.plus(entry.getKey(), entry.getValue());
+        }
+        return map;
+    }
+
+    @SafeVarargs
     public static <T> PVector<T> vector(T... values) {
         return TreePVector.from(Arrays.asList(values));
     }
diff --git a/src/main/java/org/shedlang/compiler/util/POrderedMap.java b/src/main/java/org/shedlang/compiler/util/POrderedMap.java
index 400a4413..aabec1d7 100644
--- a/src/main/java/org/shedlang/compiler/util/POrderedMap.java
+++ b/src/main/java/org/shedlang/compiler/util/POrderedMap.java
@@ -18,7 +18,6 @@ public class POrderedMap<K, V> {
         PMap<K, V> map,
         PVector<K> keys
     ) {
-
         this.map = map;
         this.keys = keys;
     }
@@ -34,4 +33,25 @@ public class POrderedMap<K, V> {
     public Collection<V> values() {
         return map.values();
     }
+
+    @Override
+    public boolean equals(Object other) {
+        if (this == other) {
+            return true;
+        }
+        if (other == null) {
+            return false;
+        }
+        if (!(other instanceof POrderedMap<?, ?> otherMap)) {
+            return false;
+        }
+        return map.equals(otherMap.map) && keys.equals(otherMap.keys);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = map.hashCode();
+        result = 31 * result + keys.hashCode();
+        return result;
+    }
 }