World age fixes for macro expansion and generated functions (#63)

* Support for expanding macros in a specific world age
* Connect macro expansion world to Core hook
* Fixes for CodeInfo generation in GeneratedFunctionStub - ensure `isva`
  and edges to bindings are set as is done in Base.
* Test non-generated branch of generated function

Author: c42f

src/eval.jl

@@ -1,14 +1,14 @@
-function lower(mod::Module, ex0, expr_compat_mode=false)
-    ctx1, ex1 = expand_forms_1(  mod,  ex0, expr_compat_mode)
+function lower(mod::Module, ex0; expr_compat_mode=false, world=Base.get_world_counter())
+    ctx1, ex1 = expand_forms_1(  mod,  ex0, expr_compat_mode, world)
     ctx2, ex2 = expand_forms_2(  ctx1, ex1)
     ctx3, ex3 = resolve_scopes(  ctx2, ex2)
     ctx4, ex4 = convert_closures(ctx3, ex3)
     ctx5, ex5 = linearize_ir(    ctx4, ex4)
     ex5
 end
 
-function macroexpand(mod::Module, ex, expr_compat_mode=false)
-    ctx1, ex1 = expand_forms_1(mod, ex, expr_compat_mode)
+function macroexpand(mod::Module, ex; expr_compat_mode=false, world=Base.get_world_counter())
+    ctx1, ex1 = expand_forms_1(mod, ex, expr_compat_mode, world)
     ex1
 end
 
@@ -354,7 +354,7 @@ function Core.eval(mod::Module, ex::SyntaxTree; expr_compat_mode::Bool=false)
         end
         return x
     end
-    linear_ir = lower(mod, ex, expr_compat_mode)
+    linear_ir = lower(mod, ex; expr_compat_mode)
     expr_form = to_lowered_expr(mod, linear_ir, 0)
     eval(mod, expr_form)
 end

src/hooks.jl

@@ -17,7 +17,7 @@ function core_lowering_hook(@nospecialize(code), mod::Module,
 
     st0 = code isa Expr ? expr_to_syntaxtree(code, LineNumberNode(line, file)) : code
     try
-        ctx1, st1 = expand_forms_1(  mod,  st0, true)
+        ctx1, st1 = expand_forms_1(  mod,  st0, true, world)
         ctx2, st2 = expand_forms_2(  ctx1, st1)
         ctx3, st3 = resolve_scopes(  ctx2, st2)
         ctx4, st4 = convert_closures(ctx3, st3)

src/macro_expansion.jl

@@ -21,11 +21,12 @@ struct MacroExpansionContext{GraphType} <: AbstractLoweringContext
     scope_layers::Vector{ScopeLayer}
     scope_layer_stack::Vector{LayerId}
     expr_compat_mode::Bool
+    macro_world::UInt
 end
 
-function MacroExpansionContext(graph::SyntaxGraph, mod::Module, expr_compat_mode::Bool)
+function MacroExpansionContext(graph::SyntaxGraph, mod::Module, expr_compat_mode::Bool, world::UInt)
     layers = ScopeLayer[ScopeLayer(1, mod, 0, false)]
-    MacroExpansionContext(graph, Bindings(), layers, LayerId[length(layers)], expr_compat_mode)
+    MacroExpansionContext(graph, Bindings(), layers, LayerId[length(layers)], expr_compat_mode, world)
 end
 
 current_layer(ctx::MacroExpansionContext) = ctx.scope_layers[last(ctx.scope_layer_stack)]
@@ -217,11 +218,10 @@ function expand_macro(ctx, ex)
     # age changes concurrently.
     # 
     # TODO: Allow this to be passed in
-    macro_world = Base.get_world_counter()
-    if hasmethod(macfunc, Tuple{typeof(mctx), typeof.(raw_args)...}; world=macro_world)
+    if hasmethod(macfunc, Tuple{typeof(mctx), typeof.(raw_args)...}; world=ctx.macro_world)
         macro_args = prepare_macro_args(ctx, mctx, raw_args)
         expanded = try
-            Base.invoke_in_world(macro_world, macfunc, macro_args...)
+            Base.invoke_in_world(ctx.macro_world, macfunc, macro_args...)
         catch exc
             newexc = exc isa MacroExpansionError ?
                 MacroExpansionError(mctx, exc.ex, exc.msg, exc.position, exc.err) :
@@ -260,14 +260,14 @@ function expand_macro(ctx, ex)
             push!(macro_args, Expr(arg))
         end
         expanded = try
-            Base.invoke_in_world(macro_world, macfunc, macro_args...)
+            Base.invoke_in_world(ctx.macro_world, macfunc, macro_args...)
         catch exc
             if exc isa MethodError && exc.f === macfunc
-                if !isempty(methods_in_world(macfunc, Tuple{typeof(mctx), Vararg{Any}}, macro_world))
+                if !isempty(methods_in_world(macfunc, Tuple{typeof(mctx), Vararg{Any}}, ctx.macro_world))
                     # If the macro has at least some methods implemented in the
                     # new style, assume the user meant to call one of those
                     # rather than any old-style macro methods which might exist
-                    exc = MethodError(macfunc, (prepare_macro_args(ctx, mctx, raw_args)..., ), macro_world)
+                    exc = MethodError(macfunc, (prepare_macro_args(ctx, mctx, raw_args)..., ), ctx.macro_world)
                 end
             end
             rethrow(MacroExpansionError(mctx, ex, "Error expanding macro", :all, exc))
@@ -280,7 +280,7 @@ function expand_macro(ctx, ex)
         # Module scope for the returned AST is the module where this particular
         # method was defined (may be different from `parentmodule(macfunc)`)
         mod_for_ast = lookup_method_instance(macfunc, macro_args,
-                                             macro_world).def.module
+                                             ctx.macro_world).def.module
         new_layer = ScopeLayer(length(ctx.scope_layers)+1, mod_for_ast,
                                current_layer_id(ctx), true)
         push!(ctx.scope_layers, new_layer)
@@ -460,18 +460,18 @@ function expand_forms_1(ctx::MacroExpansionContext, ex::SyntaxTree)
     end
 end
 
-function expand_forms_1(mod::Module, ex::SyntaxTree, expr_compat_mode::Bool)
+function expand_forms_1(mod::Module, ex::SyntaxTree, expr_compat_mode::Bool, macro_world::UInt)
     graph = ensure_attributes(syntax_graph(ex),
                               var_id=IdTag,
                               scope_layer=LayerId,
                               __macro_ctx__=Nothing,
                               meta=CompileHints)
-    ctx = MacroExpansionContext(graph, mod, expr_compat_mode)
+    ctx = MacroExpansionContext(graph, mod, expr_compat_mode, macro_world)
     ex2 = expand_forms_1(ctx, reparent(ctx, ex))
     graph2 = delete_attributes(graph, :__macro_ctx__)
     # TODO: Returning the context with pass-specific mutable data is a bad way
     # to carry state into the next pass. We might fix this by attaching such
     # data to the graph itself as global attributes?
-    ctx2 = MacroExpansionContext(graph2, ctx.bindings, ctx.scope_layers, LayerId[], expr_compat_mode)
+    ctx2 = MacroExpansionContext(graph2, ctx.bindings, ctx.scope_layers, LayerId[], expr_compat_mode, macro_world)
     return ctx2, reparent(ctx2, ex2)
 end

src/runtime.jl

@@ -338,12 +338,9 @@ function (g::GeneratedFunctionStub)(world::UInt, source::Method, @nospecialize a
     #
     # TODO: Reduce duplication where possible.
 
-    mod = parentmodule(g.gen)
-
     # Attributes from parsing
     graph = ensure_attributes(SyntaxGraph(), kind=Kind, syntax_flags=UInt16, source=SourceAttrType,
                               value=Any, name_val=String)
-
     # Attributes for macro expansion
     graph = ensure_attributes(graph,
                               var_id=IdTag,
@@ -355,8 +352,11 @@ function (g::GeneratedFunctionStub)(world::UInt, source::Method, @nospecialize a
                               is_toplevel_thunk=Bool
                               )
 
-    # Macro expansion
-    ctx1 = MacroExpansionContext(graph, mod, false)
+    # Macro expansion. Looking at Core.GeneratedFunctionStub, it seems that
+    # macros emitted by the generator are currently expanded in the latest
+    # world, so do that for compatibility.
+    macro_world = typemax(UInt)
+    ctx1 = MacroExpansionContext(graph, source.module, false, macro_world)
 
     # Run code generator - this acts like a macro expander and like a macro
     # expander it gets a MacroContext.
@@ -375,7 +375,8 @@ function (g::GeneratedFunctionStub)(world::UInt, source::Method, @nospecialize a
     # Expand any macros emitted by the generator
     ex1 = expand_forms_1(ctx1, reparent(ctx1, ex0))
     ctx1 = MacroExpansionContext(delete_attributes(graph, :__macro_ctx__),
-                                 ctx1.bindings, ctx1.scope_layers, LayerId[], false)
+                                 ctx1.bindings, ctx1.scope_layers,
+                                 LayerId[], false, macro_world)
     ex1 = reparent(ctx1, ex1)
 
     # Desugaring
@@ -398,6 +399,21 @@ function (g::GeneratedFunctionStub)(world::UInt, source::Method, @nospecialize a
     ctx5, ex5 = linearize_ir(ctx4, ex4)
     ci = to_lowered_expr(mod, ex5)
     @assert ci isa Core.CodeInfo
+
+    # See GeneratedFunctionStub code in base/expr.jl
+    ci.isva = source.isva
+    code = ci.code
+    bindings = IdSet{Core.Binding}()
+    for i = 1:length(code)
+        stmt = code[i]
+        if isa(stmt, GlobalRef)
+            push!(bindings, convert(Core.Binding, stmt))
+        end
+    end
+    if !isempty(bindings)
+        ci.edges = Core.svec(bindings...)
+    end
+
     return ci
 end
 

src/syntax_graph.jl

@@ -689,7 +689,7 @@ macro SyntaxTree(ex_old)
     ex = _find_SyntaxTree_macro(full_ex, __source__.line)
     # 4. Do the first step of JuliaLowering's syntax lowering to get
     # syntax interpolations to work
-    _, ex1 = expand_forms_1(__module__, ex, false)
+    _, ex1 = expand_forms_1(__module__, ex, false, Base.tls_world_age())
     @assert kind(ex1) == K"call" && ex1[1].value == interpolate_ast
     Expr(:call, :interpolate_ast, SyntaxTree, ex1[3][1],
          map(e->_scope_layer_1_to_esc!(Expr(e)), ex1[4:end])...)

test/demo.jl

@@ -33,7 +33,7 @@ function formatsrc(ex; kws...)
 end
 
 function debug_lower(mod, ex; expr_compat_mode=false, verbose=false, do_eval=false)
-    ctx1, ex_macroexpand = JuliaLowering.expand_forms_1(mod, ex, expr_compat_mode)
+    ctx1, ex_macroexpand = JuliaLowering.expand_forms_1(mod, ex, expr_compat_mode, Base.get_world_counter())
 
     verbose && @info "Macro expanded" formatsrc(ex_macroexpand, color_by=:scope_layer)
 

test/functions.jl

@@ -444,19 +444,32 @@ end
         function f_partially_gen(x::NTuple{N,T}) where {N,T}
             shared = :shared_stuff
             if @generated
+                if N == 2
+                    error("intentionally broken codegen (will trigger nongen branch)")
+                end
                 quote
-                    unshared = ($x, $N, $T)
+                    unshared = (:gen, ($x, $N, $T))
                 end
             else
-                # Uuuum. How do we actually test both sides of this branch???
-                unshared = :nongen # (typeof(x), N, T)
+                unshared = (:nongen, (typeof(x), N, T))
             end
             (shared, unshared)
         end
 
-        f_partially_gen((1,2,3,4,5))
+        (f_partially_gen((1,2)), f_partially_gen((1,2,3,4,5)))
     end
-    """) == (:shared_stuff, (NTuple{5,Int}, 5, Int))
+    """) == ((:shared_stuff, (:nongen, (NTuple{2,Int}, 2, Int))),
+             (:shared_stuff, (:gen, (NTuple{5,Int}, 5, Int))))
+
+    # Test generated function edges to bindings
+    # (see also https://github.com/JuliaLang/julia/pull/57230)
+    JuliaLowering.include_string(test_mod, raw"""
+    const delete_me = 4
+    @generated f_generated_return_delete_me() = return :(delete_me)
+    """)
+    @test test_mod.f_generated_return_delete_me() == 4
+    Base.delete_binding(test_mod, :delete_me)
+    @test_throws UndefVarError test_mod.f_generated_return_delete_me()
 end
 
 @testset "Broadcast" begin

test/macros.jl

@@ -1,8 +1,6 @@
-module macros
+@testset "macro tests" begin
 
-using JuliaLowering, Test
-
-module test_mod end
+test_mod = Module(:macro_test)
 
 JuliaLowering.include_string(test_mod, raw"""
 module M
@@ -110,12 +108,33 @@ M.@recursive 3
 """) == (3, (2, (1, 0)))
 
 ex = JuliaLowering.parsestmt(JuliaLowering.SyntaxTree, "M.@outer()", filename="foo.jl")
-ctx, expanded = JuliaLowering.expand_forms_1(test_mod, ex, false)
+ctx, expanded = JuliaLowering.expand_forms_1(test_mod, ex, false, Base.get_world_counter())
 @test JuliaLowering.sourcetext.(JuliaLowering.flattened_provenance(expanded[2])) == [
     "M.@outer()"
     "@inner"
     "y"
 ]
+
+# World age support for macro expansion
+JuliaLowering.include_string(test_mod, raw"""
+macro world_age_test()
+    :(world1)
+end
+""")
+world1 = Base.get_world_counter()
+JuliaLowering.include_string(test_mod, raw"""
+macro world_age_test()
+    :(world2)
+end
+""")
+world2 = Base.get_world_counter()
+
+call_world_arg_test = JuliaLowering.parsestmt(JuliaLowering.SyntaxTree, "@world_age_test()")
+@test JuliaLowering.expand_forms_1(test_mod, call_world_arg_test, false, world1)[2] ≈
+    @ast_ "world1"::K"Identifier"
+@test JuliaLowering.expand_forms_1(test_mod, call_world_arg_test, false, world2)[2] ≈
+    @ast_ "world2"::K"Identifier"
+
 # Layer parenting
 @test expanded[1].scope_layer == 2
 @test expanded[2].scope_layer == 3
@@ -233,7 +252,7 @@ end
 catch exc
     sprint(showerror, exc)
 end == """
-MacroExpansionError while expanding @oldstyle_error in module Main.macros.test_mod:
+MacroExpansionError while expanding @oldstyle_error in module Main.macro_test:
 @oldstyle_error
 └─────────────┘ ── Error expanding macro
 Caused by:
@@ -290,10 +309,10 @@ end
     err = try
         JuliaLowering.include_string(test_mod, "@sig_mismatch(1, 2, 3, 4)") === 1
     catch exc
-        sprint(showerror, exc, context=:module=>@__MODULE__)
+        sprint(showerror, exc, context=:module=>test_mod)
     end
     @test startswith(err, """
-    MacroExpansionError while expanding @sig_mismatch in module Main.macros.test_mod:
+    MacroExpansionError while expanding @sig_mismatch in module Main.macro_test:
     @sig_mismatch(1, 2, 3, 4)
     └───────────────────────┘ ── Error expanding macro
     Caused by:
@@ -329,4 +348,4 @@ end
     """) === (false, false, false, false)
 end
 
-end # module macros
+end