oc: code_typed support for optimized opaque closures

Reflection version of #53878.
JuliaLang · Apr 2, 2024 · e01c3e1 · e01c3e1
1 parent 1fedcab
commit e01c3e1
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 10 deletions.
diff --git a/base/reflection.jl b/base/reflection.jl
@@ -1584,7 +1584,7 @@ julia> code_typed(+, (Float64, Float64))
 """
 function code_typed(@nospecialize(f), @nospecialize(types=default_tt(f)); kwargs...)
     if isa(f, Core.OpaqueClosure)
-        return code_typed_opaque_closure(f; kwargs...)
+        return code_typed_opaque_closure(f, types; kwargs...)
     end
     tt = signature_type(f, types)
     return code_typed_by_type(tt; kwargs...)
@@ -1639,20 +1639,35 @@ function code_typed_by_type(@nospecialize(tt::Type);
     return asts
 end
 
-function get_oc_code_rt(@nospecialize(oc::Core.OpaqueClosure))
+function get_oc_code_rt(@nospecialize(oc::Core.OpaqueClosure), types=Tuple, optimize::Bool=true)
     ccall(:jl_is_in_pure_context, Bool, ()) && error("code reflection cannot be used from generated functions")
     m = oc.source
     if isa(m, Method)
-        code = _uncompressed_ir(m)
-        return Pair{CodeInfo,Any}(code, typeof(oc).parameters[2])
+        if isdefined(m, :source)
+            if optimize
+                tt = Tuple{typeof(oc.captures), to_tuple_type(types).parameters...}
+                mi = Core.Compiler.specialize_method(m, tt, Core.svec())
+                interp = Core.Compiler.NativeInterpreter(m.primary_world)
+                return Core.Compiler.typeinf_code(interp, mi, optimize)
+            else
+                code = _uncompressed_ir(m)
+                return Pair{CodeInfo,Any}(code, typeof(oc).parameters[2])
+            end
+        else
+            # OC constructed from optimized IR
+            codeinst = m.specializations.cache
+            return Pair{CodeInfo, Any}(codeinst.inferred, codeinst.rettype)
+        end
     else
         error("encountered invalid Core.OpaqueClosure object")
     end
 end
 
-function code_typed_opaque_closure(@nospecialize(oc::Core.OpaqueClosure);
-                                   debuginfo::Symbol=:default, _...)
-    (code, rt) = get_oc_code_rt(oc)
+function code_typed_opaque_closure(@nospecialize(oc::Core.OpaqueClosure), types=Tuple;
+                                   debuginfo::Symbol=:default,
+                                   optimize::Bool=true,
+                                   _...)
+    (code, rt) = get_oc_code_rt(oc, types)
     debuginfo === :none && remove_linenums!(code)
     return Any[Pair{CodeInfo,Any}(code, rt)]
 end
@@ -1807,7 +1822,7 @@ function return_types(@nospecialize(f), @nospecialize(types=default_tt(f));
                       interp::Core.Compiler.AbstractInterpreter=Core.Compiler.NativeInterpreter(world))
     check_generated_context(world)
     if isa(f, Core.OpaqueClosure)
-        _, rt = only(code_typed_opaque_closure(f))
+        _, rt = only(code_typed_opaque_closure(f, types))
         return Any[rt]
     end
     if isa(f, Core.Builtin)
@@ -1876,7 +1891,7 @@ function infer_return_type(@nospecialize(f), @nospecialize(types=default_tt(f));
                            interp::Core.Compiler.AbstractInterpreter=Core.Compiler.NativeInterpreter(world))
     check_generated_context(world)
     if isa(f, Core.OpaqueClosure)
-        return last(only(code_typed_opaque_closure(f)))
+        return last(only(code_typed_opaque_closure(f, types)))
     end
     if isa(f, Core.Builtin)
         return _builtin_return_type(interp, f, types)

diff --git a/test/opaque_closure.jl b/test/opaque_closure.jl
@@ -239,7 +239,13 @@ let foo::Int = 42
 end
 
 let oc = @opaque a->sin(a)
-    @test length(code_typed(oc, (Int,))) == 1
+    let opt = code_typed(oc, (Int,))
+        @test length(opt) == 1
+        @test opt[1][2] === Float64
+    end
+    let unopt = code_typed(oc, (Int,); optimize=false)
+        @test length(unopt) == 1
+    end
 end
 
 # constructing an opaque closure from IRCode
@@ -257,6 +263,7 @@ end
 let ir = first(only(Base.code_ircode(sin, (Int,))))
     @test OpaqueClosure(ir)(42) == sin(42)
     @test OpaqueClosure(ir)(42) == sin(42) # the `OpaqueClosure(::IRCode)` constructor should be non-destructive
+    @test length(code_typed(OpaqueClosure(ir))) == 1
     ir = first(only(Base.code_ircode(sin, (Float64,))))
     @test OpaqueClosure(ir)(42.) == sin(42.)
     @test OpaqueClosure(ir)(42.) == sin(42.) # the `OpaqueClosure(::IRCode)` constructor should be non-destructive