Desugaring and runtime support for generated functions

I've opted to diverge from Base's lowering of `@ generated` code
generators to add a semi-hidden `__context__` argument to the lowered
code generator function in exact analogy to the context added to macros.
This makes the code generator even more macro-like, and I've also reused
the MacroExpansionContext.

We need our own `GeneratedFunctionStub` here in order to construct the
context, to call back into JuliaLowering's version of the lowering
machinery in a precise way, and to propagate enhanced provenance
information.

This all turned out to be quite subtle, but JuliaLowering is able to
almost fully integrate with the runtime without any changes to the
runtime itself. Thus we can see the effort people have put into building
abstractions for Cassette/etc has really paid off.

Author: c42f

JuliaLowering/README.md

@@ -398,6 +398,69 @@ In total, this expands a single "function definition" into seven methods.
 Note that the above is only a sketch! There's more fiddly details when `where`
 syntax comes in
 
+### Desugaring of generated functions
+
+A brief description of how this works. Let's consider the generated function
+
+```julia
+function gen(x::NTuple{N}, y) where {N,T}
+    shared = :shared
+    # Unnecessary use of @generated, but it shows what's going on.
+    if @generated
+        quote
+            maybe_gen = ($x, $N)
+        end
+    else
+        maybe_gen = (typeof(x), N)
+    end
+    (shared, maybe_gen)
+end
+```
+
+This is desugared into the following two function definitions. First, a code
+generator which will generate code for the body of the function, given the
+static parameters `N`, `T` and the positional arguments `x`, `y`.
+(`var"#self#"::Type{typeof(gen)}` is also provided by the Julia runtime to
+complete the full signature of `gen`, though the user won't normally use this.)
+
+```julia
+function var"#gen@generator#0"(__context__::JuilaSyntax.MacroContext, N, T, var"#self#", x, y)
+    gen_stuff = quote
+        maybe_gen = ($x, $N)
+    end
+    quote
+        shared = :shared
+        $gen_stuff
+        (shared, maybe_gen)
+    end
+end
+```
+
+Second, the non-generated version, using the `if @generated` else branches, and
+containing mostly normal code.
+
+```julia
+function gen(x::NTuple{N}, y) where {N,T}
+    $(Expr(:meta, :generated,
+        Expr(:call, JuliaLowering.GeneratedFunctionStub,
+             :var"#gen@generator#0", sourceref_of_gen,
+             :(Core.svec(:var"#self", :x, :y))
+             :(Core.svec(:N, :T)))))
+    shared = :shared
+    maybe_gen = (typeof(x), N)
+    (shared, maybe_gen)
+end
+```
+
+The one extra thing added here is the `Expr(:meta, :generated)` which is an
+expression creating a callable wrapper for the user's generator, to be
+evaluated at top level. This wrapper will then be invoked by the runtime
+whenever the user calls `gen` with a new signature and it's expected that a
+`CodeInfo` be returned from it. `JuliaLowering.GeneratedFunctionStub` differs
+from `Core.GeneratedFunctionStub` in that it contains extra provenance
+information (the `sourcref_of_gen`) and expects a `SyntaxTree` to be returned
+by the user's generator code.
+
 ## Pass 3: Scope analysis / binding resolution
 
 This pass replaces variables with bindings of kind `K"BindingId"`,

JuliaLowering/src/desugaring.jl

@@ -2213,7 +2213,7 @@ function method_def_expr(ctx, srcref, callex_srcref, method_table,
             ::K"SourceLocation"(callex_srcref)
         ]
         [K"method"
-            method_table
+            isnothing(method_table) ? "nothing"::K"core" : method_table
             method_metadata
             [K"lambda"(body, is_toplevel_thunk=false)
                 [K"block" arg_names...]
@@ -2285,6 +2285,117 @@ function trim_used_typevars(ctx, arg_types, typevar_names, typevar_stmts)
     return trimmed_typevar_names
 end
 
+function is_if_generated(ex)
+    kind(ex) == K"if" && kind(ex[1]) == K"generated"
+end
+
+# Return true if a function body contains a code generator from `@generated` in
+# the form `[K"if" [K"generated"] ...]`
+function is_generated(ex)
+    if is_if_generated(ex)
+        return true
+    elseif is_quoted(ex) || kind(ex) == K"function"
+        return false
+    else
+        return any(is_generated, children(ex))
+    end
+end
+
+function split_generated(ctx, ex, gen_part)
+    if is_leaf(ex)
+        ex
+    elseif is_if_generated(ex)
+        gen_part ? @ast(ctx, ex, [K"$" ex[2]]) : ex[3]
+    else
+        mapchildren(e->split_generated(ctx, e, gen_part), ctx, ex)
+    end
+end
+
+# Split @generated function body into two parts:
+# * The code generator
+# * The non-generated function body
+function expand_function_generator(ctx, srcref, callex_srcref, func_name, func_name_str, body, arg_names, typevar_names)
+    gen_body = if is_if_generated(body)
+        body[2] # Simple case - don't need interpolation when the whole body is generated
+    else
+        expand_quote(ctx, @ast ctx body [K"block" split_generated(ctx, body, true)])
+    end
+    gen_name_str = reserve_module_binding_i(ctx.mod,
+                        "#$(isnothing(func_name_str) ? "_" : func_name_str)@generator#")
+    gen_name = new_global_binding(ctx, body, gen_name_str, ctx.mod)
+
+    # Set up the arguments for the code generator
+    gen_arg_names = SyntaxList(ctx)
+    gen_arg_types = SyntaxList(ctx)
+    # Self arg
+    push!(gen_arg_names, new_local_binding(ctx, callex_srcref, "#self#"; kind=:argument))
+    push!(gen_arg_types, @ast ctx callex_srcref [K"function_type" gen_name])
+    # Macro expansion context arg
+    if kind(func_name) != K"Identifier"
+        TODO(func_name, "Which scope do we adopt for @generated generator `__context__` in this case?")
+    end
+    push!(gen_arg_names, adopt_scope(@ast(ctx, callex_srcref, "__context__"::K"Identifier"), func_name))
+    push!(gen_arg_types, @ast(ctx, callex_srcref, MacroContext::K"Value"))
+    # Trailing arguments to the generator are provided by the Julia runtime. They are:
+    # static_parameters... parent_function arg_types...
+    first_trailing_arg = length(gen_arg_names) + 1
+    append!(gen_arg_names, typevar_names)
+    append!(gen_arg_names, arg_names)
+    # Apply nospecialize to all arguments to prevent so much codegen and add
+    # Core.Any type for them
+    for i in first_trailing_arg:length(gen_arg_names)
+        gen_arg_names[i] = setmeta(gen_arg_names[i]; nospecialize=true)
+        push!(gen_arg_types, @ast ctx gen_arg_names[i] "Any"::K"core")
+    end
+    # Code generator definition
+    gen_func_method_defs = @ast ctx srcref [K"method_defs"
+        gen_name
+        method_def_expr(ctx, srcref, callex_srcref, nothing, SyntaxList(ctx), gen_arg_names,
+                        gen_arg_types, gen_body, nothing)
+    ]
+
+    # Extract non-generated body
+    nongen_body = @ast ctx body [K"block"
+        # The Julia runtime associates the code generator with the
+        # non-generated method by adding this meta to the body. This feels like
+        # a hack though since the generator ultimately gets attached to the
+        # method rather than the CodeInfo which we're putting it inside.
+        [K"meta"
+            "generated"::K"Symbol"
+            # The following is code to be evaluated at top level and will wrap
+            # whatever code comes from the user's generator into an appropriate
+            # K"lambda" (+ K"with_static_parameters") suitable for lowering
+            # into a CodeInfo.
+            #
+            # todo: As isolated top-level code, we don't actually want to apply
+            # the normal scope rules of the surrounding function ... it should
+            # technically have scope resolved at top level.
+            [K"new"
+                GeneratedFunctionStub::K"Value" # Use stub type from JuliaLowering
+                gen_name
+                # Truncate provenance to just the source file range, as this
+                # will live permanently in the IR and we probably don't want
+                # the full provenance tree and intermediate expressions
+                # (TODO: More truncation. We certainly don't want to store the
+                #  source file either.)
+                sourceref(srcref)::K"Value"
+                [K"call"
+                    "svec"::K"core"
+                    "#self#"::K"Symbol"
+                    (n.name_val::K"Symbol"(n) for n in arg_names[2:end])...
+                ]
+                [K"call"
+                    "svec"::K"core"
+                    (n.name_val::K"Symbol"(n) for n in typevar_names)...
+                ]
+            ]
+        ]
+        split_generated(ctx, body, false)
+    ]
+
+    return gen_name, gen_func_method_defs, nongen_body
+end
+
 # Generate a method for every number of allowed optional arguments
 # For example for `f(x, y=1, z=2)` we generate two additional methods
 # f(x) = f(x, 1, 2)
@@ -2799,6 +2910,14 @@ function expand_function_def(ctx, ex, docs, rewrite_call=identity, rewrite_body=
         ]
     end
 
+    gen_func_name = nothing
+    gen_func_method_defs = nothing
+    if is_generated(body)
+        gen_func_name, gen_func_method_defs, body =
+            expand_function_generator(ctx, ex, callex, name, name_str, body, arg_names, typevar_names)
+
+    end
+
     if isnothing(keywords)
         body_func_name, kw_func_method_defs = (nothing, nothing)
         # NB: This check seems good as it statically catches any useless
@@ -2848,6 +2967,9 @@ function expand_function_def(ctx, ex, docs, rewrite_call=identity, rewrite_body=
     end
 
     @ast ctx ex [K"block"
+        if !isnothing(gen_func_name)
+            [K"function_decl"(gen_func_name) gen_func_name]
+        end
         if !isnothing(body_func_name)
             [K"function_decl"(body_func_name) body_func_name]
         end
@@ -2857,6 +2979,7 @@ function expand_function_def(ctx, ex, docs, rewrite_call=identity, rewrite_body=
         [K"scope_block"(scope_type=:hard)
             [K"block"
                 new_typevar_stmts...
+                gen_func_method_defs
                 kw_func_method_defs
                 [K"method_defs"
                     isnothing(bare_func_name) ? "nothing"::K"core" : bare_func_name
@@ -3651,7 +3774,7 @@ function expand_struct_def(ctx, ex, docs)
                 typevar_in_bounds = any(type_params[i+1:end]) do param
                     # Check the bounds of subsequent type params
                     (_,lb,ub) = analyze_typevar(ctx, param)
-                    # TODO: flisp lowering tests `lb` here so we also do. But
+                    # todo: flisp lowering tests `lb` here so we also do. But
                     # in practice this doesn't seem to constrain `typevar_name`
                     # and the generated constructor doesn't work?
                     (!isnothing(ub) && contains_identifier(ub, typevar_name)) ||

JuliaLowering/src/eval.jl

@@ -230,12 +230,8 @@ function to_lowered_expr(mod, ex, ssa_offset=0)
         Core.SSAValue(ex.var_id + ssa_offset)
     elseif k == K"return"
         Core.ReturnNode(to_lowered_expr(mod, ex[1], ssa_offset))
-    elseif is_quoted(k)
-        if k == K"inert"
-            ex[1]
-        else
-            TODO(ex, "Convert SyntaxTree to Expr")
-        end
+    elseif k == K"inert"
+        ex[1]
     elseif k == K"code_info"
         funcname = ex.is_toplevel_thunk ?
             "top-level scope" :
@@ -269,6 +265,11 @@ function to_lowered_expr(mod, ex, ssa_offset=0)
         # TODO: put allow_partial back in once we update to the latest julia
         splice!(args, 4) # allow_partial
         Expr(:new_opaque_closure, args...)
+    elseif k == K"meta"
+        args = Any[to_lowered_expr(mod, e, ssa_offset) for e in children(ex)]
+        # Unpack K"Symbol" QuoteNode as `Expr(:meta)` requires an identifier here.
+        args[1] = args[1].value
+        Expr(:meta, args...)
     else
         # Allowed forms according to https://docs.julialang.org/en/v1/devdocs/ast/
         #

JuliaLowering/src/kinds.jl

@@ -9,6 +9,8 @@ function _register_kinds()
         "BEGIN_EXTENSION_KINDS"
             # atomic fields or accesses (see `@atomic`)
             "atomic"
+            # Flag for @generated parts of a functon
+            "generated"
             # Temporary rooting of identifiers (GC.@preserve)
             "gc_preserve_begin"
             "gc_preserve_end"
@@ -46,6 +48,8 @@ function _register_kinds()
             # Catch-all for additional syntax extensions without the need to
             # extend `Kind`. Known extensions include:
             #   locals, islocal
+            # The content of an assertion is not considered to be quoted, so
+            # use K"Symbol" or K"inert" inside where necessary.
             "extension"
         "END_EXTENSION_KINDS"
 

JuliaLowering/src/linear_ir.jl

@@ -796,11 +796,22 @@ function compile(ctx::LinearIRContext, ex, needs_value, in_tail_pos)
         end
     elseif k == K"gc_preserve_begin"
         makenode(ctx, ex, k, compile_args(ctx, children(ex)))
-    elseif k == K"gc_preserve_end"
+    elseif k == K"gc_preserve_end" || k == K"global" || k == K"const"
         if needs_value
-            throw(LoweringError(ex, "misplaced label in value position"))
+            throw(LoweringError(ex, "misplaced kind $k in value position"))
         end
         emit(ctx, ex)
+        nothing
+    elseif k == K"meta"
+        emit(ctx, ex)
+        if needs_value
+            val = @ast ctx ex "nothing"::K"core"
+            if in_tail_pos
+                emit_return(ctx, val)
+            else
+                val
+            end
+        end
     elseif k == K"_while"
         end_label = make_label(ctx, ex)
         top_label = emit_label(ctx, ex)
@@ -821,12 +832,6 @@ function compile(ctx::LinearIRContext, ex, needs_value, in_tail_pos)
         if needs_value
             compile(ctx, nothing_(ctx, ex), needs_value, in_tail_pos)
         end
-    elseif k == K"global" || k == K"const"
-        if needs_value
-            throw(LoweringError(ex, "misplaced declaration"))
-        end
-        emit(ctx, ex)
-        nothing
     elseif k == K"isdefined" || k == K"captured_local" || k == K"throw_undef_if_not" ||
             k == K"boundscheck"
         if in_tail_pos
@@ -957,7 +962,12 @@ function _renumber(ctx, ssa_rewrites, slot_rewrites, label_table, ex)
             end
         end
     elseif k == K"meta"
-        TODO(ex, "_renumber $k")
+        # Somewhat-hack for Expr(:meta, :generated, gen) which has
+        # weird top-level semantics for `gen`, but we still need to translate
+        # the binding it contains to a globalref.
+        mapchildren(ctx, ex) do e
+            _renumber(ctx, ssa_rewrites, slot_rewrites, label_table, e)
+        end
     elseif is_literal(k) || is_quoted(k)
         ex
     elseif k == K"label"
@@ -968,8 +978,6 @@ function _renumber(ctx, ssa_rewrites, slot_rewrites, label_table, ex)
         mapchildren(ctx, ex) do e
             _renumber(ctx, ssa_rewrites, slot_rewrites, label_table, e)
         end
-        # TODO: foreigncall error check:
-        # "ccall function name and library expression cannot reference local variables"
     end
 end
 

JuliaLowering/src/macro_expansion.jl

@@ -66,7 +66,7 @@ end
 #--------------------------------------------------
 struct MacroContext <: AbstractLoweringContext
     graph::SyntaxGraph
-    macrocall::SyntaxTree
+    macrocall::Union{SyntaxTree,LineNumberNode,SourceRef}
     scope_layer::ScopeLayer
 end