Type check type-level params for function type

4 files changed, 91 insertions, 10 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 49072117..d7f15fc4 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1278,6 +1278,27 @@ public class TypeChecker {
         UntypedFunctionTypeNode untyped,
         TypeCheckerContext context
     ) {
+        var scopeResult = context.createFunctionTypeScope();
+        var scope = scopeResult.value();
+        context = scopeResult.context();
+
+        // TODO: better error if scope is not set
+        var typeLevelParams = untyped.typeLevelParams()
+            .map(untypedTypeLevelParams -> typeCheckTypeLevelParams(
+                untypedTypeLevelParams,
+                scope
+            ));
+
+        // TODO: Remove duplication with function signatures
+        context = context.enterFunctionSignature(
+            typeLevelParams.orElse(List.of()).stream()
+                .map(typeLevelParam -> Map.entry(
+                    typeLevelParam.name(),
+                    typeLevelParamToVariable(typeLevelParam)
+                ))
+                .collect(PCollectors.toMap())
+        );
+
         var typeCheckPositionalParamsResult = TypeCheckerResult.map(
             untyped.params().positional(),
             context,
@@ -1316,7 +1337,6 @@ public class TypeChecker {
         var effect = typeCheckEffectResult.value().orElse(Effects.NONE);
         context = typeCheckEffectResult.context();
 
-
         var typeCheckReturnTypeResult = typeCheckTypeExpression(
             untyped.returnType(),
             context
@@ -1324,9 +1344,11 @@ public class TypeChecker {
         var returnType = typeCheckReturnTypeResult.value().value();
         context = typeCheckEffectResult.context();
 
+        context = context.exit();
+
         return context.toResult(new TypedFunctionTypeNode(
             new FunctionType(
-                Optional.empty(),
+                typeLevelParams,
                 positionalParams,
                 namedParams,
                 effect,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 5dd39337..359fa410 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -24,7 +24,8 @@ public class TypeCheckerContext {
             Optional.of(ReturnTypeState.NOT_IN_FUNCTION),
             Set.of(),
             FreeVariables.builder(),
-            Empty.map()
+            Empty.map(),
+            0
         ))
     );
 
@@ -61,6 +62,7 @@ public class TypeCheckerContext {
         private final Set<String> boundVariables;
         private final FreeVariables.Builder freeVariables;
         private final PMap<String, Object> state;
+        private final int nextFunctionTypeIndex;
 
         public LexicalContext(
             Optional<NamespaceName> namespaceName,
@@ -69,7 +71,8 @@ public class TypeCheckerContext {
             Optional<ReturnTypeState> returnTypeState,
             Set<String> boundVariables,
             FreeVariables.Builder freeVariables,
-            PMap<String, Object> state
+            PMap<String, Object> state,
+            int nextFunctionTypeIndex
         ) {
             this.namespaceName = namespaceName;
             this.environment = environment;
@@ -78,6 +81,7 @@ public class TypeCheckerContext {
             this.boundVariables = boundVariables;
             this.freeVariables = freeVariables;
             this.state = state;
+            this.nextFunctionTypeIndex = nextFunctionTypeIndex;
         }
 
         public StackElement updateReturnTypeState(ReturnTypeState newReturnTypeState) {
@@ -88,7 +92,8 @@ public class TypeCheckerContext {
                 Optional.of(newReturnTypeState),
                 this.boundVariables,
                 this.freeVariables,
-                this.state
+                this.state,
+                this.nextFunctionTypeIndex
             );
         }
 
@@ -106,7 +111,8 @@ public class TypeCheckerContext {
                 this.returnTypeState,
                 this.boundVariables,
                 this.freeVariables.plus(name, referenceSource),
-                this.state
+                this.state,
+                this.nextFunctionTypeIndex
             );
         }
     }
@@ -151,6 +157,29 @@ public class TypeCheckerContext {
         );
     }
 
+    private TypeCheckerResult<LexicalContext> updateNamespaceLexicalContext(
+        Function<LexicalContext, LexicalContext> update
+    ) {
+        var newStack = this.stack;
+
+        for (var i = 0; i < this.stack.size(); i++) {
+            if (
+                this.stack.get(i) instanceof LexicalContext lexicalContext &&
+                    lexicalContext.namespaceName.isPresent()
+            ) {
+                newStack = newStack.with(i, update.apply(lexicalContext));
+
+                return new TypeCheckerContext(
+                    this.typeInfos,
+                    newStack
+                ).toResult(lexicalContext);
+            }
+        }
+
+        // TODO: better error
+        throw new RuntimeException("No namespace lexical context found");
+    }
+
     public TypeCheckerContext enterGlobal(PMap<String, Type> variables) {
         var environment = variables.entrySet().stream()
             .map(x -> Map.entry(x.getKey(), Variable.defined(x.getValue(), VariableType.VAL)))
@@ -242,7 +271,8 @@ public class TypeCheckerContext {
                 newReturnTypeState,
                 newEnvironment.keySet(),
                 FreeVariables.builder(),
-                Empty.map()
+                Empty.map(),
+                0
             ))
         );
     }
@@ -320,7 +350,8 @@ public class TypeCheckerContext {
                 lexicalContext.returnTypeState,
                 lexicalContext.boundVariables,
                 lexicalContext.freeVariables,
-                lexicalContext.state
+                lexicalContext.state,
+                lexicalContext.nextFunctionTypeIndex
             )
         );
     }
@@ -334,7 +365,8 @@ public class TypeCheckerContext {
                 lexicalContext.returnTypeState,
                 lexicalContext.boundVariables,
                 lexicalContext.freeVariables,
-                lexicalContext.state.plus(name, state)
+                lexicalContext.state.plus(name, state),
+                lexicalContext.nextFunctionTypeIndex
             )
         );
     }
