Use type level params when type checking function signature

4 files changed, 80 insertions, 10 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
index 42795805..73a57310 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
@@ -44,6 +44,17 @@ public record UntypedFunctionSignatureNode(
             );
         }
 
+        public Builder withTypeLevelParams(List<UntypedTypeLevelParamNode> typeLevelParams) {
+            return new Builder(
+                Optional.of(typeLevelParams),
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                this.body
+            );
+        }
+
         public Builder addPositionalParam(String name, UntypedTypeLevelExpressionNode type) {
             return addPositionalParam(name, Optional.of(type));
         }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 8dba748f..cc8b799f 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -855,7 +855,13 @@ public class TypeChecker {
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context
     ) {
-        var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
+        var namespaceName = context.namespaceName()
+            .orElseThrow(() -> new CannotDefineFunctionOutsideOfNamespaceError(untyped.source()));
+        var typedSignature = typeCheckFunctionSignature(
+            untyped.signature(),
+            Optional.of(new Scope.Function(namespaceName, untyped.name())),
+            context
+        );
 
         Function<TypedParamNode, PartialFunctionSignature.Param> typedParamNodeToParam = param ->
             new PartialFunctionSignature.Param(
@@ -937,7 +943,11 @@ public class TypeChecker {
         Optional<Type> hint,
         TypeCheckerContext context
     ) {
-        var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
+        var typedSignature = typeCheckFunctionSignature(
+            untyped.signature(),
+            Optional.empty(),
+            context
+        );
 
         var partialSignature = inferFunctionPartialSignature(typedSignature, hint);
 
@@ -1037,17 +1047,38 @@ public class TypeChecker {
 
     static TypedFunctionSignatureNode typeCheckFunctionSignature(
         UntypedFunctionSignatureNode untyped,
+        Optional<Scope> scope,
         TypeCheckerContext context
     ) {
+        // TODO: better error if scope is not set
+        var typeLevelParams = untyped.typeLevelParams()
+            .map(untypedTypeLevelParams -> typeCheckTypeLevelParams(
+                untypedTypeLevelParams,
+                scope.orElseThrow()
+            ))
+            .orElse(List.of());
+
+        // TODO: Add type level params to function body
+        context = context.enterFunctionSignature(
+            typeLevelParams.stream()
+                .map(typeLevelParam -> Map.entry(
+                    typeLevelParam.name(),
+                    Variable.defined(Types.metaType(typeLevelParam), VariableType.PARAM)
+                ))
+                .collect(PCollectors.toMap())
+        );
+
+        var functionSignatureContext = context;
+
         var typedParams = typeCheckParams(untyped.params(), context);
 
         var typedEffect = untyped.effect()
-            .map(untypedEffect -> typeCheckEffect(untypedEffect, context));
+            .map(untypedEffect -> typeCheckEffect(untypedEffect, functionSignatureContext));
 
         var typedReturnType = untyped.returnType()
             .map(untypedReturnType -> typeCheckTypeExpression(
                 untypedReturnType,
-                context
+                functionSignatureContext
             ));
 
         return new TypedFunctionSignatureNode(
@@ -2067,6 +2098,18 @@ public class TypeChecker {
         return context.toResult(typed);
     }
 
+    private static List<TypeParam> typeCheckTypeLevelParams(
+        List<UntypedTypeLevelParamNode> typeLevelParamNodes,
+        Scope scope
+    ) {
+        return typeLevelParamNodes.stream()
+            .map(typeLevelParamNode -> new TypeParam(
+                scope,
+                typeLevelParamNode.name()
+            ))
+            .toList();
+    }
+
     private static TypeCheckerResult<TypedExpressionNode> typeCheckUnitLiteral(
         UntypedUnitLiteralNode untyped,
         TypeCheckerContext context
@@ -2258,12 +2301,7 @@ public class TypeChecker {
     ) {
         return untypedTypeLevelParamNodes
             .map(typeLevelParamNodes -> TypeOrConstructor.constructor(new TypeConstructor<>(
-                typeLevelParamNodes.stream()
-                    .map(typeLevelParamNode -> new TypeParam(
-                        scope,
-                        typeLevelParamNode.name()
-                    ))
-                    .toList(),
+                typeCheckTypeLevelParams(typeLevelParamNodes, scope),
                 type
             )))
             .orElseGet(() -> TypeOrConstructor.type(type));
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 557b05d7..5dd39337 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -163,6 +163,17 @@ public class TypeCheckerContext {
         return enter(Optional.of(namespaceName), environment, Effects.NONE, Optional.of(ReturnTypeState.NOT_IN_FUNCTION));
     }
 
+    public TypeCheckerContext enterFunctionSignature(
+        PMap<String, Variable> environment
+    ) {
+        return enter(
+            this.namespaceName(),
+            environment,
+            Effects.NONE,
+            Optional.of(ReturnTypeState.NOT_IN_FUNCTION)
+        );
+    }
+
     public TypeCheckerContext enterFunction(
         PMap<String, Variable> environment,
         Effect effect,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/errors/CannotDefineFunctionOutsideOfNamespaceError.java b/src/main/java/org/shedlang/compiler/typechecker/errors/CannotDefineFunctionOutsideOfNamespaceError.java
new file mode 100644
index 00000000..4b52e4d8
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/errors/CannotDefineFunctionOutsideOfNamespaceError.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.typechecker.errors;
+
+import org.shedlang.compiler.errors.InternalCompilerError;
+import org.shedlang.compiler.sources.Source;
+
+public class CannotDefineFunctionOutsideOfNamespaceError extends InternalCompilerError {
+    public CannotDefineFunctionOutsideOfNamespaceError(Source source) {
+        super("Cannot define function outside of namespace", source);
+    }
+}