Adapt to change in generated functions.

Co-Authored-By: Tim Besard <383068+maleadt@users.noreply.github.com>

src/eval.jl

@@ -3,7 +3,8 @@ using Base: invokelatest
 dummy() = return
 const dummy_m = which(dummy, Tuple{})
 
-function build_codeinfo(ir::IR)
+function build_codeinfo(ir::IR, world)
+  # TODO: use `world`, if available (for looking up code, and to set edges)
   ir = copy(ir)
   ci = Base.uncompressed_ir(dummy_m)
   ci.inlineable = true
@@ -19,12 +20,34 @@ function build_codeinfo(ir::IR)
   update!(ci, ir)
 end
 
+# JuliaLang/julia#48611: world age is exposed to generated functions, and should be used
+if VERSION >= v"1.10.0-DEV.873"
+
+function func(m::Module, ir::IR)
+  generator = @eval m begin
+    function $(gensym())(world::UInt, source, self,
+                         $([Symbol(:arg, i) for i = 1:length(arguments(ir))]...))
+      return $build_codeinfo($ir, world)
+    end
+  end
+  @eval m begin
+    function $(gensym())($([Symbol(:arg, i) for i = 1:length(arguments(ir))]...))
+      $(Expr(:meta, :generated, generator))
+      $(Expr(:meta, :generated_only))
+    end
+  end
+end
+
+else
+
 function func(m::Module, ir::IR)
   @eval m (@generated function $(gensym())($([Symbol(:arg, i) for i = 1:length(arguments(ir))]...))
-    return $build_codeinfo($ir)
+    return $build_codeinfo($ir, nothing)
   end)
 end
 
+end
+
 func(ir::IR) = func(Main, ir)
 
 evalir(m::Module, ir::IR, args...) = invokelatest(func(m, ir), args...)

src/reflection/dynamo.jl

@@ -71,7 +71,7 @@ end
 # Used only for its CodeInfo
 dummy(args...) = nothing
 
-function dynamo(cache, f, args...)
+function dynamo(cache, world, f, args...)
   try
     ir = transform(f, args...)::Union{IR,Expr,Nothing}
   catch e
@@ -88,20 +88,21 @@ function dynamo(cache, f, args...)
     argnames!(m, :args)
     pushfirst!(m.code.slotnames, Symbol("#self#"))
   else
-    m = @meta dummy(1)
+    m = @meta world dummy(1)
+    m === nothing && error("Error looking up metadata for $f")
     m.code.method_for_inference_limit_heuristics = nothing
   end
   _self = splicearg!(ir)
   prewalk!(x -> x === self ? _self : x, ir)
   return update!(m.code, ir)
 end
 
-function dynamo_lambda(cache, f::Type{<:Lambda{S,I}}) where {S,I}
+function dynamo_lambda(cache, world, f::Type{<:Lambda{S,I}}) where {S,I}
   ir = cache[(S.parameters[2:end]...,)]
   ir = getlambda(ir, I)
   ir = lambdalift!(copy(ir), S, I)
   closureargs!(ir)
-  m = @meta dummy(1)
+  m = @meta world dummy(1)
   m.code.method_for_inference_limit_heuristics = nothing
   return update!(m.code, ir)
 end
@@ -116,6 +117,36 @@ function lifttype(x)
   named ? Expr(:(::), x.args[1], T) : Expr(:(::), T)
 end
 
+const caches = Dict()
+
+if VERSION >= v"1.10.0-DEV.873"
+
+function _generated_ex(world, source, ex)
+    stub = Core.GeneratedFunctionStub(identity, Core.svec(:methodinstance, :args), Core.svec())
+    stub(world, source, ex)
+end
+
+function dynamo_generator(world::UInt, source, self, args)
+  cache = if haskey(caches, self)
+    caches[self]
+  else
+    caches[self] = Dict()
+  end
+  ci = dynamo(cache, world, self, args...)
+  ci isa Expr && return _generated_ex(world, source, ci)
+  return ci
+end
+
+function dynamo_lambda_generator(world::UInt, source, self, args)
+  f = self.parameters[1].parameters[1]
+  cache = caches[f]
+  ci = dynamo_lambda(cache, world, self)
+  ci isa Expr && return _generated_ex(world, source, ci)
+  return ci
+end
+
+end
+
 macro dynamo(ex)
   @capture(shortdef(ex), (name_(args__) = body_) |
                          (name_(args__) where {Ts__} = body_)) ||
@@ -124,12 +155,32 @@ macro dynamo(ex)
   f, T = isexpr(name, :(::)) ?
     (length(name.args) == 1 ? (esc(gensym()), esc(name.args[1])) : esc.(name.args)) :
     (esc(gensym()), :(Core.Typeof($(esc(name)))))
