[JuliaLowering] Ask Nanosoldier what he thinks

JuliaLowering/src/binding_analysis.jl

@@ -238,8 +238,7 @@ function du_visit!(ctx, state::DefUseState, e)
         return false
 
     elseif k == K"function_decl"
-        # [function_decl] defines and instantiates the closure type and assigns
-        # it to its first argument (but only once per unique closure key).
+        # [function_decl] defines and instantiates the closure type
         @assert kind(e[1]) == K"BindingId"
         func_id = syntax_id(e[1])
         func_id in state.seen && return false
@@ -251,17 +250,25 @@ function du_visit!(ctx, state::DefUseState, e)
                 end
             end
         end
-        du_assign!(state, func_id)
         return false
 
     elseif k == K"method_defs"
-        # Process nested lambdas within
+        # XXX: the assignment is executed after the body, but flisp also makes
+        # the mistake of modelling the assignment as dominating the body, so we
+        # introduce boxes if it's corrected.
+        if kind(e[1]) === K"BindingId"
+            du_assign!(state, syntax_id(e[1]))
+        end
         has_label = false
         for child in children(e)
             has_label |= du_visit!(ctx, state, child)
         end
         return has_label
 
+    elseif k == K"no_method_defs"
+        du_assign!(state, syntax_id(e[1]))
+        return false
+
     elseif k == K"return"
         has_label = numchildren(e) >= 1 ? du_visit!(ctx, state, e[1]) : false
         du_kill!(state) # not necessary, but included for flisp parity

JuliaLowering/src/closure_conversion.jl

@@ -31,6 +31,7 @@ mutable struct ClosureConversionCtx <: AbstractLoweringContext
     const toplevel_pure::Bool
     const toplevel_stmts::Vector{SyntaxTree}
     const closure_infos::Dict{ClosureKey,ClosureInfo}
+    const closure_structs::Dict{ClosureKey,SyntaxTree}
 end
 
 function current_lambda_bindings(ctx::ClosureConversionCtx)
@@ -359,6 +360,7 @@ function convert_local_function_decl(ctx, ex)
                   field_val])
         end
     end
+    ctx.closure_structs[ck] = clstruct = ssavar(ctx, ex[1])
     @ast ctx ex [K"block"
         define_clstruct
         (::K"latestworld_if_toplevel")
@@ -367,30 +369,44 @@ function convert_local_function_decl(ctx, ex)
         else
             [K"call" "apply_type"::K"core" global_clstruct type_params...]
         end
-        closure_val := [K"new" closure_type init_closure_args...]
-        convert_assignment(ctx, [K"=" ex[1] closure_val])
+        [K"=" clstruct [K"new" closure_type init_closure_args...]]
         (::K"TOMBSTONE")
     ]
 end
 
-# Map the children of `ex` through _convert_closures, lifting any toplevel
-# closure definition statements to occur before the other content of `ex`.
+# We want to change the order of children as little as necessary to get all
+# top-level-only forms out to top level (extra movement is hard to reason about,
+# as there is currently a somewhat brittle ordering of forms enforced by
+# desugaring).  For top-level `st`, this means setting up a new `toplevel_stmts`
+# catcher for all children of `st` to add to.  Otherwise, expressions use their
+# parent's catcher.  An exception to "as little as necessary" is made for loops
+# for performance reasons.
 function map_cl_convert(ctx::ClosureConversionCtx, ex)
-    if ctx.toplevel
+    if !ctx.toplevel
+        mapchildren(e->_convert_closures(ctx, e), ex)
+    elseif kind(ex) === K"_while" || kind(ex) === K"_do_while"
+        mapchildren(e->_convert_closures(
+            ClosureConversionCtx(
+                ctx.bindings, ctx.mod,
+                ctx.closure_bindings, ctx.capture_rewriting, ctx.top_bindings,
+                ctx.lambda_bindings, ctx.sp_typevars, false, ctx.lifted,
+                ctx.toplevel_pure, ctx.toplevel_stmts, ctx.closure_infos,
+                ctx.closure_structs),
+            e), ex)
+    else
         toplevel_stmts = SyntaxList()
         ctx2 = ClosureConversionCtx(
             ctx.bindings, ctx.mod,
             ctx.closure_bindings, ctx.capture_rewriting, ctx.top_bindings,
             ctx.lambda_bindings, ctx.sp_typevars, true, ctx.lifted,
-            ctx.toplevel_pure, toplevel_stmts, ctx.closure_infos)
+            ctx.toplevel_pure, toplevel_stmts, ctx.closure_infos,
+            ctx.closure_structs)
         res = mapchildren(e->_convert_closures(ctx2, e), ex)
         if isempty(toplevel_stmts)
             res
         else
             @ast ctx ex [K"block" toplevel_stmts... res]
         end
