Compile function definition into stack instructions

3 files changed, 159 insertions, 33 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
index cdfde02f..4bf840f9 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -1,8 +1,12 @@
 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;
+import org.shedlang.compiler.types.Types;
 
 import java.util.List;
 import java.util.Optional;
@@ -15,5 +19,100 @@ public record TypedFunctionDefinitionNode(
     List<TypedFunctionStatementNode> body,
     Source source
 ) implements TypedNamespaceStatementNode {
+    public static Builder builder() {
+        return new Builder(
+            "f",
+            Empty.vector(),
+            Empty.vector(),
+            Optional.empty(),
+            new TypedTypeLevelReferenceNode<>("Unit", Types.metaType(Types.UNIT), Sources.NULL),
+            Empty.vector()
+        );
+    }
 
+    public record Builder(
+        String name,
+        PVector<TypedParamNode> positionalParams,
+        PVector<TypedParamNode> namedParams,
+        Optional<TypedTypeLevelExpressionNode<Effect>> effect,
+        TypedTypeLevelExpressionNode<Type> returnType,
+        PVector<TypedFunctionStatementNode> body
+    ) {
+        public TypedFunctionDefinitionNode build() {
+            return new TypedFunctionDefinitionNode(
+                name,
+                new TypedParamsNode(positionalParams, namedParams, Sources.NULL),
+                effect,
+                returnType,
+                body,
+                Sources.NULL
+            );
+        }
+
+        public TypedFunctionDefinitionNode.Builder withName(String name) {
+            return new TypedFunctionDefinitionNode.Builder(
+                name,
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                this.body
+            );
+        }
+
+        public TypedFunctionDefinitionNode.Builder addPositionalParam(String name, TypedTypeLevelExpressionNode<Type> type) {
+            return new TypedFunctionDefinitionNode.Builder(
+                this.name,
+                this.positionalParams.plus(new TypedParamNode(name, type, Sources.NULL)),
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                this.body
+            );
+        }
+
+        public TypedFunctionDefinitionNode.Builder addNamedParam(String name, TypedTypeLevelExpressionNode<Type> type) {
+            return new TypedFunctionDefinitionNode.Builder(
+                this.name,
+                this.positionalParams,
+                this.namedParams.plus(new TypedParamNode(name, type, Sources.NULL)),
+                this.effect,
+                this.returnType,
+                this.body
+            );
+        }
+
+        public TypedFunctionDefinitionNode.Builder withEffect(TypedTypeLevelExpressionNode<Effect> effect) {
+            return new TypedFunctionDefinitionNode.Builder(
+                this.name,
+                this.positionalParams,
+                this.namedParams,
+                Optional.of(effect),
+                this.returnType,
+                this.body
+            );
+        }
+
+        public TypedFunctionDefinitionNode.Builder withReturnType(TypedTypeLevelExpressionNode<Type> returnType) {
+            return new TypedFunctionDefinitionNode.Builder(
+                this.name,
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                returnType,
+                this.body
+            );
+        }
+
+        public TypedFunctionDefinitionNode.Builder addStatement(TypedFunctionStatementNode statement) {
+            return new TypedFunctionDefinitionNode.Builder(
+                this.name,
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                this.body.plus(statement)
+            );
+        }
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
index a3c5447a..7c01256d 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
@@ -49,66 +49,66 @@ public record UntypedFunctionDefinitionNode(
         public Builder withName(String name) {
             return new Builder(
                 name,
-                positionalParams,
-                namedParams,
-                effect,
-                returnType,
-                body
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                this.body
             );
         }
 
         public Builder addPositionalParam(String name, UntypedTypeLevelExpressionNode type) {
             return new Builder(
-                name,
-                positionalParams.plus(new UntypedParamNode(name, type, Sources.NULL)),
-                namedParams,
-                effect,
-                returnType,
-                body
+                this.name,
+                this.positionalParams.plus(new UntypedParamNode(name, type, Sources.NULL)),
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                this.body
             );
         }
 
         public Builder addNamedParam(String name, UntypedTypeLevelExpressionNode type) {
             return new Builder(
-                name,
-                positionalParams,
-                namedParams.plus(new UntypedParamNode(name, type, Sources.NULL)),
-                effect,
-                returnType,
-                body
+                this.name,
+                this.positionalParams,
+                this.namedParams.plus(new UntypedParamNode(name, type, Sources.NULL)),
+                this.effect,
+                this.returnType,
+                this.body
             );
         }
 
         public Builder withEffect(UntypedTypeLevelExpressionNode effect) {
             return new Builder(
-                name,
-                positionalParams,
-                namedParams,
+                this.name,
+                this.positionalParams,
+                this.namedParams,
                 Optional.of(effect),
-                returnType,
-                body
+                this.returnType,
+                this.body
             );
         }
 
         public Builder withReturnType(UntypedTypeLevelExpressionNode returnType) {
             return new Builder(
-                name,
-                positionalParams,
-                namedParams,
-                effect,
+                this.name,
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
                 returnType,
-                body
+                this.body
             );
         }
 
         public Builder addStatement(UntypedFunctionStatementNode statement) {
             return new Builder(
-                name,
-                positionalParams,
-                namedParams,
-                effect,
-                returnType,
-                body.plus(statement)
+                this.name,
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                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 4623bc5d..406da12c 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -59,6 +59,26 @@ public class StackCompiler {
     private static PVector<Instruction> compileExpressionStatement(TypedExpressionStatementNode node) {
         return compileExpression(node.expression()).plus(new DiscardInstruction(node.source()));
     }
+
+    private static PVector<Instruction> compileFunctionDefinition(TypedFunctionDefinitionNode node) {
+        var defineFunctionInstruction = new DefineFunctionInstruction(
+            node.name(),
+            node.params().positional().stream()
+                .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
+                .toList(),
+            node.params().named().stream()
+                .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
+                .toList(),
+            node.body().stream()
+                .flatMap(statement -> compileFunctionStatement(statement).stream())
+                .toList(),
+            node.source()
+        );
+
+        var storeInstruction = new LocalStoreInstruction(node.name(), node.source());
+
+        return P.vector(defineFunctionInstruction, storeInstruction);
+    }
     
     public static PVector<Instruction> compileFunctionStatement(TypedFunctionStatementNode node) {
         return switch (node) {
@@ -71,6 +91,13 @@ public class StackCompiler {
         return P.vector(new PushIntInstruction(node.value(), node.source()));
     }
 
+    public static PVector<Instruction> compileNamespaceStatement(TypedNamespaceStatementNode node) {
+        return switch (node) {
+            case TypedFunctionDefinitionNode functionDefinition ->
+                compileFunctionDefinition(functionDefinition);
+        };
+    }
+
     private static PVector<Instruction> compileReference(TypedReferenceNode node) {
         return P.vector(new LocalLoadInstruction(node.name(), node.source()));
     }