-  gendef = quote
-    local cache = Dict()
-    @generated ($f::$T)($(esc(:args))...) where $(Ts...) =
-      return IRTools.dynamo(cache, $f, args...)
-    @generated (f::IRTools.Inner.Lambda{<:Tuple{<:$T,Vararg{Any}}})(args...) where $(Ts...) =
-      return IRTools.Inner.dynamo_lambda(cache, f)
+  gendef = if VERSION >= v"1.10.0-DEV.873"
+    quote
+      function ($f::$T)($(esc(:args))...) where $(Ts...)
+        $(Expr(:meta, :generated, dynamo_generator))
+        $(Expr(:meta, :generated_only))
+      end
+      function (f::IRTools.Inner.Lambda{<:Tuple{<:$T,Vararg{Any}}})(args...) where $(Ts...)
+        $(Expr(:meta, :generated, dynamo_lambda_generator))
+        $(Expr(:meta, :generated_only))
+      end
+    end
+  else
+    quote
+      @generated function ($f::$T)($(esc(:args))...) where $(Ts...)
+        cache = if haskey($caches, $T)
+          $caches[$T]
+        else
+          $caches[$T] = Dict()
+        end
+        return IRTools.dynamo(cache, nothing, $f, args...)
+      end
+      @generated function (f::IRTools.Inner.Lambda{<:Tuple{<:$T,Vararg{Any}}})(args...) where $(Ts...)
+        cache = $caches[$T]
+        return IRTools.Inner.dynamo_lambda(cache, nothing, f)
+      end
+    end
   end
   quote
     $(isexpr(name, :(::)) || esc(:(function $name end)))

src/reflection/reflection.jl

@@ -40,13 +40,27 @@ See also [`@meta`](@ref).
     julia> IRTools.meta(Tuple{typeof(gcd),Int,Int})
     Metadata for gcd(a::T, b::T) where T<:Union{Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8} in Base at intfuncs.jl:31
 """
-function meta(T; types = T, world = worldcounter())
+function meta(T; types = T, world=nothing)
+  if world === nothing
+    world = worldcounter()
+  end
   F = T.parameters[1]
   F == typeof(invoke) && return invoke_meta(T; world = world)
   F isa DataType && (F.name.module === Core.Compiler ||
                      F <: Core.Builtin ||
                      F <: Core.Builtin) && return nothing
-  _methods = Base._methods_by_ftype(T, -1, world)
+  min_world = Ref{UInt}(typemin(UInt))
+  max_world = Ref{UInt}(typemax(UInt))
+  has_ambig = Ptr{Int32}(C_NULL)  # don't care about ambiguous results
+  _methods = if VERSION >= v"1.7.0-DEV.1297"
+      Base._methods_by_ftype(T, #=mt=# nothing, #=lim=# -1,
+                             world, #=ambig=# false,
+                             min_world, max_world, has_ambig)
+  else
+      Base._methods_by_ftype(T, #=lim=# -1,
+                             world, #=ambig=# false,
+                             min_world, max_world, has_ambig)
+  end
   _methods === nothing && return nothing
   _methods isa Bool && return nothing
   length(_methods) == 0 && return nothing
@@ -88,17 +102,25 @@ function invoke_meta(T; world)
 end
 
 """
-    @meta f(args...)
+    @meta [world] f(args...)
 
 Convenience macro for retrieving metadata without writing a full type signature.
 
     julia> IRTools.@meta gcd(10, 5)
     Metadata for gcd(a::T, b::T) where T<:Union{Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8} in Base at intfuncs.jl:31
 """
-macro meta(ex)
-  isexpr(ex, :call) || error("@meta f(args...)")
-  f, args = ex.args[1], ex.args[2:end]
-  :(meta(typesof($(esc.((f, args...))...))))
+macro meta(ex...)
+  if length(ex) == 1
+    world = nothing
+    call = ex[1]
+  elseif length(ex) == 2
+    world, call = ex
+  else
+    error("@meta [world] f(args...)")
+  end
+  isexpr(call, :call) || error("@meta [world] f(args...)")
+  f, args = call.args[1], call.args[2:end]
+  :(meta(typesof($(esc.((f, args...))...)); world=$(esc(world))))
 end
 
 function code_ir(f, T)

src/reflection/utils.jl

@@ -56,7 +56,6 @@ function slots!(ci::CodeInfo)
       x isa Slot || return x
       haskey(ss, x) && return ss[x]
       @static if VERSION >= v"1.10.0-DEV.870"
-        isnothing(ci.slottypes) && (ci.slottypes = Any[])
         push!(ci.slottypes, x.type)
       end
       push!(ci.slotnames, x.id)
@@ -134,6 +133,7 @@ end
 
 @static if VERSION >= v"1.10.0-DEV.870"
   function replace_code_newstyle!(ci, ir, _)
+    isnothing(ci.slottypes) && (ci.slottypes = Any[])
     return Core.Compiler.replace_code_newstyle!(ci, ir)
   end
 elseif VERSION < v"1.8.0-DEV.267"