@@ -378,7 +410,8 @@ public class TypeCheckerContext {
                 lexicalContext.returnTypeState,
                 lexicalContext.boundVariables,
                 lexicalContext.freeVariables,
-                lexicalContext.state
+                lexicalContext.state,
+                lexicalContext.nextFunctionTypeIndex
             )
         );
     }
@@ -489,6 +522,22 @@ public class TypeCheckerContext {
             .orElseThrow(() -> new MissingUserDefinedEffectInfoError(effect, source));
     }
 
+    public TypeCheckerResult<Scope> createFunctionTypeScope() {
+        return this.updateNamespaceLexicalContext(lexicalContext -> new LexicalContext(
+            lexicalContext.namespaceName,
+            lexicalContext.environment,
+            lexicalContext.effect,
+            lexicalContext.returnTypeState,
+            lexicalContext.boundVariables,
+            lexicalContext.freeVariables,
+            lexicalContext.state,
+            lexicalContext.nextFunctionTypeIndex + 1
+        )).map(lexicalContext -> new Scope.FunctionType(
+            lexicalContext.namespaceName.orElseThrow(),
+            lexicalContext.nextFunctionTypeIndex
+        ));
+    }
+
     public <T> TypeCheckerResult<T> toResult(T value) {
         return new TypeCheckerResult<>(value, this);
     }
diff --git a/src/main/java/org/shedlang/compiler/types/Scope.java b/src/main/java/org/shedlang/compiler/types/Scope.java
index 3dd66554..1a7fde3f 100644
--- a/src/main/java/org/shedlang/compiler/types/Scope.java
+++ b/src/main/java/org/shedlang/compiler/types/Scope.java
@@ -3,6 +3,7 @@ package org.shedlang.compiler.types;
 public sealed interface Scope {
     // TODO: handle methods
     record Function(NamespaceName namespaceName, String name) implements Scope {}
+    record FunctionType(NamespaceName namespaceName, int functionTypeIndex) implements Scope {}
     record NativeType(SimpleNativeType type) implements Scope {}
     record Struct(SimpleStructType type) implements Scope {}
     record Sum(SimpleSumType type) implements Scope {}
diff --git a/src/main/java/org/shedlang/compiler/types/TypeParam.java b/src/main/java/org/shedlang/compiler/types/TypeParam.java
index c996dd97..f76f5add 100644
--- a/src/main/java/org/shedlang/compiler/types/TypeParam.java
+++ b/src/main/java/org/shedlang/compiler/types/TypeParam.java
@@ -9,6 +9,15 @@ public record TypeParam(Scope scope, String name) implements Type, TypeLevelPara
         return new TypeParam(new Scope.Function(namespaceName, functionName), name);
     }
 
+    public static TypeParam forFunctionType(
+        NamespaceName namespaceName,
+        int functionTypeIndex,
+        String name
+    ) {
+        var scope = new Scope.FunctionType(namespaceName, functionTypeIndex);
+        return new TypeParam(scope, name);
+    }
+
     public static TypeParam forNativeType(SimpleNativeType nativeType, String name) {
         return new TypeParam(new Scope.NativeType(nativeType), name);
     }