-    else
-        mapchildren(e->_convert_closures(ctx, e), ex)
     end
 end
 
@@ -471,12 +487,7 @@ function _convert_closures(ctx::ClosureConversionCtx, ex)
         if haskey(ctx.closure_bindings, closure_key(ctx, func_name))
             convert_local_function_decl(ctx, ex)
         else
-            # Single-arg K"method" has the side effect of creating a global
-            # binding for `func_name` if it doesn't exist.
-            @ast ctx ex [K"block"
-                [K"method" func_name]
-                (::K"TOMBSTONE") # <- function_decl should not be used in value position
-            ]
+            @ast ctx ex [K"block" [K"method" func_name] (::K"TOMBSTONE")]
         end
     elseif k == K"method"
         @jl_assert ctx.lifted ex
@@ -536,22 +547,25 @@ function _convert_closures(ctx::ClosureConversionCtx, ex)
             ctx.closure_bindings, cap_rewrite,
             ctx.top_bindings, ctx.lambda_bindings, ctx.sp_typevars,
             ctx.toplevel, true, ctx.toplevel_pure, ctx.toplevel_stmts,
-            ctx.closure_infos)
+            ctx.closure_infos, ctx.closure_structs)
         tvs = map_cl_convert(ctx2, ex[2])
-        if !ctx.toplevel
+        assign_fname = !is_closure ? nothing : let ck = closure_key(ctx, name)
+            convert_assignment(ctx, @ast ctx ex [K"=" name ctx.closure_structs[ck]])
+        end
+        if is_closure && !ctx.toplevel
             push!(ctx2.toplevel_stmts, tvs)
-            tvs = @ast ctx ex[2] (::K"TOMBSTONE")
+            push!(ctx2.toplevel_stmts, map_cl_convert(ctx2, ex[3]))
+            @ast ctx ex [K"block" assign_fname (::K"TOMBSTONE")]
+        else
+            @ast ctx ex [K"block" tvs map_cl_convert(ctx2, ex[3]) assign_fname]
         end
-        body = map_cl_convert(ctx2, ex[3])
-        if is_closure
-            if ctx.toplevel
-                @ast ctx ex [K"block" tvs body]
-            else
-                push!(ctx2.toplevel_stmts, body)
-                @ast ctx ex (::K"TOMBSTONE")
-            end
+    elseif k == K"no_method_defs"
+        name = ex[1]
+        if kind(name) == K"BindingId" && get_binding(ctx, name).kind === :local
+            ck = closure_key(ctx, name)
+            convert_assignment(ctx, @ast ctx ex [K"=" name ctx.closure_structs[ck]])
         else
-            @ast ctx ex [K"block" tvs body (::K"TOMBSTONE")]
+            @ast ctx ex (::K"TOMBSTONE")
         end
     elseif k == K"_opaque_closure"
         ck = closure_key(ctx, ex[1])
@@ -565,8 +579,8 @@ function _convert_closures(ctx::ClosureConversionCtx, ex)
             ctx.bindings, ctx.mod,
             ctx.closure_bindings, capture_rewrites, ctx.top_bindings,
             ctx.lambda_bindings, ctx.sp_typevars, false, false,
-            ctx.toplevel_pure, ctx.toplevel_stmts, ctx.closure_infos)
-
+            ctx.toplevel_pure, ctx.toplevel_stmts, ctx.closure_infos,
+            ctx.closure_structs)
         argt = _convert_closures(ctx, ex[2])
         rt_lb = _convert_closures(ctx, ex[3])
         rt_ub = _convert_closures(ctx, ex[4])
@@ -616,7 +630,7 @@ function closure_convert_lambda(ctx, ex, sps)
         lbs, ctx.sp_typevars,
         k === K"toplevel_lambda", k === K"toplevel_lambda",
         ctx.toplevel_pure && k == K"generated_lambda",
-        ctx.toplevel_stmts, ctx.closure_infos)
+        ctx.toplevel_stmts, ctx.closure_infos, ctx.closure_structs)
     lambda_children = SyntaxList()
     push!(lambda_children, ex[1])
     push!(lambda_children, ex[2])
