improve return-type reflections of opaque closure (JuliaLang#44743)

aviatesk · pull[bot] · commit 871e3d7746d4 · 2023-08-31T23:35:56.000Z
diff --git a/base/reflection.jl b/base/reflection.jl
@@ -1276,16 +1276,16 @@ function code_typed_by_type(@nospecialize(tt::Type);
     return asts
 end
 
-function code_typed_opaque_closure(@nospecialize(closure::Core.OpaqueClosure);
-        optimize=true,
-        debuginfo::Symbol=:default,
-        interp = Core.Compiler.NativeInterpreter(closure.world))
+function code_typed_opaque_closure(@nospecialize(oc::Core.OpaqueClosure);
+    debuginfo::Symbol=:default, __...)
     ccall(:jl_is_in_pure_context, Bool, ()) && error("code reflection cannot be used from generated functions")
-    m = closure.source
+    m = oc.source
     if isa(m, Method)
         code = _uncompressed_ir(m, m.source)
         debuginfo === :none && remove_linenums!(code)
-        return Any[(code => code.rettype)]
+        # intersect the declared return type and the inferred return type (if available)
+        rt = typeintersect(code.rettype, typeof(oc).parameters[2])
+        return Any[code => rt]
     else
         error("encountered invalid Core.OpaqueClosure object")
     end
@@ -1295,6 +1295,10 @@ function return_types(@nospecialize(f), @nospecialize(types=default_tt(f));
                       world = get_world_counter(),
                       interp = Core.Compiler.NativeInterpreter(world))
     ccall(:jl_is_in_pure_context, Bool, ()) && error("code reflection cannot be used from generated functions")
+    if isa(f, Core.OpaqueClosure)
+        _, rt = only(code_typed_opaque_closure(f))
+        return Any[rt]
+    end
     types = to_tuple_type(types)
     rt = []
     for match in _methods(f, types, -1, world)::Vector
diff --git a/test/opaque_closure.jl b/test/opaque_closure.jl
@@ -227,6 +227,15 @@ const GLOBAL_OPAQUE_CLOSURE = @opaque () -> 123
 call_global_opaque_closure() = GLOBAL_OPAQUE_CLOSURE()
 @test call_global_opaque_closure() == 123
 
+let foo::Int = 42
+    Base.Experimental.@force_compile
+    oc = Base.Experimental.@opaque a::Int->sin(a) + cos(foo)
+
+    @test only(Base.return_types(oc, (Int,))) === Float64
+    code, rt = first(code_typed(oc, (Int,)))
+    @test rt === Float64
+end
+
 let oc = @opaque a->sin(a)
     @test length(code_typed(oc, (Int,))) == 1
 end
