Extract functions to store and load free variables

1 files changed, 56 insertions, 21 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index abc89e10..9e5b2b1d 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -96,24 +96,14 @@ class PatioCodeGenClosures {
                 bodyContext = bodyContext.defineShidParam(shidParam.name());
             }
 
-            for (var freeVariable : freeVariables) {
-                var freeVariableLocalResult = bodyContext.declareShidVal(
-                    freeVariable.name(),
-                    freeVariableType(freeVariable, bodyContextParent)
-                );
-                var freeVariableLocal = freeVariableLocalResult.value();
-                bodyContext = freeVariableLocalResult.context();
-
-                bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
-                    freeVariableLocal,
-                    closurePointer,
-                    layout,
-                    freeVariableOffset(freeVariable),
-                    source
-                ));
-
-                bodyContext = bodyContext.initShidLocal(freeVariable.name());
-            }
+            bodyContext = loadFreeVariables(
+                freeVariables,
+                closurePointer,
+                layout,
+                source,
+                bodyContext,
+                bodyContextParent
+            );
 
             bodyContext = generateBody.apply(bodyContext);
 
@@ -144,23 +134,68 @@ class PatioCodeGenClosures {
             source
         ));
 
+        context = storeFreeVariables(
+            freeVariables,
+            closurePointer,
+            layout,
+            source,
+            context
+        );
+
+        return context;
+    }
+
+    private static PatioCodeGenContext storeFreeVariables(
+        FreeVariables freeVariables,
+        PatioLocalRef closurePointer,
+        PatioDataLayout closureLayout,
+        Source source,
+        PatioCodeGenContext context
+    ) {
         for (var freeVariable : freeVariables) {
             context = context.onceShidLocalDefined(
                 freeVariable.name(),
                 c -> c.addInstruction(new PatioMemoryStore(
                     closurePointer,
-                    layout,
+                    closureLayout,
                     freeVariableOffset(freeVariable),
                     c.shidLocalNameToLocal(freeVariable.name()),
                     source
                 ))
             );
-
         }
-
         return context;
     }
 
+    private static PatioCodeGenContext loadFreeVariables(
+        FreeVariables freeVariables,
+        PatioLocalRef closurePointer,
+        PatioDataLayout closureLayout,
+        Source source,
+        PatioCodeGenContext bodyContext,
+        PatioCodeGenContext bodyContextParent
+    ) {
+        for (var freeVariable : freeVariables) {
+            var freeVariableLocalResult = bodyContext.declareShidVal(
+                freeVariable.name(),
+                freeVariableType(freeVariable, bodyContextParent)
+            );
+            var freeVariableLocal = freeVariableLocalResult.value();
+            bodyContext = freeVariableLocalResult.context();
+
+            bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
+                freeVariableLocal,
+                closurePointer,
+                closureLayout,
+                freeVariableOffset(freeVariable),
+                source
+            ));
+
+            bodyContext = bodyContext.initShidLocal(freeVariable.name());
+        }
+        return bodyContext;
+    }
+
     static PatioCodeGenContext call(
         PatioLocalRef target,
         PatioLocalRef closure,