Preserve captured environments in nested AOT functions

Author: ingydotnet

pkg/runtime/codegen.go

@@ -1159,7 +1159,12 @@ func (g *Generator) generateFnFunc(fn lang.FnFunc) string {
 func (g *Generator) generateFn(fn *Fn) string {
 	// Save and restore current environment
 	prevEnv := g.currentFnEnv
-	g.currentFnEnv = fn.GetEnvironment() // Set the captured environment for this function
+	// Runtime function values carry their captured environment. Functions
+	// constructed directly from nested AST nodes do not; in that case they
+	// inherit the environment of the function currently being generated.
+	if env := fn.GetEnvironment(); env != nil {
+		g.currentFnEnv = env
+	}
 	defer func() { g.currentFnEnv = prevEnv }()
 
 	astNode := fn.ASTNode()

pkg/runtime/codegen_internal_test.go

@@ -41,6 +41,25 @@ func TestGenerateStandardFileHandles(t *testing.T) {
 	}
 }
 
+func TestGenerateNestedClosureCapturedAtLoadTime(t *testing.T) {
+	ns := lang.FindOrCreateNamespace(lang.NewSymbol("codegen.nested-capture"))
+	ns.ReferAllSnapshot(lang.NSCore, nil)
+	lang.PushThreadBindings(lang.NewMap(lang.VarCurrentNS, ns))
+	defer lang.PopThreadBindings()
+
+	ReadEval(`
+		(defn make-closure [keep]
+		  (fn [_]
+		    (let [finish (fn [] keep)]
+		      (finish))))
+		(def captured (make-closure true))`)
+
+	var output bytes.Buffer
+	if err := NewGenerator(&output).Generate(ns); err != nil {
+		t.Fatalf("generate nested captured closure: %v", err)
+	}
+}
+
 func TestGenerateResolvedHostReference(t *testing.T) {
 	generator := NewGenerator(&bytes.Buffer{})
 	node := ast.MakeNode(ast.OpConst, nil)