Type check type-level params for function type

5 files changed, 138 insertions, 19 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);
     }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionTypeTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionTypeTests.java
index f4a2a6a2..22d994ea 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionTypeTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionTypeTests.java
@@ -5,11 +5,10 @@ import org.shedlang.compiler.ast.untyped.Untyped;
 import org.shedlang.compiler.ast.untyped.UntypedFunctionTypeNode;
 import org.shedlang.compiler.typechecker.errors.DuplicateVariableNameError;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
-import org.shedlang.compiler.types.Effects;
-import org.shedlang.compiler.types.FunctionType;
-import org.shedlang.compiler.types.Types;
+import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
 
+import java.util.List;
 import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -18,6 +17,8 @@ import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.equalTo;
 
 public class TypeCheckerFunctionTypeTests {
+    private static final NamespaceName NAMESPACE_NAME = NamespaceName.of("Example", "Project");
+
     @Test
     public void returnTypeIsTypeChecked() {
         var untyped = UntypedFunctionTypeNode.builder()
@@ -26,7 +27,8 @@ public class TypeCheckerFunctionTypeTests {
         var context = TypeCheckerContext.EMPTY
             .enterGlobal(P.map(
                 Map.entry("String", Types.metaType(Types.STRING))
-            ));
+            ))
+            .enterNamespace(NAMESPACE_NAME, P.map());
 
         var result = TypeChecker.typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped, context);
 
@@ -48,7 +50,8 @@ public class TypeCheckerFunctionTypeTests {
         var context = TypeCheckerContext.EMPTY
             .enterGlobal(P.map(
                 Map.entry("String", Types.metaType(Types.STRING))
-            ));
+            ))
+            .enterNamespace(NAMESPACE_NAME, P.map());
 
         var result = TypeChecker.typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped, context);
 
@@ -73,7 +76,8 @@ public class TypeCheckerFunctionTypeTests {
             .enterGlobal(P.map(
                 Map.entry("Io", Types.effectType(Effects.IO)),
                 Map.entry("String", Types.metaType(Types.STRING))
-            ));
+            ))
+            .enterNamespace(NAMESPACE_NAME, P.map());
 
         var result = TypeChecker.typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped, context);
 
@@ -100,7 +104,8 @@ public class TypeCheckerFunctionTypeTests {
                 Map.entry("Bool", Types.metaType(Types.BOOL)),
                 Map.entry("Int", Types.metaType(Types.INT)),
                 Map.entry("String", Types.metaType(Types.STRING))
-            ));
+            ))
+            .enterNamespace(NAMESPACE_NAME, P.map());
 
         var result = TypeChecker.typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped, context);
 
@@ -128,7 +133,8 @@ public class TypeCheckerFunctionTypeTests {
                 Map.entry("Bool", Types.metaType(Types.BOOL)),
                 Map.entry("Int", Types.metaType(Types.INT)),
                 Map.entry("String", Types.metaType(Types.STRING))
-            ));
+            ))
+            .enterNamespace(NAMESPACE_NAME, P.map());
 
         var result = TypeChecker.typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped, context);
 
@@ -156,7 +162,8 @@ public class TypeCheckerFunctionTypeTests {
                 Map.entry("Bool", Types.metaType(Types.BOOL)),
                 Map.entry("Int", Types.metaType(Types.INT)),
                 Map.entry("String", Types.metaType(Types.STRING))
-            ));
+            ))
+            .enterNamespace(NAMESPACE_NAME, P.map());
 
         var error = assertThrows(
             DuplicateVariableNameError.class,
@@ -165,4 +172,35 @@ public class TypeCheckerFunctionTypeTests {
 
         assertThat(error.variableName(), equalTo("a"));
     }
+
+    @Test
+    public void typeLevelParamsAreTypeChecked() {
+        var namespaceName = NamespaceName.of("Example", "Project");
+        var untyped = UntypedFunctionTypeNode.builder()
+            .withTypeLevelParams(List.of(
+                Untyped.typeParam("T"),
+                Untyped.typeParam("R")
+            ))
+            .addPositionalParam(Untyped.typeLevelReference("T"))
+            .withReturnType(Untyped.typeLevelReference("R"))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterNamespace(NAMESPACE_NAME, P.map());
+
+        var result = TypeChecker.typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped, context);
+
+        assertThat(
+            result.value(),
+            isTypedFunctionTypeNode(equalTo(
+                FunctionType.builder()
+                    .withTypeLevelParams(List.of(
+                        TypeParam.forFunctionType(namespaceName, 0, "T"),
+                        TypeParam.forFunctionType(namespaceName, 0, "R")
+                    ))
+                    .addPositionalParam(TypeParam.forFunctionType(namespaceName, 0, "T"))
+                    .withReturnType(TypeParam.forFunctionType(namespaceName, 0, "R"))
+                    .build()
+            ))
+        );
+    }
 }