Make param type optional in typed AST

3 files changed, 14 insertions, 16 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
index ad22b7ea..276ccb9a 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
@@ -40,7 +40,7 @@ public record TypedFunctionSignatureNode(
 
         public Builder addPositionalParam(String name, TypedTypeLevelExpressionNode<Type> type) {
             return new Builder(
-                this.positionalParams.plus(new TypedParamNode(name, type, Sources.NULL)),
+                this.positionalParams.plus(new TypedParamNode(name, Optional.of(type), Sources.NULL)),
                 this.namedParams,
                 this.effect,
                 this.returnType
@@ -50,7 +50,7 @@ public record TypedFunctionSignatureNode(
         public Builder addNamedParam(String name, TypedTypeLevelExpressionNode<Type> type) {
             return new Builder(
                 this.positionalParams,
-                this.namedParams.plus(new TypedParamNode(name, type, Sources.NULL)),
+                this.namedParams.plus(new TypedParamNode(name, Optional.of(type), Sources.NULL)),
                 this.effect,
                 this.returnType
             );
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedParamNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedParamNode.java
index 6c6b4765..2c4264cd 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedParamNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedParamNode.java
@@ -3,9 +3,11 @@ package org.shedlang.compiler.ast.typed;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.types.Type;
 
+import java.util.Optional;
+
 public record TypedParamNode(
     String name,
-    TypedTypeLevelExpressionNode<Type> type,
+    Optional<TypedTypeLevelExpressionNode<Type>> type,
     Source source
 ) implements TypedNode {
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index e0bc492b..e4b83516 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -474,7 +474,8 @@ public class TypeChecker {
         for (var param : typedSignature.params().all()) {
             bodyEnvironment = bodyEnvironment.add(
                 param.name(),
-                Variable.defined(param.type().value(), VariableType.PARAM),
+                // TODO: handle params without explicit type
+                Variable.defined(param.type().orElseThrow().value(), VariableType.PARAM),
                 param.source()
             );
         }
@@ -732,13 +733,13 @@ public class TypeChecker {
 
     private static List<Type> functionSignatureToPositionalParams(TypedFunctionSignatureNode signature) {
         return signature.params().positional().stream()
-            .map(param -> param.type().value())
+            .map(param -> param.type().orElseThrow().value())
             .toList();
     }
 
     private static List<NamedParam> functionSignatureToNamedParams(TypedFunctionSignatureNode signature) {
         return signature.params().named().stream()
-            .map(param -> new NamedParam(param.name(), param.type().value()))
+            .map(param -> new NamedParam(param.name(), param.type().orElseThrow().value()))
             .toList();
     }
 
@@ -784,10 +785,7 @@ public class TypeChecker {
             .toList();
 
         var namedParams = untyped.params().named().stream()
-            .map(param -> {
-                var typedParamNode = typeCheckFunctionTypeNamedParam(param, context);
-                return new NamedParam(typedParamNode.name(), typedParamNode.type().value());
-            })
+            .map(param -> typeCheckFunctionTypeNamedParam(param, context))
             .toList();
 
         var seenParamNames = new HashSet<String>();
@@ -820,14 +818,13 @@ public class TypeChecker {
         );
     }
 
-    private static TypedParamNode typeCheckFunctionTypeNamedParam(
+    private static NamedParam typeCheckFunctionTypeNamedParam(
         UntypedFunctionTypeNamedParamNode untyped,
         TypeCheckerContext context
     ) {
-        return new TypedParamNode(
+        return new NamedParam(
             untyped.name(),
-            typeCheckTypeExpression(untyped.type(), context),
-            untyped.source()
+            typeCheckTypeExpression(untyped.type(), context).value()
         );
     }
 
@@ -1046,8 +1043,7 @@ public class TypeChecker {
     ) {
         return new TypedParamNode(
             untyped.name(),
-            // TODO: handle param without type
-            typeCheckTypeExpression(untyped.type().orElseThrow(), context),
+            untyped.type().map(p -> typeCheckTypeExpression(p, context)),
             untyped.source()
         );
     }