Backport #30385 + dependent changes to 1.0.4

base/compiler/abstractinterpretation.jl

@@ -104,7 +104,7 @@ function abstract_call_gf_by_type(@nospecialize(f), argtypes::Vector{Any}, @nosp
         # if there's a possibility we could constant-propagate a better result
         # (hopefully without doing too much work), try to do that now
         # TODO: it feels like this could be better integrated into abstract_call_method / typeinf_edge
-        const_rettype = abstract_call_method_with_const_args(f, argtypes, applicable[nonbot]::SimpleVector, sv)
+        const_rettype = abstract_call_method_with_const_args(rettype, f, argtypes, applicable[nonbot]::SimpleVector, sv)
         if const_rettype ⊑ rettype
             # use the better result, if it's a refinement of rettype
             rettype = const_rettype
@@ -142,24 +142,24 @@ function abstract_call_gf_by_type(@nospecialize(f), argtypes::Vector{Any}, @nosp
     return rettype
 end
 
-function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector{Any}, match::SimpleVector, sv::InferenceState)
+function abstract_call_method_with_const_args(@nospecialize(rettype), @nospecialize(f), argtypes::Vector{Any}, match::SimpleVector, sv::InferenceState)
     method = match[3]::Method
     nargs::Int = method.nargs
     method.isva && (nargs -= 1)
     length(argtypes) >= nargs || return Any
     haveconst = false
     for a in argtypes
-        a = maybe_widen_conditional(a)
-        if isa(a, Const) && !isdefined(typeof(a.val), :instance) && !(isa(a.val, Type) && issingletontype(a.val))
+        a = widenconditional(a)
+        if has_nontrivial_const_info(a)
             # have new information from argtypes that wasn't available from the signature
-            if isa(a.val, Symbol) || isa(a.val, Type) || (!isa(a.val, String) && isimmutable(a.val))
+            if !isa(a, Const) || (isa(a.val, Symbol) || isa(a.val, Type) || (!isa(a.val, String) && isimmutable(a.val)))
                 # don't consider mutable values or Strings useful constants
                 haveconst = true
                 break
             end
         end
     end
-    haveconst || return Any
+    haveconst || improvable_via_constant_propagation(rettype) || return Any
     sig = match[1]
     sparams = match[2]::SimpleVector
     code = code_for_method(method, sig, sparams, sv.params.world)
@@ -181,7 +181,7 @@ function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector
         if !istopfunction(f, :getproperty) && !istopfunction(f, :setproperty!)
             # in this case, see if all of the arguments are constants
             for a in argtypes
-                a = maybe_widen_conditional(a)
+                a = widenconditional(a)
                 if !isa(a, Const) && !isconstType(a)
                     return Any
                 end
@@ -190,25 +190,7 @@ function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector
     end
     inf_result = cache_lookup(code, argtypes, sv.params.cache)
     if inf_result === nothing
-        inf_result = InferenceResult(code)
-        atypes = get_argtypes(inf_result)
-        if method.isva
-            vargs = argtypes[(nargs + 1):end]
-            for i in 1:length(vargs)
-                a = maybe_widen_conditional(vargs[i])
-                if i > length(inf_result.vargs)
-                    push!(inf_result.vargs, a)
-                elseif a isa Const
-                    inf_result.vargs[i] = a
-                end
-            end
-        end
-        for i in 1:nargs
-            a = maybe_widen_conditional(argtypes[i])
-            if a isa Const
-                atypes[i] = a # inject Const argtypes into inference
-            end
-        end
+        inf_result = InferenceResult(code, argtypes)
         frame = InferenceState(inf_result, #=cache=#false, sv.params)
         frame.limited = true
         frame.parent = sv
@@ -400,6 +382,10 @@ end
 # Union of Tuples of the same length is converted to Tuple of Unions.
 # returns an array of types
 function precise_container_type(@nospecialize(arg), @nospecialize(typ), vtypes::VarTable, sv::InferenceState)
+    if isa(typ, PartialTuple)
+        return typ.fields
+    end
+
     if isa(typ, Const)
         val = typ.val
         if isa(val, SimpleVector) || isa(val, Tuple)
@@ -408,14 +394,9 @@ function precise_container_type(@nospecialize(arg), @nospecialize(typ), vtypes::
     end
 
     arg = ssa_def_expr(arg, sv)
-    if is_specializable_vararg_slot(arg, sv.nargs, sv.result.vargs)
-        return sv.result.vargs
-    end
-
     tti0 = widenconst(typ)
     tti = unwrap_unionall(tti0)
-    if isa(arg, Expr) && arg.head === :call && (abstract_evals_to_constant(arg.args[1], svec, vtypes, sv) ||
-                                                abstract_evals_to_constant(arg.args[1], tuple, vtypes, sv))
+    if isa(arg, Expr) && arg.head === :call && abstract_evals_to_constant(arg.args[1], svec, vtypes, sv)
         aa = arg.args
         result = Any[ abstract_eval(aa[j],vtypes,sv) for j=2:length(aa) ]
         if _any(isvarargtype, result)
@@ -551,7 +532,7 @@ end
 
 function pure_eval_call(@nospecialize(f), argtypes::Vector{Any}, @nospecialize(atype), sv::InferenceState)
     for i = 2:length(argtypes)
-        a = maybe_widen_conditional(argtypes[i])
+        a = widenconditional(argtypes[i])
         if !(isa(a, Const) || isconstType(a))
             return false
         end
@@ -570,7 +551,7 @@ function pure_eval_call(@nospecialize(f), argtypes::Vector{Any}, @nospecialize(a
         return false
     end
 
-    args = Any[ (a = maybe_widen_conditional(argtypes[i]); isa(a, Const) ? a.val : a.parameters[1]) for i in 2:length(argtypes) ]
+    args = Any[ (a = widenconditional(argtypes[i]); isa(a, Const) ? a.val : a.parameters[1]) for i in 2:length(argtypes) ]
     try
         value = Core._apply_pure(f, args)
         # TODO: add some sort of edge(s)
@@ -1108,8 +1089,8 @@ function typeinf_local(frame::InferenceState)
                 end
             elseif hd === :return
                 pc´ = n + 1
-                rt = maybe_widen_conditional(abstract_eval(stmt.args[1], s[pc], frame))
-                if !isa(rt, Const) && !isa(rt, Type)
+                rt = widenconditional(abstract_eval(stmt.args[1], s[pc], frame))
+                if !isa(rt, Const) && !isa(rt, Type) && !isa(rt, PartialTuple)
                     # only propagate information we know we can store
                     # and is valid inter-procedurally
                     rt = widenconst(rt)

base/compiler/inferenceresult.jl

@@ -4,90 +4,122 @@ const EMPTY_VECTOR = Vector{Any}()
 
 mutable struct InferenceResult
     linfo::MethodInstance
-    args::Vector{Any}
-    vargs::Vector{Any} # Memoize vararg type info w/Consts here when calling get_argtypes
-                       # on the InferenceResult, so that the optimizer can use this info
-                       # later during inlining.
+    argtypes::Vector{Any}
+    overridden_by_const::BitVector
     result # ::Type, or InferenceState if WIP
     src #::Union{CodeInfo, OptimizationState, Nothing} # if inferred copy is available
-    function InferenceResult(linfo::MethodInstance)
+    function InferenceResult(linfo::MethodInstance, given_argtypes = nothing)
         if isdefined(linfo, :inferred_const)
             result = Const(linfo.inferred_const)
         else
             result = linfo.rettype
         end
-        return new(linfo, EMPTY_VECTOR, Any[], result, nothing)
+        argtypes, overridden_by_const = matching_cache_argtypes(linfo, given_argtypes)
+        return new(linfo, argtypes, overridden_by_const, result, nothing)
     end
 end
 
-function get_argtypes(result::InferenceResult)
-    result.args === EMPTY_VECTOR || return result.args # already cached
-    argtypes, vargs = get_argtypes(result.linfo)
-    result.args = argtypes
-    if vargs !== nothing
-        result.vargs = vargs
+function is_argtype_match(@nospecialize(given_argtype),
+                          @nospecialize(cache_argtype),
+                          overridden_by_const::Bool)
+    if isa(given_argtype, Const) || isa(given_argtype, PartialTuple)
+        return is_lattice_equal(given_argtype, cache_argtype)
     end
-    return argtypes
+    return !overridden_by_const
 end
 
-function get_argtypes(linfo::MethodInstance)
+# In theory, there could be a `cache` containing a matching `InferenceResult`
+# for the provided `linfo` and `given_argtypes`. The purpose of this function is
+# to return a valid value for `cache_lookup(linfo, argtypes, cache).argtypes`,
+# so that we can construct cache-correct `InferenceResult`s in the first place.
+function matching_cache_argtypes(linfo::MethodInstance, given_argtypes::Vector)
+    @assert isa(linfo.def, Method) # ensure the next line works
+    nargs::Int = linfo.def.nargs
+    @assert length(given_argtypes) >= (nargs - 1)
+    given_argtypes = anymap(widenconditional, given_argtypes)
+    if linfo.def.isva
+        isva_given_argtypes = Vector{Any}(undef, nargs)
+        for i = 1:(nargs - 1)
+            isva_given_argtypes[i] = given_argtypes[i]
+        end
+        isva_given_argtypes[nargs] = tuple_tfunc(given_argtypes[nargs:end])
+        given_argtypes = isva_given_argtypes
+    end
+    cache_argtypes, overridden_by_const = matching_cache_argtypes(linfo, nothing)
+    if nargs === length(given_argtypes)
+        for i in 1:nargs
+            given_argtype = given_argtypes[i]
+            cache_argtype = cache_argtypes[i]
+            if !is_argtype_match(given_argtype, cache_argtype, overridden_by_const[i])
+                # prefer the argtype we were given over the one computed from `linfo`
+                cache_argtypes[i] = given_argtype
+                overridden_by_const[i] = true
+            end
+        end
+    end
+    return cache_argtypes, overridden_by_const
+end
+
+function matching_cache_argtypes(linfo::MethodInstance, ::Nothing)
     toplevel = !isa(linfo.def, Method)
-    atypes::SimpleVector = unwrap_unionall(linfo.specTypes).parameters
+    linfo_argtypes = Any[unwrap_unionall(linfo.specTypes).parameters...]
     nargs::Int = toplevel ? 0 : linfo.def.nargs
-    args = Vector{Any}(undef, nargs)
-    vargs = nothing
+    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 `PartialTuple` instance.
     if !toplevel && linfo.def.isva
         if linfo.specTypes == Tuple
             if nargs > 1
-                atypes = svec(Any[ Any for i = 1:(nargs - 1) ]..., Tuple.parameters[1])
+                linfo_argtypes = svec(Any[Any for i = 1:(nargs - 1)]..., Tuple.parameters[1])
             end
-            vararg_type = Tuple
+            vargtype = Tuple
         else
-            laty = length(atypes)
-            if nargs > laty
-                va = atypes[laty]
+            linfo_argtypes_length = length(linfo_argtypes)
+            if nargs > linfo_argtypes_length
+                va = linfo_argtypes[linfo_argtypes_length]
                 if isvarargtype(va)
                     new_va = rewrap_unionall(unconstrain_vararg_length(va), linfo.specTypes)
-                    vararg_type_vec = Any[new_va]
-                    vararg_type = Tuple{new_va}
+                    vargtype_elements = Any[new_va]
+                    vargtype = Tuple{new_va}
                 else
-                    vararg_type_vec = Any[]
-                    vararg_type = Tuple{}
+                    vargtype_elements = Any[]
+                    vargtype = Tuple{}
                 end
             else
-                vararg_type_vec = Any[]
-                for p in atypes[nargs:laty]
+                vargtype_elements = Any[]
+                for p in linfo_argtypes[nargs:linfo_argtypes_length]
                     p = isvarargtype(p) ? unconstrain_vararg_length(p) : p
-                    push!(vararg_type_vec, rewrap_unionall(p, linfo.specTypes))
+                    push!(vargtype_elements, rewrap(p, linfo.specTypes))
                 end
-                vararg_type = tuple_tfunc(Tuple{vararg_type_vec...})
-                for i in 1:length(vararg_type_vec)
-                    atyp = vararg_type_vec[i]
+                for i in 1:length(vargtype_elements)
+                    atyp = vargtype_elements[i]
                     if isa(atyp, DataType) && isdefined(atyp, :instance)
                         # replace singleton types with their equivalent Const object
-                        vararg_type_vec[i] = Const(atyp.instance)
+                        vargtype_elements[i] = Const(atyp.instance)
                     elseif isconstType(atyp)
-                        vararg_type_vec[i] = Const(atyp.parameters[1])
+                        vargtype_elements[i] = Const(atyp.parameters[1])
                     end
                 end
+                vargtype = tuple_tfunc(vargtype_elements)
             end
-            vargs = vararg_type_vec
         end
-        args[nargs] = vararg_type
+        cache_argtypes[nargs] = vargtype
         nargs -= 1
     end
-    laty = length(atypes)
-    if laty > 0
-        if laty > nargs
-            laty = nargs
-        end
+    # Now, we propagate type info from `linfo_argtypes` into `cache_argtypes`, improving some
+    # type info as we go (where possible). Note that if we're dealing with a varargs method,
+    # we already handled the last element of `cache_argtypes` (and decremented `nargs` so that
+    # we don't overwrite the result of that work here).
+    linfo_argtypes_length = length(linfo_argtypes)
+    if linfo_argtypes_length > 0
+        n = linfo_argtypes_length > nargs ? nargs : linfo_argtypes_length
+        tail_index = n
         local lastatype
-        atail = laty
-        for i = 1:laty
-            atyp = atypes[i]
-            if i == laty && isvarargtype(atyp)
+        for i = 1:n
+            atyp = linfo_argtypes[i]
+            if i == n && isvarargtype(atyp)
                 atyp = unwrapva(atyp)
-                atail -= 1
+                tail_index -= 1
             end
             while isa(atyp, TypeVar)
                 atyp = atyp.ub
@@ -98,42 +130,45 @@ function get_argtypes(linfo::MethodInstance)
             elseif isconstType(atyp)
                 atyp = Const(atyp.parameters[1])
             else
-                atyp = rewrap_unionall(atyp, linfo.specTypes)
+                atyp = rewrap(atyp, linfo.specTypes)
             end
-            i == laty && (lastatype = atyp)
-            args[i] = atyp
+            i == n && (lastatype = atyp)
+            cache_argtypes[i] = atyp
         end
-        for i = (atail + 1):nargs
-            args[i] = lastatype
+        for i = (tail_index + 1):nargs
+            cache_argtypes[i] = lastatype
         end
     else
         @assert nargs == 0 "invalid specialization of method" # wrong number of arguments
     end
-    return args, vargs
+    return cache_argtypes, falses(length(cache_argtypes))
 end
 
-function cache_lookup(code::MethodInstance, argtypes::Vector{Any}, cache::Vector{InferenceResult})
-    method = code.def::Method
+function cache_lookup(linfo::MethodInstance, given_argtypes::Vector{Any}, cache::Vector{InferenceResult})
+    method = linfo.def::Method
     nargs::Int = method.nargs
     method.isva && (nargs -= 1)
-    for cache_code in cache
-        # try to search cache first
-        cache_args = cache_code.args
-        cache_vargs = cache_code.vargs
-        if cache_code.linfo === code && length(argtypes) === (length(cache_vargs) + nargs)
-            cache_match = true
-            for i in 1:length(argtypes)
-                a = maybe_widen_conditional(argtypes[i])
-                ca = i <= nargs ? cache_args[i] : cache_vargs[i - nargs]
-                # verify that all Const argument types match between the call and cache
-                if (isa(a, Const) || isa(ca, Const)) && !(a === ca)
-                    cache_match = false
-                    break
-                end
+    length(given_argtypes) >= nargs || return nothing
+    for cached_result in cache
+        cached_result.linfo === linfo || continue
+        cache_match = true
+        cache_argtypes = cached_result.argtypes
+        cache_overridden_by_const = cached_result.overridden_by_const
+        for i in 1:nargs
+            if !is_argtype_match(given_argtypes[i],
+                                 cache_argtypes[i],
+                                 cache_overridden_by_const[i])
+                cache_match = false
+                break
             end
-            cache_match || continue
-            return cache_code
         end
+        if method.isva && cache_match
+            cache_match = is_argtype_match(tuple_tfunc(given_argtypes[(nargs + 1):end]),
+                                           cache_argtypes[end],
+                                           cache_overridden_by_const[end])
+        end
+        cache_match || continue
+        return cached_result
     end
     return nothing
 end

base/compiler/inferencestate.jl

@@ -59,7 +59,7 @@ mutable struct InferenceState
 
         # initial types
         nslots = length(src.slotnames)
-        argtypes = get_argtypes(result)
+        argtypes = result.argtypes
         nargs = length(argtypes)
         s_argtypes = VarTable(undef, nslots)
         slottypes = Vector{Any}(undef, nslots)