Add free variables to define function instruction

3 files changed, 15 insertions, 0 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 79d13152..9bec5e5c 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -115,6 +115,18 @@ public record TypedFunctionDefinitionNode(
             );
         }
 
+        public Builder addFreeVariable(String freeVariable) {
+            return new TypedFunctionDefinitionNode.Builder(
+                this.name,
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                this.freeVariables.plus(freeVariable),
+                this.body
+            );
+        }
+
         public TypedFunctionDefinitionNode.Builder addStatement(TypedFunctionStatementNode statement) {
             return new TypedFunctionDefinitionNode.Builder(
                 this.name,
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 55cd628f..aa59bc81 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -69,6 +69,7 @@ public class StackCompiler {
             node.params().named().stream()
                 .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
                 .toList(),
+            node.freeVariables(),
             node.body().stream()
                 .flatMap(statement -> compileFunctionStatement(statement).stream())
                 .toList(),
diff --git a/src/main/java/org/shedlang/compiler/stackir/instructions/DefineFunctionInstruction.java b/src/main/java/org/shedlang/compiler/stackir/instructions/DefineFunctionInstruction.java
index fef0583f..379f1f68 100644
--- a/src/main/java/org/shedlang/compiler/stackir/instructions/DefineFunctionInstruction.java
+++ b/src/main/java/org/shedlang/compiler/stackir/instructions/DefineFunctionInstruction.java
@@ -3,11 +3,13 @@ package org.shedlang.compiler.stackir.instructions;
 import org.shedlang.compiler.sources.Source;
 
 import java.util.List;
+import java.util.Set;
 
 public record DefineFunctionInstruction(
     String name,
     List<Param> positionalParams,
     List<Param> namedParams,
+    Set<String> freeVariables,
     List<Instruction> body,
     Source source
 ) implements Instruction {