Add params to environment in function

4 files changed, 69 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedParamsNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedParamsNode.java
index 2c864f90..f2601644 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedParamsNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedParamsNode.java
@@ -3,10 +3,14 @@ package org.shedlang.compiler.ast.typed;
 import org.shedlang.compiler.sources.Source;
 
 import java.util.List;
+import java.util.stream.Stream;
 
 public record TypedParamsNode(
     List<TypedParamNode> positional,
     List<TypedParamNode> named,
     Source source
 ) implements TypedNode {
+    public Stream<TypedParamNode> all() {
+        return Stream.concat(positional.stream(), named.stream());
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index b6bfb399..aa7ad39b 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -7,8 +7,10 @@ import org.shedlang.compiler.typechecker.errors.*;
 import org.shedlang.compiler.typechecker.typesets.TypeSetTypeLevelValue;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
 import org.shedlang.compiler.types.*;
+import org.shedlang.compiler.util.PCollectors;
 
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Optional;
 
 public class TypeChecker {
@@ -204,19 +206,24 @@ public class TypeChecker {
 
         var typedReturnType = typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped.returnType(), context);
 
-        var bodyResult = context.mapWithContext(
+        var bodyEnvironment = typedParams.all()
+            .map(param -> Map.entry(param.name(), param.type().value()))
+            .collect(PCollectors.toMap());
+        var bodyContext = context.enter(bodyEnvironment);
+
+        var bodyResult = bodyContext.mapWithContext(
             untyped.body(),
             TypeChecker::typeCheckFunctionStatement
         );
         var typedBody = bodyResult.value();
-        context = bodyResult.context();
+        bodyContext = bodyResult.context();
 
         return new TypedFunctionDefinitionNode(
             untyped.name(),
             typedParams,
             typedEffect,
             typedReturnType,
-            context.freeVariables(),
+            bodyContext.freeVariables(),
             typedBody,
             untyped.source()
         );
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 91d1daac..da51b959 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -14,24 +14,37 @@ import java.util.function.BiFunction;
 public class TypeCheckerContext {
     public static final TypeCheckerContext EMPTY = new TypeCheckerContext(
         Empty.map(),
+        Set.of(),
         Empty.orderedSet()
     );
 
     private final PMap<String, Type> environment;
+    private final Set<String> boundVariables;
     private final POrderedSet<String> freeVariables;
 
-    public TypeCheckerContext(PMap<String, Type> environment, POrderedSet<String> freeVariables) {
+    public TypeCheckerContext(
+        PMap<String, Type> environment,
+        Set<String> boundVariables,
+        POrderedSet<String> freeVariables
+    ) {
         this.environment = environment;
+        this.boundVariables = boundVariables;
         this.freeVariables = freeVariables;
     }
 
     public TypeCheckerContext enter(PMap<String, Type> environment) {
-        return new TypeCheckerContext(this.environment.plusAll(environment), Empty.orderedSet());
+        return new TypeCheckerContext(
+            this.environment.plusAll(environment),
+            environment.keySet(),
+            Empty.orderedSet()
+        );
     }
 
     public TypeCheckerResult<Optional<Type>> runtimeLookup(String name) {
         var type = Optional.ofNullable(environment.get(name));
-        return withFreeVariable(name).toResult(type);
+
+        var context = boundVariables.contains(name) ? this : withFreeVariable(name);
+        return context.toResult(type);
     }
 
     public Optional<Type> compileTimeLookup(String name) {
@@ -39,7 +52,7 @@ public class TypeCheckerContext {
     }
     
     private TypeCheckerContext withFreeVariable(String name) {
-        return new TypeCheckerContext(environment, freeVariables.plus(name));
+        return new TypeCheckerContext(environment, boundVariables, freeVariables.plus(name));
     }
 
     public <T> TypeCheckerResult<T> toResult(T value) {
diff --git a/src/main/java/org/shedlang/compiler/util/PCollectors.java b/src/main/java/org/shedlang/compiler/util/PCollectors.java
new file mode 100644
index 00000000..6c2d8208
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/util/PCollectors.java
@@ -0,0 +1,38 @@
+package org.shedlang.compiler.util;
+
+import org.pcollections.Empty;
+import org.pcollections.PMap;
+
+import java.util.Map;
+import java.util.stream.Collector;
+
+public class PCollectors {
+    private PCollectors() {
+    }
+
+    public static <K, V> Collector<Map.Entry<K, V>, PMapCollector<K, V>, PMap<K, V>> toMap() {
+        return Collector.of(
+            PMapCollector::new,
+            PMapCollector::add,
+            PMapCollector::addAll,
+            PMapCollector::toMap
+        );
+    }
+
+    private static class PMapCollector<K, V> {
+        private PMap<K, V> map = Empty.map();
+
+        void add(Map.Entry<K, V> entry) {
+            map = map.plus(entry.getKey(), entry.getValue());
+        }
+
+        PMapCollector<K, V> addAll(PMapCollector<K, V> other) {
+            map = map.plusAll(other.map);
+            return this;
+        }
+
+        PMap<K, V> toMap() {
+            return map;
+        }
+    }
+}