Conversion of assignments to closure captures

Convert captured variables to `Box`es for both read and write, both
inside and outside the closure. Also ensure we emit newvar nodes for
variables which might be used uninitialized (currently pessimistic).

Author: c42f

src/ast.jl

@@ -63,6 +63,7 @@ struct BindingInfo
     type::Union{Nothing,SyntaxTree} # Type, for bindings declared like x::T = 10
     is_const::Bool            # Constant, cannot be reassigned
     is_ssa::Bool              # Single assignment, defined before use
+    is_captured::Bool         # Variable is captured by some lambda
     is_always_defined::Bool   # A local that we know has an assignment that dominates all usages (is never undef)
     is_internal::Bool         # True for internal bindings generated by the compiler
     is_ambiguous_local::Bool  # Local, but would be global in soft scope (ie, the REPL)
@@ -74,11 +75,12 @@ function BindingInfo(name::AbstractString, kind::Symbol;
                      type::Union{Nothing,SyntaxTree} = nothing,
                      is_const::Bool = false,
                      is_ssa::Bool = false,
+                     is_captured::Bool = false,
                      is_always_defined::Bool = is_ssa,
                      is_internal::Bool = false,
                      is_ambiguous_local::Bool = false,
                      is_nospecialize::Bool = false)
-    BindingInfo(name, kind, mod, type, is_const, is_ssa, is_always_defined,
+    BindingInfo(name, kind, mod, type, is_const, is_ssa, is_captured, is_always_defined,
                 is_internal, is_ambiguous_local, is_nospecialize)
 end
 
@@ -113,7 +115,7 @@ function _binding_id(ex::SyntaxTree)
 end
 
 function update_binding!(bindings::Bindings, x;
-        type=nothing, is_const=nothing, is_always_defined=nothing)
+        type=nothing, is_const=nothing, is_always_defined=nothing, is_captured=nothing)
     id = _binding_id(x)
     b = lookup_binding(bindings, id)
     bindings.info[id] = BindingInfo(
@@ -123,6 +125,7 @@ function update_binding!(bindings::Bindings, x;
         isnothing(type) ? b.type : type,
         isnothing(is_const) ? b.is_const : is_const,
         b.is_ssa,
+        isnothing(is_captured) ? b.captured : is_captured,
         isnothing(is_always_defined) ? b.is_always_defined : is_always_defined,
         b.is_internal,
         b.is_ambiguous_local,
@@ -224,6 +227,8 @@ function makeleaf(ctx, srcref, k::Kind, value; kws...)
         makeleaf(graph, srcref, k; id=value, kws...)
     elseif k == K"symbolic_label"
         makeleaf(graph, srcref, k; name_val=value, kws...)
+    elseif k == K"TOMBSTONE"
+        makeleaf(graph, srcref, k; kws...)
     else
         val = k == K"Integer" ? convert(Int,     value) :
               k == K"Float"   ? convert(Float64, value) :
@@ -342,8 +347,15 @@ end
 function _expand_ast_tree(ctx, srcref, tree)
     if Meta.isexpr(tree, :(::))
         # Leaf node
-        _match_kind(srcref, tree.args[2]) do kind, srcref, kws
-            :(makeleaf($ctx, $srcref, $kind, $(esc(tree.args[1])), $(kws...)))
+        if length(tree.args) == 2
+            val = esc(tree.args[1])
+            kindspec = tree.args[2]
+        else
+            val = nothing
+            kindspec = tree.args[1]
+        end
+        _match_kind(srcref, kindspec) do kind, srcref, kws
+            :(makeleaf($ctx, $srcref, $kind, $(val), $(kws...)))
         end
     elseif Meta.isexpr(tree, :call) && tree.args[1] === :(=>)
         # Leaf node with copied attributes

src/closure_conversion.jl

@@ -31,6 +31,49 @@ function add_lambda_local!(ctx::ClosureConversionCtx, id)
     init_lambda_binding(ctx.lambda_bindings, id)
 end
 
+# Access captured variable from inside a closure
+function captured_var_access(ctx, ex)
+    cinfo = ctx.closure_info
+    field_sym = cinfo.field_syms[cinfo.field_name_inds[ex.var_id]]
+    @ast ctx ex [K"call"
+        "getfield"::K"core"
+        # FIXME: attributing the self binding to srcref=ex gives misleading printing.
+        # We should carry provenance with each binding to fix this.
+        ctx.lambda_bindings.self::K"BindingId"
+        field_sym
+    ]
+end
+
+function get_box_contents(ctx::ClosureConversionCtx, var, box_ex)
+    undef_var = new_mutable_var(ctx, var, lookup_binding(ctx, var.var_id).name)
+    @ast ctx var [K"block"
+        box := box_ex
+        # Lower in an UndefVar check to a similarly named variable
+        # (ref #20016) so that closure lowering Box introduction
+        # doesn't impact the error message and the compiler is expected
+        # to fold away the extraneous null check
+        #
+        # TODO: Ideally the runtime would rely on provenance info for
+        # this error and we can remove isdefined check.
+        [K"if" [K"call"
+                "isdefined"::K"core"
+                box
+                "contents"::K"Symbol"
+            ]
+            ::K"TOMBSTONE"
+            [K"block"
+                 [K"newvar" undef_var]
+                 undef_var
+            ]
+        ]
+        [K"call"
+            "getfield"::K"core"
+            box
+            "contents"::K"Symbol"
+        ]
+    ]
+end
+
 # Convert `ex` to `type` by calling `convert(type, ex)` when necessary.
 #
 # Used for converting the right hand side of an assignment to a typed local or
@@ -118,37 +161,32 @@ function convert_assignment(ctx, ex)
     if binfo.kind == :global
         convert_global_assignment(ctx, ex, var, rhs0)
     else
-        closed   = false # TODO
-        captured = false # TODO
-        @assert binfo.kind == :local
-        if isnothing(binfo.type) && !closed && !captured
+        @assert binfo.kind == :local || binfo.kind == :argument
+        lbinfo = get(ctx.lambda_bindings.bindings, var.var_id, nothing)
+        self_captured = !isnothing(lbinfo) && lbinfo.is_captured
+        captured      = binfo.is_captured
+        if isnothing(binfo.type) && !self_captured && !captured
             @ast ctx ex [K"=" var rhs0]
         else
-            @assert binfo.kind == :local
             # Typed local
-            tmp_rhs0 = is_simple_atom(ctx, rhs0) ? nothing : ssavar(ctx, rhs0)
-            rhs1 = isnothing(tmp_rhs0) ? rhs0 : tmp_rhs0
-            rhs = isnothing(binfo.type) ? rhs1 :
-                  convert_for_type_decl(ctx, ex, rhs1, _convert_closures(ctx, binfo.type), true)
-            assgn = if closed
-                @assert false # TODO
-            elseif captured
-                @assert false # TODO
-            else
-                @ast ctx ex [K"=" var rhs]
-            end
-            if isnothing(tmp_rhs0)
-                @ast ctx ex [K"block"
-                    assgn
-                    rhs0
+            tmp_rhs0 = ssavar(ctx, rhs0)
+            rhs = isnothing(binfo.type) ? tmp_rhs0 :
+                  convert_for_type_decl(ctx, ex, tmp_rhs0, _convert_closures(ctx, binfo.type), true)
+            assignment = if self_captured || captured
+                @ast ctx ex [K"call"
+                    "setfield!"::K"core"
+                    self_captured ? captured_var_access(ctx, var) : var
+                    "contents"::K"Symbol"
+                    rhs
                 ]
             else
-                @ast ctx ex [K"block"
-                    [K"=" tmp_rhs0 rhs0]
-                    assgn
-                    tmp_rhs0
-                ]
+                @ast ctx ex [K"=" var rhs]
             end
+            @ast ctx ex [K"block"
+                [K"=" tmp_rhs0 rhs0]
+                assignment
+                tmp_rhs0
+            ]
         end
     end
 end
@@ -237,40 +275,10 @@ function _convert_closures(ctx::ClosureConversionCtx, ex)
     if k == K"BindingId"
         id = ex.var_id
         lbinfo = get(ctx.lambda_bindings.bindings, id, nothing)
-        if !isnothing(lbinfo) && lbinfo.is_captured
-            cinfo = ctx.closure_info
-            field_sym = cinfo.field_syms[cinfo.field_name_inds[id]]
-            undef_var = new_mutable_var(ctx, ex, lookup_binding(ctx, id).name)
-            @ast ctx ex [K"block"
-                box := [K"call"
-                    "getfield"::K"core"
-                    ctx.lambda_bindings.self::K"BindingId"
-                    field_sym
-                ]
-                # Lower in an UndefVar check to a similarly named variable
-                # (ref #20016) so that closure lowering Box introduction
-                # doesn't impact the error message and the compiler is expected
-                # to fold away the extraneous null check
-                #
-                # TODO: Ideally the runtime would rely on provenance info for
-                # this error and we can remove isdefined check.
-                [K"if" [K"call"
-                        "isdefined"::K"core"
-                        box
-                        "contents"::K"Symbol"
-                    ]
-                    "nothing"::K"core"
-                    [K"block"
-                         [K"newvar" undef_var]
-                         undef_var
-                    ]
-                ]
-                [K"call"
-                    "getfield"::K"core"
-                    box
-                    "contents"::K"Symbol"
-                ]
-            ]
+        if !isnothing(lbinfo) && lbinfo.is_captured # TODO: && vinfo:asgn cv ??
+            get_box_contents(ctx, ex, captured_var_access(ctx, ex))
+        elseif lookup_binding(ctx, id).is_captured # TODO: && vinfo:asgn vi
+            get_box_contents(ctx, ex, ex)
         else
             ex
         end
@@ -295,6 +303,16 @@ function _convert_closures(ctx::ClosureConversionCtx, ex)
                 _convert_closures(ctx, ex[2])
             ]
         end
+    elseif k == K"local"
+        var = ex[1]
+        binfo = lookup_binding(ctx, var)
+        if binfo.is_captured
+            @ast ctx ex [K"=" var [K"call" "Box"::K"core"]]
+        elseif !binfo.is_always_defined
+            @ast ctx ex [K"newvar" var]
+        else
+            makeleaf(ctx, ex, K"TOMBSTONE")
+        end
     elseif k == K"::"
         _convert_closures(ctx,
             @ast ctx ex [K"call"
@@ -322,10 +340,7 @@ function _convert_closures(ctx::ClosureConversionCtx, ex)
                 ctx.closure_infos[func_name_id] = closure_info
                 init_closure_args = SyntaxList(ctx)
                 for id in field_orig_bindings
-                    # FIXME: This isn't actually correct: we need to convert
-                    # all outer references to boxes too!
-                    push!(init_closure_args, _convert_closures(ctx,
-                          @ast ctx ex [K"call" "Box"::K"core" id::K"BindingId"]))
+                    push!(init_closure_args, @ast ctx ex id::K"BindingId")
                 end
                 @ast ctx ex [K"block"
                     [K"=" func_name
@@ -380,18 +395,30 @@ function closure_convert_lambda(ctx, ex)
     @assert kind(ex) == K"lambda"
     body_stmts = SyntaxList(ctx)
     toplevel_stmts = ex.is_toplevel_thunk ? body_stmts : ctx.toplevel_stmts
+    lambda_bindings = ex.lambda_bindings
     ctx2 = ClosureConversionCtx(ctx.graph, ctx.bindings, ctx.mod,
-                                ctx.closure_bindings, ctx.closure_info, ex.lambda_bindings,
+                                ctx.closure_bindings, ctx.closure_info, lambda_bindings,
                                 toplevel_stmts, ctx.closure_infos)
     lambda_children = SyntaxList(ctx)
-    push!(lambda_children, _convert_closures(ctx2, ex[1]))
-    push!(lambda_children, _convert_closures(ctx2, ex[2]))
+    args = ex[1]
+    push!(lambda_children, args)
+    push!(lambda_children, ex[2])
 
-    # Convert body. This is done as a special case to allow inner calls to
-    # _convert_closures to also add to body_stmts in the case that
-    # ex.is_toplevel_thunk is true.
-    in_body_stmts = kind(ex[3]) != K"block" ? ex[3:3] : ex[3][1:end]
-    for e in in_body_stmts
+    # Add box initializations for arguments which are captured by an inner lambda
+    for arg in children(args)
+        kind(arg) != K"Placeholder" || continue
+        binfo = lookup_binding(ctx, arg)
+        if binfo.is_captured # TODO: && binfo.is_assigned
+            push!(body_stmts, @ast ctx arg [K"="
+                arg
+                [K"call" "Box"::K"core" arg]
+            ])
+        end
+    end
+    # Convert body. Note that _convert_closures may call `push!(body_stmts, e)`
+    # internally for any expressions `e` which need to be moved to top level.
+    input_body_stmts = kind(ex[3]) != K"block" ? ex[3:3] : ex[3][1:end]
+    for e in input_body_stmts
         push!(body_stmts, _convert_closures(ctx2, e))
     end
     push!(lambda_children, @ast ctx2 ex[3] [K"block" body_stmts...])

src/scope_analysis.jl

@@ -64,9 +64,9 @@ function _find_scope_vars!(assignments, locals, destructured_args, globals, used
 end
 
 # Find names of all identifiers used in the given expression, grouping them
-# into sets by type.
+# into sets by type of usage.
 #
-# NB: This only works propery after desugaring has already processed assignments
+# NB: This only works propery after desugaring
 function find_scope_vars(ex)
     ExT = typeof(ex)
     assignments = Dict{NameKey,ExT}()
@@ -422,6 +422,7 @@ function analyze_scope(ctx, ex, scope_type, is_toplevel_global_scope=false,
                 if !haskey(lambda_bindings.bindings, id)
                     # Used vars from a scope *outside* the current lambda are captured
                     init_lambda_binding(lambda_bindings, id, is_captured=true, is_read=true)
+                    update_binding!(ctx, id; is_captured=true)
                 else
                     update_lambda_binding!(lambda_bindings, id, is_read=true)
                 end
@@ -437,6 +438,7 @@ function analyze_scope(ctx, ex, scope_type, is_toplevel_global_scope=false,
                 if !haskey(lambda_bindings.bindings, id)
                     # Assigned vars from a scope *outside* the current lambda are captured
                     init_lambda_binding(lambda_bindings, id, is_captured=true, is_assigned=true)
+                    update_binding!(ctx, id; is_captured=true)
                 else
                     update_lambda_binding!(lambda_bindings, id, is_assigned=true)
                 end
@@ -457,6 +459,17 @@ function analyze_scope(ctx, ex, scope_type, is_toplevel_global_scope=false,
                      is_hard_scope, var_ids, lambda_bindings)
 end
 
+function add_local_decls!(ctx, stmts, srcref, scope)
+    # Add local decls to start of block so that closure conversion can
+    # initialize if necessary.
+    for id in values(scope.var_ids)
+        binfo = lookup_binding(ctx, id)
+        if binfo.kind == :local
+            push!(stmts, @ast ctx srcref [K"local" id::K"BindingId"])
+        end
+    end
+end
+
 # Do some things which are better done after converting to BindingId.
 function maybe_update_bindings!(ctx, ex)
     k = kind(ex)
@@ -522,12 +535,20 @@ function _resolve_scopes(ctx, ex::SyntaxTree)
         push!(ctx.scope_stack, scope)
         arg_bindings = _resolve_scopes(ctx, ex[1])
         sparm_bindings = _resolve_scopes(ctx, ex[2])
+        body_stmts = SyntaxList(ctx)
+        add_local_decls!(ctx, body_stmts, ex, scope)
         body = _resolve_scopes(ctx, ex[3])
+        if kind(body) == K"block"
+            append!(body_stmts, children(body))
+        else
+            push!(body_stmts, body)
+        end
         ret_var = numchildren(ex) == 4 ? _resolve_scopes(ctx, ex[4]) : nothing
         pop!(ctx.scope_stack)
 
         lambda_bindings = scope.lambda_bindings
         if !is_toplevel_thunk
+            # Record all lambdas for the same closure type in one place
             func_name = last(ctx.method_def_stack)
             if kind(func_name) == K"BindingId"
                 func_name_id = func_name.var_id
@@ -544,19 +565,21 @@ function _resolve_scopes(ctx, ex::SyntaxTree)
                                is_toplevel_thunk=is_toplevel_thunk)
             arg_bindings
             sparm_bindings
-            body
+            [K"block"
+                body_stmts...
+            ]
             ret_var
         ]
     elseif k == K"scope_block"
         scope = analyze_scope(ctx, ex, ex.scope_type)
         push!(ctx.scope_stack, scope)
-        body = SyntaxList(ctx)
+        stmts = SyntaxList(ctx)
+        add_local_decls!(ctx, stmts, ex, scope)
         for e in children(ex)
-            push!(body, _resolve_scopes(ctx, e))
+            push!(stmts, _resolve_scopes(ctx, e))
         end
-        body
         pop!(ctx.scope_stack)
-        @ast ctx ex [K"block" body...]
+        @ast ctx ex [K"block" stmts...]
     elseif k == K"extension"
         etype = extension_type(ex)
         if etype == "islocal"

src/utils.jl

@@ -78,6 +78,9 @@ function print_ir(io::IO, ex, indent="")
             println(io, indent, lno, " --- method ", string(e[1]), " ", string(e[2]))
             @assert kind(e[3]) == K"lambda" || kind(e[3]) == K"code_info"
             print_ir(io, e[3], indent*"    ")
+        elseif kind(e) == K"code_info" && e.is_toplevel_thunk
+            println(io, indent, lno, " --- thunk")
+            print_ir(io, e, indent*"    ")
         else
             code = string(e)
             println(io, indent, lno, " ", code)