Infer type of positional param from hint

5 files changed, 74 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java
index c624aab7..4944f4cd 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java
@@ -5,6 +5,8 @@ import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
 
+import java.util.Optional;
+
 public record UntypedFunctionExpressionNode(
     UntypedFunctionSignatureNode signature,
     UntypedBlockNode body,
@@ -30,6 +32,13 @@ public record UntypedFunctionExpressionNode(
         }
 
         public Builder addPositionalParam(String name, UntypedTypeLevelExpressionNode type) {
+            return addPositionalParam(name, Optional.of(type));
+        }
+
+        public Builder addPositionalParam(
+            String name,
+            Optional<UntypedTypeLevelExpressionNode> type
+        ) {
             return new Builder(
                 this.signature.addPositionalParam(name, type),
                 this.body
@@ -37,6 +46,13 @@ public record UntypedFunctionExpressionNode(
         }
 
         public Builder addNamedParam(String name, UntypedTypeLevelExpressionNode type) {
+            return addNamedParam(name, Optional.of(type));
+        }
+
+        public Builder addNamedParam(
+            String name,
+            Optional<UntypedTypeLevelExpressionNode> type
+        ) {
             return new Builder(
                 this.signature.addNamedParam(name, type),
                 this.body
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 d15c699d..641a9411 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
@@ -58,7 +58,11 @@ public record UntypedFunctionSignatureNode(
         }
 
         public Builder addNamedParam(String name, UntypedTypeLevelExpressionNode type) {
-            var paramNode = new UntypedParamNode(name, Optional.of(type), Sources.NULL);
+            return addNamedParam(name, Optional.of(type));
+        }
+
+        public Builder addNamedParam(String name, Optional<UntypedTypeLevelExpressionNode> type) {
+            var paramNode = new UntypedParamNode(name, type, Sources.NULL);
             return new Builder(
                 this.positionalParams,
                 this.namedParams.plus(paramNode),
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index ff416edc..eec4a8be 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -8,6 +8,8 @@ import org.shedlang.compiler.typechecker.errors.*;
 import org.shedlang.compiler.typechecker.typesets.TypeSetTypeLevelValue;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
 import org.shedlang.compiler.types.*;
+import org.shedlang.compiler.util.Lists;
+import org.shedlang.compiler.util.Optionals;
 import org.shedlang.compiler.util.P;
 import org.shedlang.compiler.util.PCollectors;
 
@@ -665,13 +667,36 @@ public class TypeChecker {
     ) {
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
-        var partialSignatureParams = typedSignature.params().all().stream()
-            .map(paramNode -> new PartialFunctionSignature.Param(
+        var functionTypeHint = Optionals.filterToType(FunctionType.class, hint);
+
+        var positionalParams = new ArrayList<PartialFunctionSignature.Param>();
+
+        for (var i = 0; i < typedSignature.params().positional().size(); i++) {
+            var paramIndex = i;
+            var paramNode = typedSignature.params().positional().get(paramIndex);
+            var paramHint = functionTypeHint.flatMap(h -> Lists.tryGet(h.positionalParams(), paramIndex));
+
+            var param = new PartialFunctionSignature.Param(
                 paramNode.name(),
-                paramNode.type().orElseThrow().value(),
+                paramNode.type()
+                    .map(x -> x.value())
+                    .or(() -> paramHint)
+                    // TODO: handle no explicit type nor hint
+                    .orElseThrow(),
                 paramNode.source()
-            ))
-            .toList();
+            );
+            positionalParams.add(param);
+        }
+
+        var partialSignatureParams = new ArrayList<>(positionalParams);
+
+        for (var namedParam : typedSignature.params().named()) {
+            partialSignatureParams.add(new PartialFunctionSignature.Param(
+                namedParam.name(),
+                namedParam.type().orElseThrow().value(),
+                namedParam.source()
+            ));
+        }
 
         var effect = functionSignatureToEffect(typedSignature);
 
@@ -690,7 +715,9 @@ public class TypeChecker {
         context = typedBodyResult.context();
 
         var type = new FunctionType(
-            functionSignatureToPositionalParams(typedSignature),
+            positionalParams.stream()
+                .map(x -> x.type())
+                .toList(),
             functionSignatureToNamedParams(typedSignature),
             effect,
             typedBody.type()
diff --git a/src/main/java/org/shedlang/compiler/util/Lists.java b/src/main/java/org/shedlang/compiler/util/Lists.java
index 527b5eb1..b33169df 100644
--- a/src/main/java/org/shedlang/compiler/util/Lists.java
+++ b/src/main/java/org/shedlang/compiler/util/Lists.java
@@ -16,6 +16,10 @@ public class Lists {
         return result;
     }
 
+    public static <T> Optional<T> tryGet(List<T> list, int index) {
+        return index < list.size() ? Optional.of(list.get(index)) : Optional.empty();
+    }
+
     public static <T> Optional<T> last(List<T> list) {
         if (list.isEmpty()) {
             return Optional.empty();
diff --git a/src/main/java/org/shedlang/compiler/util/Optionals.java b/src/main/java/org/shedlang/compiler/util/Optionals.java
new file mode 100644
index 00000000..58c92c09
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/util/Optionals.java
@@ -0,0 +1,16 @@
+package org.shedlang.compiler.util;
+
+import java.util.Optional;
+
+public class Optionals {
+    private Optionals() {
+    }
+
+    @SuppressWarnings("unchecked")
+    public static <T, R> Optional<R> filterToType(
+        Class<R> clazz,
+        Optional<T> optional
+    ) {
+        return (Optional<R>) optional.filter(x -> clazz.isInstance(x));
+    }
+}