Ensure outer closures capture variables nested inside inner closures

Author: c42f

src/closure_conversion.jl

@@ -391,8 +391,13 @@ function _convert_closures(ctx::ClosureConversionCtx, ex)
                 ctx.closure_infos[func_name_id] = closure_info
                 type_params = SyntaxList(ctx)
                 init_closure_args = SyntaxList(ctx)
-                for (id,boxed) in zip(field_orig_bindings, field_is_box)
+                for (id, boxed) in zip(field_orig_bindings, field_is_box)
                     field_val = binding_ex(ctx, id)
+                    if is_self_captured(ctx, field_val)
+                        # Access from outer closure if necessary but do not
+                        # unbox to feed into the inner nested closure.
+                        field_val = captured_var_access(ctx, field_val)
+                    end
                     push!(init_closure_args, field_val)
                     if !boxed
                         push!(type_params, @ast ctx ex [K"call"

src/scope_analysis.jl

@@ -710,6 +710,20 @@ function analyze_variables!(ctx, ex)
         end
         ctx2 = VariableAnalysisContext(ctx.graph, ctx.bindings, ctx.mod, lambda_bindings,
                                        ctx.method_def_stack, ctx.closure_bindings)
+        # Add any captured bindings to the enclosing lambda, if necessary.
+        for (id,lbinfo) in pairs(lambda_bindings.bindings)
+            if lbinfo.is_captured
+                outer_lbinfo = lookup_lambda_binding(ctx.lambda_bindings, id)
+                if isnothing(outer_lbinfo)
+                    # Inner lambda captures a variable. If it's not yet present
+                    # in the outer lambda, the outer lambda must capture it as
+                    # well so that the closure associated to the inner lambda
+                    # can be initialized when `function_decl` is hit.
+                    init_lambda_binding(ctx.lambda_bindings, id, is_captured=true, is_read=true)
+                end
+            end
+        end
+
         # TODO: Types of any assigned captured vars will also be used and might be captured.
         foreach(e->analyze_variables!(ctx2, e), ex[3:end])
     else

test/closures.jl

@@ -110,6 +110,22 @@ end
 @test f_closure_local_var_types(2.0) == 1.0
 @test_throws MethodError f_closure_local_var_types("hi")
 
+# Multiply nested closures. In this case g_nest needs to capture `x` in order
+# to construct an instance of `h_nest()` inside it.
+@test JuliaLowering.include_string(test_mod, """
+begin
+    function f_nest(x)
+        function g_nest(y)
+            function h_nest(z)
+                (x,y,z)
+            end
+        end
+    end
+
+    f_nest(1)(2)(3)
+end
+""") === (1,2,3)
+
 # Anon function syntax
 @test JuliaLowering.include_string(test_mod, """
 begin

test/closures_ir.jl

@@ -285,7 +285,8 @@ end
 19  (return %₁₈)
 
 ########################################
-# FIXME: Nested captures of arguments
+# Nested captures - here `g` captures `x` because it is needed to initialize
+# the closure `h` which captures both `x` and `y`.
 function f(x)
     function g(y)
         function h(z)
@@ -294,18 +295,58 @@ function f(x)
     end
 end
 #---------------------
-LoweringError:
-function f(x)
-#          ╙ ── Found unexpected binding of kind argument
-    function g(y)
-        function h(z)
-
-Detailed provenance:
-#₈/x
-└─ x
-   └─ x
-      └─ @ :1
-
+1   (method TestMod.f)
+2   (call core.svec :x)
+3   (call core.svec false)
+4   (call JuliaLowering.eval_closure_type TestMod :#f#g##4 %₂ %₃)
+5   (call core.svec :x :y)
+6   (call core.svec false false)
+7   (call JuliaLowering.eval_closure_type TestMod :#f#g#h##0 %₅ %₆)
+8   TestMod.#f#g#h##0
+9   (call core.svec %₈ core.Any)
+10  (call core.svec)
+11  SourceLocation::3:18
+12  (call core.svec %₉ %₁₀ %₁₁)
+13  --- method core.nothing %₁₂
+    slots: [slot₁/#self#(!read) slot₂/z]
+    1   (call core.getfield slot₁/#self# :x)
+    2   (call core.getfield slot₁/#self# :y)
+    3   (call core.tuple %₁ %₂ slot₂/z)
+    4   (return %₃)
+14  TestMod.#f#g##4
+15  (call core.svec %₁₄ core.Any)
+16  (call core.svec)
+17  SourceLocation::2:14
+18  (call core.svec %₁₅ %₁₆ %₁₇)
+19  --- method core.nothing %₁₈
+    slots: [slot₁/#self#(!read) slot₂/y(!read) slot₃/h]
+    1   TestMod.#f#g#h##0
+    2   (call core.getfield slot₁/#self# :x)
+    3   (call core.typeof %₂)
+    4   (call core.typeof slot₂/y)
+    5   (call core.apply_type %₁ %₃ %₄)
+    6   (call core.getfield slot₁/#self# :x)
+    7   (new %₅ %₆ slot₂/y)
+    8   (= slot₃/h %₇)
+    9   slot₃/h
+    10  (return %₉)
+20  TestMod.f
+21  (call core.Typeof %₂₀)
+22  (call core.svec %₂₁ core.Any)
+23  (call core.svec)
+24  SourceLocation::1:10
+25  (call core.svec %₂₂ %₂₃ %₂₄)
+26  --- method core.nothing %₂₅
+    slots: [slot₁/#self#(!read) slot₂/x(!read) slot₃/g]
+    1   TestMod.#f#g##4
+    2   (call core.typeof slot₂/x)
+    3   (call core.apply_type %₁ %₂)
+    4   (new %₃ slot₂/x)
+    5   (= slot₃/g %₄)
+    6   slot₃/g
+    7   (return %₆)
+27  TestMod.f
+28  (return %₂₇)
 
 ########################################
 # Global method capturing local variables

test/generators.jl

@@ -55,4 +55,22 @@ Tuple{Int,Int}[(x,y) for x in 1:2, y in 1:3]
 """) == [(1,1) (1,2) (1,3)
          (2,1) (2,2) (2,3)]
 
+# Triply nested comprehension
+@test JuliaLowering.include_string(test_mod, """
+[(x,y,z) for x in 1:3 for y in 4:5 for z in 6:7]
+""") == [
+    (1, 4, 6)
+    (1, 4, 7)
+    (1, 5, 6)
+    (1, 5, 7)
+    (2, 4, 6)
+    (2, 4, 7)
+    (2, 5, 6)
+    (2, 5, 7)
+    (3, 4, 6)
+    (3, 4, 7)
+    (3, 5, 6)
+    (3, 5, 7)
+]
+
 end

test/generators_ir.jl

@@ -282,18 +282,3 @@ T[(x,y) for x in xs, y in ys]
 49  (goto label₁₅)
 50  (return %₇)
 
-########################################
-# FIXME - error in nested closure conversion: Triply nested generator
-((x,y,z) for x in 1:3 for y in 4:5 for z in 6:7)
-#---------------------
-LoweringError:
-((x,y,z) for x in 1:3 for y in 4:5 for z in 6:7)
-#            ╙ ── Found unexpected binding of kind argument
-
-Detailed provenance:
-#₁₃/x
-└─ x
-   └─ x
-      └─ @ :1
-
-