Type check return type of function type

3 files changed, 81 insertions, 5 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionTypeNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionTypeNode.java
new file mode 100644
index 00000000..bcab7ac5
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionTypeNode.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.FunctionType;
+
+public record TypedFunctionTypeNode(
+    FunctionType value,
+    Source source
+) implements TypedTypeLevelExpressionNode<FunctionType> {
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java
index 6d4d9502..ca8404e3 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java
@@ -1,6 +1,9 @@
 package org.shedlang.compiler.ast.untyped;
 
+import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.sources.Sources;
+import org.shedlang.compiler.util.P;
 
 import java.util.Optional;
 
@@ -10,4 +13,37 @@ public record UntypedFunctionTypeNode(
     UntypedTypeLevelExpressionNode returnType,
     Source source
 ) implements UntypedTypeLevelExpressionNode {
+    public static Builder builder() {
+        return new Builder(
+            P.vector(),
+            P.vector(),
+            Optional.empty(),
+            new UntypedTypeLevelReferenceNode("Unit", Sources.NULL)
+        );
+    }
+
+    public record Builder(
+        PVector<UntypedTypeLevelExpressionNode> positionalParams,
+        PVector<UntypedParamNode> namedParams,
+        Optional<UntypedTypeLevelExpressionNode> effect,
+        UntypedTypeLevelExpressionNode returnType
+    ) {
+        public UntypedFunctionTypeNode build() {
+            return new UntypedFunctionTypeNode(
+                new UntypedParamTypesNode(this.positionalParams, this.namedParams, Sources.NULL),
+                this.effect,
+                this.returnType,
+                Sources.NULL
+            );
+        }
+
+        public Builder withReturnType(UntypedTypeLevelReferenceNode returnType) {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                returnType
+            );
+        }
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 40647834..aeed2edd 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -390,6 +390,9 @@ public class TypeChecker {
             case UntypedFieldAccessNode untypedFieldAccess ->
                 typeCheckFieldAccess(untypedFieldAccess, context);
 
+            case UntypedFunctionExpressionNode untypedFunctionExpression ->
+                throw new UnsupportedOperationException("TODO");
+
             case UntypedIfNode untypedIf ->
                 typeCheckIf(untypedIf, context);
 
@@ -639,6 +642,22 @@ public class TypeChecker {
         };
     }
 
+    private static TypedFunctionTypeNode typeCheckFunctionType(
+        UntypedFunctionTypeNode untyped,
+        TypeCheckerContext context
+    ) {
+        return new TypedFunctionTypeNode(
+            new FunctionType(
+                List.of(),
+                List.of(),
+                Effects.NONE,
+                typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped.returnType(), context)
+                    .value()
+            ),
+            untyped.source()
+        );
+    }
+
     private static TypeCheckerResult<? extends TypedExpressionNode> typeCheckIf(
         UntypedIfNode untyped,
         TypeCheckerContext context
@@ -1013,18 +1032,29 @@ public class TypeChecker {
         TypeCheckerContext context
     ) {
         return switch (untyped) {
-            case UntypedFunctionTypeNode functionType ->
-                throw new UnsupportedOperationException("TODO");
+            case UntypedFunctionTypeNode functionType -> {
+                var typedFunctionTypeNode = typeCheckFunctionType(functionType, context);
+
+                if (valueType.valueClass().equals(Type.class)) {
+                    yield (TypedTypeLevelExpressionNode<T>) typedFunctionTypeNode;
+                } else {
+                    throw new UnexpectedTypeError(
+                        valueType,
+                        Types.metaType(typedFunctionTypeNode.value()),
+                        untyped.source()
+                    );
+                }
+            }
 
             case UntypedTupleTypeNode tuple -> {
-                var typedTupleNode = typeCheckTupleType(tuple, context);
+                var typedTupleTypeNode = typeCheckTupleType(tuple, context);
 
                 if (valueType.valueClass().equals(Type.class)) {
-                    yield (TypedTypeLevelExpressionNode<T>) typedTupleNode;
+                    yield (TypedTypeLevelExpressionNode<T>) typedTupleTypeNode;
                 } else {
                     throw new UnexpectedTypeError(
                         valueType,
-                        Types.metaType(typedTupleNode.value()),
+                        Types.metaType(typedTupleTypeNode.value()),
                         untyped.source()
                     );
                 }