Type check function expressions

4 files changed, 148 insertions, 3 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java
index 2f479f25..25422b1e 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java
@@ -4,6 +4,6 @@ import org.shedlang.compiler.types.Type;
 
 public sealed interface TypedExpressionNode
     extends TypedNode
-    permits TypedBinaryOperationNode, TypedBoolLiteralNode, TypedCallNode, TypedFieldAccessNode, TypedIfNode, TypedIntLiteralNode, TypedReferenceNode, TypedStringLiteralNode, TypedTupleNode, TypedWhenNode {
+    permits TypedBinaryOperationNode, TypedBoolLiteralNode, TypedCallNode, TypedFieldAccessNode, TypedFunctionExpressionNode, TypedIfNode, TypedIntLiteralNode, TypedReferenceNode, TypedStringLiteralNode, TypedTupleNode, TypedWhenNode {
     Type type();
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java
new file mode 100644
index 00000000..0f8ad4bc
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java
@@ -0,0 +1,119 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.pcollections.Empty;
+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.Type;
+
+public record TypedFunctionExpressionNode(
+    TypedFunctionSignatureNode signature,
+    TypedFunctionBodyNode body,
+    Source source
+) implements TypedExpressionNode {
+    @Override
+    public Type type() {
+        return signature.type();
+    }
+
+    public static Builder builder() {
+        return new Builder(
+            TypedFunctionSignatureNode.builder(),
+            FreeVariables.builder(),
+            Empty.vector(),
+            Empty.vector()
+        );
+    }
+
+    public record Builder(
+        TypedFunctionSignatureNode.Builder signature,
+        FreeVariables.Builder freeVariables,
+        PVector<String> boundVariableNames,
+        PVector<TypedFunctionStatementNode> body
+    ) {
+        public TypedFunctionExpressionNode build() {
+            return new TypedFunctionExpressionNode(
+                signature.build(),
+                new TypedFunctionBodyNode(
+                    freeVariables.build(),
+                    new TypedBlockNode(body, boundVariableNames, Sources.NULL),
+                    Sources.NULL
+                ),
+                Sources.NULL
+            );
+        }
+
+        public Builder withName(String name) {
+            return new Builder(
+                this.signature,
+                this.freeVariables,
+                this.boundVariableNames,
+                this.body
+            );
+        }
+
+        public Builder addPositionalParam(String name, TypedTypeLevelExpressionNode<Type> type) {
+            return new Builder(
+                this.signature.addPositionalParam(name, type),
+                this.freeVariables,
+                this.boundVariableNames,
+                this.body
+            ).addBoundVariable(name);
+        }
+
+        public Builder addNamedParam(String name, TypedTypeLevelExpressionNode<Type> type) {
+            return new Builder(
+                this.signature.addNamedParam(name, type),
+                this.freeVariables,
+                this.boundVariableNames,
+                this.body
+            ).addBoundVariable(name);
+        }
+
+        public Builder withEffect(TypedTypeLevelExpressionNode<Effect> effect) {
+            return new Builder(
+                this.signature.withEffect(effect),
+                this.freeVariables,
+                this.boundVariableNames,
+                this.body
+            );
+        }
+
+        public Builder withReturnType(TypedTypeLevelExpressionNode<Type> returnType) {
+            return new Builder(
+                this.signature.withReturnType(returnType),
+                this.freeVariables,
+                this.boundVariableNames,
+                this.body
+            );
+        }
+
+        public Builder addFreeVariable(String freeVariable) {
+            return new Builder(
+                this.signature,
+                this.freeVariables.plus(freeVariable, Sources.NULL),
+                this.boundVariableNames,
+                this.body
+            );
+        }
+
+        public Builder addBoundVariable(String name) {
+            return new Builder(
+                this.signature,
+                this.freeVariables,
+                this.boundVariableNames.plus(name),
+                this.body
+            );
+        }
+
+        public Builder addStatement(TypedFunctionStatementNode statement) {
+            return new Builder(
+                this.signature,
+                this.freeVariables,
+                this.boundVariableNames,
+                this.body.plus(statement)
+            );
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 2bd1229a..bb138c84 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -127,6 +127,9 @@ public class StackCompiler {
             case TypedFieldAccessNode typedFieldAccessNode ->
                 compileFieldAccess(typedFieldAccessNode, context);
 
+            case TypedFunctionExpressionNode typedFunctionExpressionNode ->
+                throw new UnsupportedOperationException("TODO");
+
             case TypedIfNode typedIfNode ->
                 compileIf(typedIfNode, context);
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index d51674a8..fac340ed 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -391,7 +391,7 @@ public class TypeChecker {
                 typeCheckFieldAccess(untypedFieldAccess, context);
 
             case UntypedFunctionExpressionNode untypedFunctionExpression ->
-                throw new UnsupportedOperationException("TODO");
+                typeCheckFunctionExpression(untypedFunctionExpression, context);
 
             case UntypedIfNode untypedIf ->
                 typeCheckIf(untypedIf, context);
@@ -464,7 +464,7 @@ public class TypeChecker {
         }
     }
 
-    private static TypeCheckerResult<TypedFunctionBodyNode> typeCheckFunctionBody(
+    static TypeCheckerResult<TypedFunctionBodyNode> typeCheckFunctionBody(
         UntypedBlockNode untyped,
         TypedFunctionSignatureNode typedSignature,
         TypeCheckerContext context
@@ -586,6 +586,29 @@ public class TypeChecker {
         return context.toResult(typed);
     }
 
+    private static TypeCheckerResult<TypedFunctionExpressionNode> typeCheckFunctionExpression(
+        UntypedFunctionExpressionNode untyped,
+        TypeCheckerContext context
+    ) {
+        var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
+
+        var typedBodyResult = typeCheckFunctionBody(
+            untyped.body(),
+            typedSignature,
+            context
+        );
+        var typedBody = typedBodyResult.value();
+        context = typedBodyResult.context();
+
+        var typed = new TypedFunctionExpressionNode(
+            typedSignature,
+            typedBody,
+            untyped.source()
+        );
+
+        return context.toResult(typed);
+    }
+
     private static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionReturn(
         UntypedFunctionReturnNode untyped,
         TypeCheckerContext context