Add PartialOpaque lattice element for OpaqueClosure (#39512)

This adds a lattice element for tracking OpaqueClosures in inference,
but does not yet do anything with it. The reason I'm separating
this out is that just the introduction of the lattice element
raises some tricky issues. In particular, the lattice element
refers back to the OpaqueClosure method, which we currently
don't support in the serializer. I played with several ways
of adding support for that, but in the end it all ended up
super complicated for questionable benefit, so in this PR,
CodeInstances that get inferred to `PartialOpaque` get
omitted during serialization (i.e. they will be reinfered
upon loading the .ji).

Author: Keno

base/boot.jl

@@ -424,6 +424,7 @@ eval(Core, :(CodeInstance(mi::MethodInstance, @nospecialize(rettype), @nospecial
                     mi, rettype, inferred_const, inferred, const_flags, min_world, max_world)))
 eval(Core, :(Const(@nospecialize(v)) = $(Expr(:new, :Const, :v))))
 eval(Core, :(PartialStruct(@nospecialize(typ), fields::Array{Any, 1}) = $(Expr(:new, :PartialStruct, :typ, :fields))))
+eval(Core, :(PartialOpaque(@nospecialize(typ), @nospecialize(env), isva::Bool, parent::MethodInstance, source::Method) = $(Expr(:new, :PartialOpaque, :typ, :env, :isva, :parent, :source))))
 eval(Core, :(MethodMatch(@nospecialize(spec_types), sparams::SimpleVector, method::Method, fully_covers::Bool) =
     $(Expr(:new, :MethodMatch, :spec_types, :sparams, :method, :fully_covers))))
 

base/compiler/abstractinterpretation.jl

@@ -227,6 +227,7 @@ function const_prop_profitable(@nospecialize(arg))
             const_prop_profitable(b) && return true
         end
     end
+    isa(arg, PartialOpaque) && return true
     isa(arg, Const) || return true
     val = arg.val
     # don't consider mutable values or Strings useful constants
@@ -268,7 +269,7 @@ function abstract_call_method_with_const_args(interp::AbstractInterpreter, @nosp
     # see if any or all of the arguments are constant and propagating constants may be worthwhile
     for a in argtypes
         a = widenconditional(a)
-        if allconst && !isa(a, Const) && !isconstType(a) && !isa(a, PartialStruct)
+        if allconst && !isa(a, Const) && !isconstType(a) && !isa(a, PartialStruct) && !isa(a, PartialOpaque)
             allconst = false
         end
         if !haveconst && has_nontrivial_const_info(a) && const_prop_profitable(a)
@@ -688,7 +689,7 @@ end
 function abstract_apply(interp::AbstractInterpreter, @nospecialize(itft), @nospecialize(aft), aargtypes::Vector{Any}, sv::InferenceState,
                         max_methods::Int = InferenceParams(interp).MAX_METHODS)
     aftw = widenconst(aft)
-    if !isa(aft, Const) && (!isType(aftw) || has_free_typevars(aftw))
+    if !isa(aft, Const) && !isa(aft, PartialOpaque) && (!isType(aftw) || has_free_typevars(aftw))
         if !isconcretetype(aftw) || (aftw <: Builtin)
             add_remark!(interp, sv, "Core._apply_iterate called on a function of a non-concrete type")
             # bail now, since it seems unlikely that abstract_call will be able to do any better after splitting
@@ -1057,6 +1058,20 @@ function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
     return abstract_call_gf_by_type(interp, f, argtypes, atype, sv, max_methods)
 end
 
+function abstract_call_opaque_closure(interp::AbstractInterpreter, closure::PartialOpaque, argtypes::Vector{Any}, sv::InferenceState)
+    return CallMeta(Any, nothing)
+end
+
+function most_general_argtypes(closure::PartialOpaque)
+    ret = Any[]
+    cc = widenconst(closure)
+    argt = unwrap_unionall(cc).parameters[1]
+    if !isa(argt, DataType) || argt.name !== typename(Tuple)
+        argt = Tuple
+    end
+    return most_general_argtypes(closure.source, argt, closure.isva)
+end
+
 # call where the function is any lattice element
 function abstract_call(interp::AbstractInterpreter, fargs::Union{Nothing,Vector{Any}}, argtypes::Vector{Any},
                        sv::InferenceState, max_methods::Int = InferenceParams(interp).MAX_METHODS)
@@ -1068,10 +1083,14 @@ function abstract_call(interp::AbstractInterpreter, fargs::Union{Nothing,Vector{
         f = ft.parameters[1]
     elseif isa(ft, DataType) && isdefined(ft, :instance)
         f = ft.instance
+    elseif isa(ft, PartialOpaque)
+        return abstract_call_opaque_closure(interp, ft, argtypes, sv)
+    elseif isa(unwrap_unionall(ft), DataType) && unwrap_unionall(ft).name === typename(Core.OpaqueClosure)
+        return CallMeta(rewrap_unionall(unwrap_unionall(ft).parameters[2], ft), false)
     else
         # non-constant function, but the number of arguments is known
         # and the ft is not a Builtin or IntrinsicFunction
-        if typeintersect(widenconst(ft), Builtin) != Union{}
+        if typeintersect(widenconst(ft), Union{Builtin, Core.OpaqueClosure}) != Union{}
             add_remark!(interp, sv, "Could not identify method table for call")
             return CallMeta(Any, false)
         end
@@ -1173,25 +1192,34 @@ function abstract_eval_value(interp::AbstractInterpreter, @nospecialize(e), vtyp
     end
 end
 
+function collect_argtypes(interp::AbstractInterpreter, ea::Vector{Any}, vtypes::VarTable, sv::InferenceState)
+    n = length(ea)
+    argtypes = Vector{Any}(undef, n)
+    @inbounds for i = 1:n
+        ai = abstract_eval_value(interp, ea[i], vtypes, sv)
+        if bail_out_statement(interp, ai, sv)
+            return Bottom
+        end
+        argtypes[i] = ai
+    end
+    return argtypes
+end
+
 function abstract_eval_statement(interp::AbstractInterpreter, @nospecialize(e), vtypes::VarTable, sv::InferenceState)
     if !isa(e, Expr)
         return abstract_eval_special_value(interp, e, vtypes, sv)
     end
     e = e::Expr
     if e.head === :call
         ea = e.args
-        n = length(ea)
-        argtypes = Vector{Any}(undef, n)
-        @inbounds for i = 1:n
-            ai = abstract_eval_value(interp, ea[i], vtypes, sv)
-            if bail_out_statement(interp, ai, sv)
-                return Bottom
-            end
-            argtypes[i] = ai
+        argtypes = collect_argtypes(interp, ea, vtypes, sv)
+        if argtypes === Bottom
+            t = Bottom
+        else
+            callinfo = abstract_call(interp, ea, argtypes, sv)
+            sv.stmt_info[sv.currpc] = callinfo.info
+            t = callinfo.rt
         end
-        callinfo = abstract_call(interp, ea, argtypes, sv)
-        sv.stmt_info[sv.currpc] = callinfo.info
-        t = callinfo.rt
     elseif e.head === :new
         t = instanceof_tfunc(abstract_eval_value(interp, e.args[1], vtypes, sv))[1]
         if isconcretetype(t) && !t.mutable
@@ -1242,6 +1270,24 @@ function abstract_eval_statement(interp::AbstractInterpreter, @nospecialize(e),
                 t = PartialStruct(t, at.fields)
             end
         end
+    elseif e.head === :new_opaque_closure
+        t = Union{}
+        if length(e.args) >= 5
+            ea = e.args
+            argtypes = collect_argtypes(interp, ea, vtypes, sv)
+            if argtypes === Bottom
+                t = Bottom
+            else
+                t = _opaque_closure_tfunc(argtypes[1], argtypes[2], argtypes[3],
+                    argtypes[4], argtypes[5], argtypes[6:end], sv.linfo)
+                if isa(t, PartialOpaque)
+                    # Infer this now so that the specialization is available to
+                    # optimization.
+                    abstract_call_opaque_closure(interp, t,
+                        most_general_argtypes(t), sv)
+                end
+            end
+        end
     elseif e.head === :foreigncall
         abstract_eval_value(interp, e.args[1], vtypes, sv)
         t = sp_type_rewrap(e.args[2], sv.linfo, true)
@@ -1404,7 +1450,7 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
             elseif isa(stmt, ReturnNode)
                 pc´ = n + 1
                 rt = widenconditional(abstract_eval_value(interp, stmt.val, s[pc], frame))
-                if !isa(rt, Const) && !isa(rt, Type) && !isa(rt, PartialStruct)
+                if !isa(rt, Const) && !isa(rt, Type) && !isa(rt, PartialStruct) && !isa(rt, PartialOpaque)
                     # only propagate information we know we can store
                     # and is valid inter-procedurally
                     rt = widenconst(rt)

base/compiler/compiler.jl

@@ -6,7 +6,7 @@ using Core.Intrinsics, Core.IR
 
 import Core: print, println, show, write, unsafe_write, stdout, stderr,
              _apply_iterate, svec, apply_type, Builtin, IntrinsicFunction,
-             MethodInstance, CodeInstance, MethodMatch
+             MethodInstance, CodeInstance, MethodMatch, PartialOpaque
 
 const getproperty = Core.getfield
 const setproperty! = Core.setfield!

base/compiler/inferenceresult.jl

@@ -45,15 +45,20 @@ function matching_cache_argtypes(linfo::MethodInstance, given_argtypes::Vector)
     return cache_argtypes, overridden_by_const
 end
 
-function matching_cache_argtypes(linfo::MethodInstance, ::Nothing)
-    toplevel = !isa(linfo.def, Method)
-    linfo_argtypes = Any[unwrap_unionall(linfo.specTypes).parameters...]
-    nargs::Int = toplevel ? 0 : linfo.def.nargs
+function most_general_argtypes(method::Union{Method, Nothing}, @nospecialize(specTypes),
+    isva::Bool)
+    toplevel = method === nothing
+    linfo_argtypes = Any[unwrap_unionall(specTypes).parameters...]
+    nargs::Int = toplevel ? 0 : method.nargs
+    if !toplevel && method.is_for_opaque_closure
+        # For opaque closure, the closure environment is processed elsewhere
+        nargs -= 1
+    end
     cache_argtypes = Vector{Any}(undef, nargs)
     # First, if we're dealing with a varargs method, then we set the last element of `args`
     # to the appropriate `Tuple` type or `PartialStruct` instance.
-    if !toplevel && linfo.def.isva
-        if linfo.specTypes == Tuple
+    if !toplevel && isva
+        if specTypes == Tuple
             if nargs > 1
                 linfo_argtypes = svec(Any[Any for i = 1:(nargs - 1)]..., Tuple.parameters[1])
             end
@@ -63,7 +68,7 @@ function matching_cache_argtypes(linfo::MethodInstance, ::Nothing)
             if nargs > linfo_argtypes_length
                 va = linfo_argtypes[linfo_argtypes_length]
                 if isvarargtype(va)
-                    new_va = rewrap_unionall(unconstrain_vararg_length(va), linfo.specTypes)
+                    new_va = rewrap_unionall(unconstrain_vararg_length(va), specTypes)
                     vargtype_elements = Any[new_va]
                     vargtype = Tuple{new_va}
                 else
@@ -74,7 +79,7 @@ function matching_cache_argtypes(linfo::MethodInstance, ::Nothing)
                 vargtype_elements = Any[]
                 for p in linfo_argtypes[nargs:linfo_argtypes_length]
                     p = isvarargtype(p) ? unconstrain_vararg_length(p) : p
-                    push!(vargtype_elements, rewrap(p, linfo.specTypes))
+                    push!(vargtype_elements, rewrap(p, specTypes))
                 end
                 for i in 1:length(vargtype_elements)
                     atyp = vargtype_elements[i]
@@ -115,7 +120,7 @@ function matching_cache_argtypes(linfo::MethodInstance, ::Nothing)
             elseif isconstType(atyp)
                 atyp = Const(atyp.parameters[1])
             else
-                atyp = rewrap(atyp, linfo.specTypes)
+                atyp = rewrap(atyp, specTypes)
             end
             i == n && (lastatype = atyp)
             cache_argtypes[i] = atyp
@@ -126,6 +131,13 @@ function matching_cache_argtypes(linfo::MethodInstance, ::Nothing)
     else
         @assert nargs == 0 "invalid specialization of method" # wrong number of arguments
     end
+    cache_argtypes
+end
+
+function matching_cache_argtypes(linfo::MethodInstance, ::Nothing)
+    mthd = isa(linfo.def, Method) ? linfo.def::Method : nothing
+    cache_argtypes = most_general_argtypes(mthd, linfo.specTypes, isa(mthd, Method) ?
+            mthd.isva : false)
     return cache_argtypes, falses(length(cache_argtypes))
 end
 

base/compiler/tfuncs.jl

@@ -1362,6 +1362,27 @@ add_tfunc(arrayref, 3, INT_INF, arrayref_tfunc, 20)
 add_tfunc(const_arrayref, 3, INT_INF, arrayref_tfunc, 20)
 add_tfunc(arrayset, 4, INT_INF, (@nospecialize(boundscheck), @nospecialize(a), @nospecialize(v), @nospecialize i...)->a, 20)
 
+function _opaque_closure_tfunc(@nospecialize(arg), @nospecialize(isva),
+        @nospecialize(lb), @nospecialize(ub), @nospecialize(source), env::Vector{Any},
+        linfo::MethodInstance)
+
+    argt, argt_exact = instanceof_tfunc(arg)
+    lbt, lb_exact = instanceof_tfunc(lb)
+    if !lb_exact
+        lbt = Union{}
+    end
+
+    ubt, ub_exact = instanceof_tfunc(ub)
+
+    t = (argt_exact ? Core.OpaqueClosure{argt, T} : Core.OpaqueClosure{<:argt, T}) where T
+    t = lbt == ubt ? t{ubt} : (t{T} where lbt <: T <: ubt)
+
+    (isa(source, Const) && isa(source.val, Method)) || return t
+    (isa(isva, Const) && isa(isva.val, Bool)) || return t
+
+    return PartialOpaque(t, tuple_tfunc(env), isva.val, linfo, source.val)
+end
+
 function array_type_undefable(@nospecialize(a))
     if isa(a, Union)
         return array_type_undefable(a.a) || array_type_undefable(a.b)

base/compiler/typeinfer.jl

@@ -289,6 +289,9 @@ function CodeInstance(result::InferenceResult, @nospecialize(inferred_result::An
         if isa(result_type, Const)
             rettype_const = result_type.val
             const_flags = 0x2
+        elseif isa(result_type, PartialOpaque)
+            rettype_const = result_type
+            const_flags = 0x2
         elseif isconstType(result_type)
             rettype_const = result_type.parameters[1]
             const_flags = 0x2
@@ -773,6 +776,8 @@ function typeinf_edge(interp::AbstractInterpreter, method::Method, @nospecialize
         if isdefined(code, :rettype_const)
             if isa(code.rettype_const, Vector{Any}) && !(Vector{Any} <: code.rettype)
                 return PartialStruct(code.rettype, code.rettype_const), mi
+            elseif code.rettype <: Core.OpaqueClosure && isa(code.rettype_const, PartialOpaque)
+                return code.rettype_const, mi
             else
                 return Const(code.rettype_const), mi
             end

base/compiler/typelattice.jl

@@ -185,6 +185,14 @@ function ⊑(@nospecialize(a), @nospecialize(b))
         end
         return false
     end
+    if isa(a, PartialOpaque)
+        if isa(b, PartialOpaque)
+            (a.parent === b.parent && a.source === b.source) || return false
+            return (widenconst(a) <: widenconst(b)) &&
+                ⊑(a.env, b.env)
+        end
+        return widenconst(a) ⊑ b
+    end
     if isa(a, Const)
         if isa(b, Const)
             return a.val === b.val
@@ -223,6 +231,13 @@ function is_lattice_equal(@nospecialize(a), @nospecialize(b))
     isa(b, PartialStruct) && return false
     a isa Const && return false
     b isa Const && return false
+    if isa(a, PartialOpaque)
+        isa(b, PartialOpaque) || return false
+        widenconst(a) == widenconst(b) || return false
+        a.source === b.source || return false
+        a.parent === b.parent || return false
+        return is_lattice_equal(a.env, b.env)
+    end
     return a ⊑ b && b ⊑ a
 end
 
@@ -240,6 +255,7 @@ end
 widenconst(m::MaybeUndef) = widenconst(m.typ)
 widenconst(c::PartialTypeVar) = TypeVar
 widenconst(t::PartialStruct) = t.typ
+widenconst(t::PartialOpaque) = t.typ
 widenconst(t::Type) = t
 widenconst(t::TypeVar) = t
 widenconst(t::Core.TypeofVararg) = t

base/compiler/typelimits.jl

@@ -368,6 +368,15 @@ function tmerge(@nospecialize(typea), @nospecialize(typeb))
        return anyconst ? PartialStruct(widenconst(typea), fields) :
             widenconst(typea)
     end
+    if isa(typea, PartialOpaque) && isa(typeb, PartialOpaque) && widenconst(typea) == widenconst(typeb)
+        if !(typea.source === typeb.source &&
+             typea.isva === typeb.isva &&
+             typea.parent === typeb.parent)
+            return widenconst(typea)
+        end
+        return PartialOpaque(typea.typ, tmerge(typea.env, typeb.env),
+            typea.isva, typea.parent, typea.source)
+    end
     # no special type-inference lattice, join the types
     typea, typeb = widenconst(typea), widenconst(typeb)
     if !isa(typea, Type) || !isa(typeb, Type)

base/compiler/typeutils.jl

@@ -34,6 +34,7 @@ end
 
 function has_nontrivial_const_info(@nospecialize t)
     isa(t, PartialStruct) && return true
+    isa(t, PartialOpaque) && return true
     isa(t, Const) || return false
     val = t.val
     return !isdefined(typeof(val), :instance) && !(isa(val, Type) && hasuniquerep(val))

base/compiler/utilities.jl

@@ -104,11 +104,20 @@ function invoke_api(li::CodeInstance)
     return ccall(:jl_invoke_api, Cint, (Any,), li)
 end
 
-function get_staged(li::MethodInstance)
-    may_invoke_generator(li) || return nothing
+function has_opaque_closure(c::CodeInfo)
+    for i = 1:length(c.code)
+        stmt = c.code[i]
+        (isa(stmt, Expr) && stmt.head === :new_opaque_closure) && return true
+    end
+    return false
+end
+
+function get_staged(mi::MethodInstance)
+    may_invoke_generator(mi) || return nothing
     try
         # user code might throw errors – ignore them
-        return ccall(:jl_code_for_staged, Any, (Any,), li)::CodeInfo
+        ci = ccall(:jl_code_for_staged, Any, (Any,), mi)::CodeInfo
+        return ci
     catch
         return nothing
     end