Implement OpaqueClosure return type narrowing (#39917)

Allows the optimizer to rewrite the return type parameter
of the OpaqueClosure based on inference results of the
partially specialized (i.e. specialized on the closure
environment, but not on the argument types of the opaque
closure). This helps by forcing an inference barrier to
occur if the PartialOpaque-ness information gets lost,
causing a re-infer with at least the rt information we
have from inference.

Author: Keno

base/compiler/abstractinterpretation.jl

@@ -1467,8 +1467,9 @@ function abstract_eval_statement(interp::AbstractInterpreter, @nospecialize(e),
                 if isa(t, PartialOpaque)
                     # Infer this now so that the specialization is available to
                     # optimization.
-                    abstract_call_opaque_closure(interp, t,
+                    callinfo = abstract_call_opaque_closure(interp, t,
                         most_general_argtypes(t), sv)
+                    sv.stmt_info[sv.currpc] = OpaqueClosureCreateInfo(callinfo)
                 end
             end
         end

base/compiler/ssair/inlining.jl

@@ -1048,6 +1048,32 @@ function inline_invoke!(ir::IRCode, idx::Int, sig::Signature, info::InvokeCallIn
     return nothing
 end
 
+function narrow_opaque_closure!(ir::IRCode, stmt::Expr, @nospecialize(info), state::InliningState)
+    if isa(info, OpaqueClosureCreateInfo)
+        unspec_call_info = info.unspec.info
+        if isa(unspec_call_info, ConstCallInfo)
+            unspec_call_info = unspec_call_info.call
+        end
+        isa(unspec_call_info, OpaqueClosureCallInfo) || return
+        lbt = argextype(stmt.args[3], ir, ir.sptypes)
+        lb, exact = instanceof_tfunc(lbt)
+        exact || return
+        ubt = argextype(stmt.args[4], ir, ir.sptypes)
+        ub, exact = instanceof_tfunc(ubt)
+        exact || return
+        # Narrow opaque closure type
+
+        newT = widenconst(tmeet(tmerge(lb, info.unspec.rt), ub))
+        if newT != ub
+            # N.B.: Narrowing the ub requires a backdge on the mi whose type
+            # information we're using, since a change in that function may
+            # invalidate ub result.
+            push!(state.et, unspec_call_info.mi)
+            stmt.args[4] = newT
+        end
+    end
+end
+
 # Handles all analysis and inlining of intrinsics and builtins. In particular,
 # this method does not access the method table or otherwise process generic
 # functions.
@@ -1057,6 +1083,9 @@ function process_simple!(ir::IRCode, todo::Vector{Pair{Int, Any}}, idx::Int, sta
     if stmt.head === :splatnew
         inline_splatnew!(ir, idx)
         return nothing
+    elseif stmt.head === :new_opaque_closure
+        narrow_opaque_closure!(ir, stmt, ir.stmts[idx][:info], state)
+        return nothing
     end
 
     stmt.head === :call || return nothing

base/compiler/stmtinfo.jl

@@ -114,6 +114,10 @@ struct OpaqueClosureCallInfo
     mi::MethodInstance
 end
 
+struct OpaqueClosureCreateInfo
+    unspec::CallMeta
+end
+
 # Stmt infos that are used by external consumers, but not by optimization.
 # These are not produced by default and must be explicitly opted into by
 # the AbstractInterpreter.

test/opaque_closure.jl

@@ -20,7 +20,7 @@ let ci = @code_lowered const_int()
             Expr(:opaque_closure_method, 0, lno, ci)))
     end
 end
-@test_broken isa(oc_simple_inf(), Core.OpaqueClosure{Tuple{}, Int})
+@test isa(oc_simple_inf(), Core.OpaqueClosure{Tuple{}, Int})
 @test oc_simple_inf()() == 1
 
 struct OcClos2Int
@@ -71,8 +71,8 @@ let ci = @code_lowered OcClos1Any(1)()
             :x))
     end
 end
-@test_broken isa(oc_infer_pass_clos(1), Core.OpaqueClosure{Tuple{}, typeof(1)})
-@test_broken isa(oc_infer_pass_clos("a"), Core.OpaqueClosure{Tuple{}, typeof("a")})
+@test isa(oc_infer_pass_clos(1), Core.OpaqueClosure{Tuple{}, typeof(1)})
+@test isa(oc_infer_pass_clos("a"), Core.OpaqueClosure{Tuple{}, typeof("a")})
 @test oc_infer_pass_clos(1)() == 1
 @test oc_infer_pass_clos("a")() == "a"
 
@@ -131,12 +131,14 @@ function test_oc_world_age end
 mk_oc_world_age() = @opaque ()->test_oc_world_age()
 g_world_age = @opaque ()->test_oc_world_age()
 h_world_age = mk_oc_world_age()
+@test isa(h_world_age, Core.OpaqueClosure{Tuple{}, Union{}})
 test_oc_world_age() = 1
 @test_throws MethodError g_world_age()
 @test_throws MethodError h_world_age()
 @test mk_oc_world_age()() == 1
 g_world_age = @opaque ()->test_oc_world_age()
 @test g_world_age() == 1
+@test isa(mk_oc_world_age(), Core.OpaqueClosure{Tuple{}, Int})
 
 # Evil, dynamic Vararg stuff (don't do this - made to work for consistency)
 function maybe_opaque(isva::Bool)