Add type-level params into function body environment

2 files changed, 41 insertions, 1 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
index e8e86885..74e4a11a 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
@@ -6,12 +6,15 @@ import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.types.Effect;
 import org.shedlang.compiler.types.Effects;
 import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.types.TypeLevelParam;
 import org.shedlang.compiler.util.Lists;
 import org.shedlang.compiler.util.P;
 
 import java.util.List;
+import java.util.Optional;
 
 public record PartialFunctionSignature(
+    Optional<List<TypeLevelParam>> typeLevelParams,
     List<Param> positionalParams,
     List<Param> namedParams,
     Effect effect,
@@ -29,6 +32,7 @@ public record PartialFunctionSignature(
 
     public PartialFunctionSignature withEffect(Effect effect) {
         return new PartialFunctionSignature(
+            this.typeLevelParams,
             this.positionalParams,
             this.namedParams,
             effect,
@@ -38,6 +42,7 @@ public record PartialFunctionSignature(
 
     public PartialFunctionSignature withReturnTypeState(ReturnTypeState returnTypeState) {
         return new PartialFunctionSignature(
+            this.typeLevelParams,
             this.positionalParams,
             this.namedParams,
             this.effect,
@@ -47,6 +52,7 @@ public record PartialFunctionSignature(
 
     public static Builder builder() {
         return new Builder(
+            Optional.empty(),
             P.vector(),
             P.vector(),
             Effects.NONE,
@@ -55,6 +61,7 @@ public record PartialFunctionSignature(
     }
 
     public record Builder(
+        Optional<List<TypeLevelParam>> typeLevelParams,
         PVector<Param> positionalParams,
         PVector<Param> namedParams,
         Effect effect,
@@ -62,6 +69,7 @@ public record PartialFunctionSignature(
     ) {
         public PartialFunctionSignature build() {
             return new PartialFunctionSignature(
+                typeLevelParams,
                 positionalParams,
                 namedParams,
                 effect,
@@ -69,8 +77,19 @@ public record PartialFunctionSignature(
             );
         }
 
+        public Builder withTypeLevelParams(List<TypeLevelParam> typeLevelParams) {
+            return new Builder(
+                Optional.of(typeLevelParams),
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                this.returnTypeState
+            );
+        }
+
         public Builder addPositionalParam(String name, Type type) {
             return new Builder(
+                this.typeLevelParams,
                 this.positionalParams.plus(new Param(name, type, Sources.NULL)),
                 this.namedParams,
                 this.effect,
@@ -80,6 +99,7 @@ public record PartialFunctionSignature(
 
         public Builder addNamedParam(String name, Type type) {
             return new Builder(
+                this.typeLevelParams,
                 this.positionalParams,
                 this.namedParams.plus(new Param(name, type, Sources.NULL)),
                 this.effect,
@@ -89,6 +109,7 @@ public record PartialFunctionSignature(
 
         public Builder withEffect(Effect effect) {
             return new Builder(
+                this.typeLevelParams,
                 this.positionalParams,
                 this.namedParams,
                 effect,
@@ -98,6 +119,7 @@ public record PartialFunctionSignature(
 
         public Builder withReturnType(Type returnType) {
             return new Builder(
+                this.typeLevelParams,
                 this.positionalParams,
                 this.namedParams,
                 this.effect,
@@ -107,6 +129,7 @@ public record PartialFunctionSignature(
 
         public Builder withoutReturnType() {
             return new Builder(
+                this.typeLevelParams,
                 this.positionalParams,
                 this.namedParams,
                 this.effect,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 13319913..8d325d0f 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -732,6 +732,15 @@ public class TypeChecker {
     ) {
         var bodyEnvironment = EnvironmentBuilder.EMPTY;
 
+        for (var typeLevelParam : signature.typeLevelParams().orElse(List.of())) {
+            bodyEnvironment = bodyEnvironment.add(
+                typeLevelParam.name(),
+                typeLevelParamToVariable(typeLevelParam),
+                // TODO: better source
+                untyped.source()
+            );
+        }
+
         for (var param : signature.params()) {
             bodyEnvironment = bodyEnvironment.add(
                 param.name(),
@@ -902,6 +911,7 @@ public class TypeChecker {
             .orElseThrow(() -> new FunctionSignatureMissingReturnTypeError(untyped.signature().source()));
 
         var partialSignature = new PartialFunctionSignature(
+            typedSignature.typeLevelParams(),
             positionalParams,
             namedParams,
             effect,
@@ -1078,7 +1088,7 @@ public class TypeChecker {
             typeLevelParams.orElse(List.of()).stream()
                 .map(typeLevelParam -> Map.entry(
                     typeLevelParam.name(),
-                    Variable.defined(Types.typeLevelValueType(typeLevelParam), VariableType.PARAM)
+                    typeLevelParamToVariable(typeLevelParam)
                 ))
                 .collect(PCollectors.toMap())
         );
@@ -1105,6 +1115,10 @@ public class TypeChecker {
         );
     }
 
+    private static Variable typeLevelParamToVariable(TypeLevelParam typeLevelParam) {
+        return Variable.defined(Types.typeLevelValueType(typeLevelParam), VariableType.PARAM);
+    }
+
     // TODO: lift out general Params record?
     private record InferredParamsTypes(
         List<PartialFunctionSignature.Param> positionalParams,
@@ -1171,6 +1185,7 @@ public class TypeChecker {
         var effect = functionSignatureToEffect(typedSignature);
 
         return new PartialFunctionSignature(
+            Optional.empty(),
             inferredParamTypes.positionalParams,
             inferredParamTypes.namedParams,
             effect,
@@ -2082,6 +2097,7 @@ public class TypeChecker {
             .orElse(expectedResumeType);
 
         var partialSignature = new PartialFunctionSignature(
+            Optional.empty(),
             inferredParamTypes.positionalParams,
             inferredParamTypes.namedParams,
             effect,
@@ -2145,6 +2161,7 @@ public class TypeChecker {
         List<UntypedTypeLevelParamNode> typeLevelParamNodes,
         Scope scope
     ) {
+        // TODO: check for duplicates
         return typeLevelParamNodes.stream()
             .map(typeLevelParamNode -> typeCheckTypeLevelParam(typeLevelParamNode, scope))
             .toList();