Pass partial signature when type checking function body

4 files changed, 139 insertions, 23 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 f7273497..2a5ea89f 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedParamsNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedParamsNode.java
@@ -11,7 +11,7 @@ public record TypedParamsNode(
     List<TypedParamNode> named,
     Source source
 ) implements TypedNode {
-    public Iterable<TypedParamNode> all() {
+    public List<TypedParamNode> all() {
         return Stream.concat(positional.stream(), named.stream())
             .collect(Collectors.toList());
     }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedParamsNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedParamsNode.java
index 4f964402..c6ce76ad 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedParamsNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedParamsNode.java
@@ -11,7 +11,7 @@ public record UntypedParamsNode(
     List<UntypedParamNode> named,
     Source source
 ) implements UntypedNode {
-    public Iterable<UntypedParamNode> all() {
+    public List<UntypedParamNode> all() {
         return Stream.concat(positional.stream(), named.stream())
             .collect(Collectors.toList());
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
new file mode 100644
index 00000000..e86a2840
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
@@ -0,0 +1,82 @@
+package org.shedlang.compiler.typechecker;
+
+import org.pcollections.PVector;
+import org.shedlang.compiler.sources.Source;
+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.util.P;
+
+import java.util.List;
+import java.util.Optional;
+
+public record PartialFunctionSignature(
+    List<Param> params,
+    Effect effect,
+    Optional<Type> returnType
+) {
+    public record Param(String name, Type type, Source source) {}
+
+    public static Builder builder() {
+        return new Builder(
+            P.vector(),
+            Effects.NONE,
+            Optional.empty()
+        );
+    }
+
+    public record Builder(
+        PVector<Param> params,
+        Effect effect,
+        Optional<Type> returnType
+    ) {
+        public PartialFunctionSignature build() {
+            return new PartialFunctionSignature(
+                params,
+                effect,
+                returnType
+            );
+        }
+
+        public Builder addPositionalParam(String name, Type type) {
+            return new Builder(
+                this.params.plus(new Param(name, type, Sources.NULL)),
+                this.effect,
+                this.returnType
+            );
+        }
+
+        public Builder addNamedParam(String name, Type type) {
+            return new Builder(
+                this.params.plus(new Param(name, type, Sources.NULL)),
+                this.effect,
+                this.returnType
+            );
+        }
+
+        public Builder withEffect(Effect effect) {
+            return new Builder(
+                this.params,
+                effect,
+                this.returnType
+            );
+        }
+
+        public Builder withReturnType(Type returnType) {
+            return new Builder(
+                this.params,
+                this.effect,
+                Optional.of(returnType)
+            );
+        }
+
+        public Builder withoutReturnType() {
+            return new Builder(
+                this.params,
+                this.effect,
+                Optional.empty()
+            );
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index e4b83516..58285e58 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -466,16 +466,15 @@ public class TypeChecker {
 
     static TypeCheckerResult<TypedFunctionBodyNode> typeCheckFunctionBody(
         UntypedBlockNode untyped,
-        TypedFunctionSignatureNode typedSignature,
+        PartialFunctionSignature signature,
         TypeCheckerContext context
     ) {
         var bodyEnvironment = EnvironmentBuilder.EMPTY;
 
-        for (var param : typedSignature.params().all()) {
+        for (var param : signature.params()) {
             bodyEnvironment = bodyEnvironment.add(
                 param.name(),
-                // TODO: handle params without explicit type
-                Variable.defined(param.type().orElseThrow().value(), VariableType.PARAM),
+                Variable.defined(param.type(), VariableType.PARAM),
                 param.source()
             );
         }
@@ -487,11 +486,11 @@ public class TypeChecker {
         );
         bodyEnvironment = bodyEnvironment.addUndefined(declaredVariables);
 
-        var explicitReturnType = functionSignatureToReturnType(typedSignature);
+        var explicitReturnType = signature.returnType();
 
         context = context.enterFunction(
             bodyEnvironment.build(),
-            functionSignatureToEffect(typedSignature),
+            signature.effect(),
             explicitReturnType
         );
 
@@ -569,7 +568,8 @@ public class TypeChecker {
     }
 
     private record TypeCheckFunctionDefinitionState(
-        TypedFunctionSignatureNode signature,
+        TypedFunctionSignatureNode typedSignatureNode,
+        PartialFunctionSignature partialSignature,
         FunctionType type
     ) {
     }
@@ -578,23 +578,41 @@ public class TypeChecker {
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context
     ) {
-        for (var param : untyped.signature().params().all()) {
-            if (param.type().isEmpty()) {
-                throw new FunctionSignatureMissingParamTypeError(param.source());
-            }
-        }
-
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
+        var partialSignatureParams = typedSignature.params().all().stream()
+            .map(param -> {
+                if (param.type().isEmpty()) {
+                    throw new FunctionSignatureMissingParamTypeError(param.source());
+                } else {
+                    return new PartialFunctionSignature.Param(
+                        param.name(),
+                        param.type().get().value(),
+                        param.source()
+                    );
+                }
+            })
+            .toList();
+
+        var effect = functionSignatureToEffect(typedSignature);
+
+        var returnType = functionSignatureToReturnType(typedSignature)
+            .orElseThrow(() -> new FunctionSignatureMissingReturnTypeError(untyped.signature().source()));
+
+        var partialSignature = new PartialFunctionSignature(
+            partialSignatureParams,
+            effect,
+            Optional.of(returnType)
+        );
+
         var type = new FunctionType(
             functionSignatureToPositionalParams(typedSignature),
             functionSignatureToNamedParams(typedSignature),
-            functionSignatureToEffect(typedSignature),
-            functionSignatureToReturnType(typedSignature)
-                .orElseThrow(() -> new FunctionSignatureMissingReturnTypeError(untyped.signature().source()))
+            effect,
+            returnType
         );
 
-        var state = new TypeCheckFunctionDefinitionState(typedSignature, type);
+        var state = new TypeCheckFunctionDefinitionState(typedSignature, partialSignature, type);
 
         return context
             .declareVariableType(untyped.name(), type, untyped.source())
@@ -606,12 +624,12 @@ public class TypeChecker {
         TypeCheckerContext context
     ) {
         var state = context.getLexicalState(TypeCheckFunctionDefinitionState.class, untyped.name());
-        var typedSignature = state.signature;
+        var typedSignature = state.typedSignatureNode;
         var type = state.type;
 
         var typedBodyResult = typeCheckFunctionBody(
             untyped.body(),
-            typedSignature,
+            state.partialSignature,
             context
         );
         var typedBody = typedBodyResult.value();
@@ -634,9 +652,25 @@ public class TypeChecker {
     ) {
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
+        var partialSignatureParams = typedSignature.params().all().stream()
+            .map(paramNode -> new PartialFunctionSignature.Param(
+                paramNode.name(),
+                paramNode.type().orElseThrow().value(),
+                paramNode.source()
+            ))
+            .toList();
+
+        var effect = functionSignatureToEffect(typedSignature);
+
+        var partialSignature = new PartialFunctionSignature(
+            partialSignatureParams,
+            effect,
+            typedSignature.returnType().map(node -> node.value())
+        );
+
         var typedBodyResult = typeCheckFunctionBody(
             untyped.body(),
-            typedSignature,
+            partialSignature,
             context
         );
         var typedBody = typedBodyResult.value();
@@ -645,7 +679,7 @@ public class TypeChecker {
         var type = new FunctionType(
             functionSignatureToPositionalParams(typedSignature),
             functionSignatureToNamedParams(typedSignature),
-            functionSignatureToEffect(typedSignature),
+            effect,
             typedBody.type()
         );