@@ -663,12 +677,15 @@ end
 
 
 """
-Closure conversion and lowering of bindings
+For each local function decl with closure key `ck`, we:
+1. Declare the closure type, populating `closure_infos[ck]`
+2. Define all methods
+3. Instantiate the closure with `new`, storing it in `closure_structs[ck]`, and
+   assigning this to the function name
 
-This pass does a few things:
+Also in this pass:
 * Deal with typed variables (K"decl") and their assignments
 * Deal with const and non-const global assignments
-* Convert closures into types
 * Lower variables captured by closures into boxes, etc, as necessary
 
 Invariants:
@@ -685,7 +702,8 @@ Invariants:
                                    ctx.closure_bindings, nothing,
                                    lbs, lbs, ctx.sp_typevars,
                                    false, true, true, SyntaxList(),
-                                   Dict{ClosureKey,ClosureInfo}())
+                                   Dict{ClosureKey,ClosureInfo}(),
+                                   Dict{ClosureKey,SyntaxTree}())
     ex_out = closure_convert_lambda(ctx_out, ex, children(ex[3]))
     if !isempty(ctx_out.toplevel_stmts)
         throw(LoweringError(first(ctx_out.toplevel_stmts), "Top level code was found outside any top level context. `@generated` functions may not contain closures, including `do` syntax and generators/comprehension"))

JuliaLowering/src/compat.jl

@@ -82,18 +82,6 @@ function _expr_to_est(@nospecialize(e), src::SourceAttrType)
     elseif e isa QuoteNode
         cid, _ = _expr_to_est(e.value, src)
         newnode(src, K"inert", SyntaxList(cid))
-    elseif e isa Expr && e.head === :lambda && length(e.args) == 2
-        argnames = e.args[1]::Vector
-        arg_cs = SyntaxTree[]
-        for name in argnames
-            id = newleaf(src, K"Identifier", String(name::Symbol))
-            push!(arg_cs, id)
-        end
-        body_id, src = _expr_to_est(e.args[2], src)
-        args_block = newnode(src, K"block", arg_cs)
-        tvars_block = newnode(src, K"block", SyntaxTree[])
-        st = newnode(src, K"lambda",
-                     SyntaxTree[args_block, tvars_block, body_id])
     elseif e isa Expr
         head_s = string(e.head)
         st_k = find_kind(head_s)
@@ -504,7 +492,7 @@ function collect_body_meta(st)
                 km = kind(m)
                 if km === K"purity"
                     push!(mmetas, m)
-                elseif syntax_name(m) in (
+                elseif kind(m) === K"Identifier" && syntax_name(m) in (
                     "inline", "noinline", "propagate_inbounds",
                     "nospecializeinfer", "aggressive_constprop", "no_constprop")
                     push!(mmetas, @mknode(m; kind=K"Symbol"))
@@ -527,6 +515,12 @@ function _dst_function_body(ctx, st, r, method_metas)
     isnothing(method_metas) ? r2 : setmeta(r2, :method_metas, method_metas)
 end
 
+function dst_raw_lambda(ctx, st, sps)
+    argl = map(x->_expr_to_est(x::Symbol, st[1])[1], st[1].value::Vector)
+    @ast _ st [K"lambda" [K"block" argl...] [K"block" sps...]
+        est_to_dst(with(ctx; toplevel=false), st[2])]
+end
+
 """
 Convert the Expr-like tree (EST) coming from macro expansion to the tree
 desugaring expects (DST), where some forms have SyntaxNode structure and others
@@ -640,6 +634,14 @@ function est_to_dst(ctx::SyntaxCompatContext, st::SyntaxTree)
                 @ast _ st [K"call" "collect"::K"top" arg]
             end
         end
+        [K"typed_comprehension" t g] -> let
+            arg = rec(ctx, g)
+            if kind(arg) === K"generator"
+                @ast _ st [K"typed_comprehension" t arg]
+            else
+                @ast _ st [K"call" "collect"::K"top" t arg]
+            end
+        end
         # hack: `[_ for _ in rhs]`, `[f(_) for _ in rhs]` works
         ([K"generator" body [K"=" u2 rhs]],
          when=is_flisp_compat(st) &&
@@ -741,8 +743,9 @@ function est_to_dst(ctx::SyntaxCompatContext, st::SyntaxTree)
         [K"core" x] -> newleaf(st, K"core", syntax_name(x))
         [K"top" x] -> newleaf(st, K"top", syntax_name(x))
         [K"static_parameter" x] -> newleaf(st, K"static_parameter", x.value::IdTag)
-        [K"lambda" args sps body] -> @mknode(st; children=SyntaxTree[
-            args, sps, rec(with(ctx; toplevel=false), body)])
+        [K"with-static-parameters" lam sps...] ->
+            dst_raw_lambda(ctx, lam, sps)
+        [K"lambda" _ _] -> dst_raw_lambda(ctx, st, SyntaxTree[])
         [K"copyast" [K"inert" ex]] -> @ast _ st [K"call"
             interpolate_expr::K"Value"
             [K"inert"(st[1]) ex]