Move nargs/isva to CodeInfo

This changes the canonical source of truth for va handling from
`Method` to `CodeInfo`. There are multiple goals for this change:

1. This addresses a longstanding complaint about the way that
   CodeInfo-returning generated functions work. Previously, the
   va-ness or not of the returned CodeInfo always had to match that
   of the generator. For Cassette-like transforms that generally have
   one big generator function that is varargs (while then looking up
   lowered code that is not varargs), this could become quite annoying.
   It's possible to workaround, but there is really no good reason to tie
   the two together. As we observed when we implemented OpaqueClosures, the
   vararg-ness of the signature and the `vararg arguments`->`tuple`
   transformation are mostly independent concepts. With this PR, generated
   functions can return CodeInfos with whatever combination of nargs/isva
   is convenient.

2. This change requires clarifying where the va processing boundary is
   in inference. #54076 was already moving in that direction for irinterp,
   and this essentially does much of the same for regular inference. As a
   consequence the constprop cache is now using non-va-cooked signatures,
   which I think is preferable.

3. This further decouples codegen from the presence of a `Method` (which
   is already not assumed, since the code being generated could be a
   toplevel thunk, but some codegen features are only available to things
   that come from Methods). There are a number of upcoming features that
   will require codegen of things that are not quite method specializations
   (See design doc linked in #52797 and things like #50641). This helps
   pave the road for that.

4. I've previously considered expanding the kinds of vararg signatures that
   can be described (see e.g. #53851), which also requires a decoupling of
   the signature and ast notions of vararg. This again lays the groundwork
   for that, although I have no immediate plans to implement this change.

Impact wise, this adds an internal field, which is not too breaking,
but downstream clients vary in how they construct their `CodeInfo`s and
the current way they're doing it will likely be incorrect after this change,
so they will require a small two-line adjustment. We should perhaps consider
pulling out some of the more common patterns into a more stable package, since
interface in most of the last few releases, but that's a separate issue.

Author: Keno

base/compiler/abstractinterpretation.jl

@@ -479,7 +479,12 @@ function conditional_argtype(𝕃ᵢ::AbstractLattice, @nospecialize(rt), @nospe
     if isa(rt, InterConditional) && rt.slot == i
         return rt
     else
-        thentype = elsetype = tmeet(𝕃ᵢ, widenslotwrapper(argtypes[i]), fieldtype(sig, i))
+        argt = widenslotwrapper(argtypes[i])
+        if isvarargtype(argt)
+            @assert fieldcount(sig) == i
+            argt = unwrapva(argt)
+        end
+        thentype = elsetype = tmeet(𝕃ᵢ, argt, fieldtype(sig, i))
         condval = maybe_extract_const_bool(rt)
         condval === true && (elsetype = Bottom)
         condval === false && (thentype = Bottom)
@@ -986,15 +991,12 @@ function maybe_get_const_prop_profitable(interp::AbstractInterpreter,
         # N.B. remarks are emitted within `const_prop_entry_heuristic`
         return nothing
     end
-    nargs::Int = method.nargs
-    method.isva && (nargs -= 1)
-    length(arginfo.argtypes) < nargs && return nothing
     if !const_prop_argument_heuristic(interp, arginfo, sv)
         add_remark!(interp, sv, "[constprop] Disabled by argument and rettype heuristics")
         return nothing
     end
     all_overridden = is_all_overridden(interp, arginfo, sv)
-    if !force && !const_prop_function_heuristic(interp, f, arginfo, nargs, all_overridden, sv)
+    if !force && !const_prop_function_heuristic(interp, f, arginfo, all_overridden, sv)
         add_remark!(interp, sv, "[constprop] Disabled by function heuristic")
         return nothing
     end
@@ -1113,9 +1115,9 @@ function force_const_prop(interp::AbstractInterpreter, @nospecialize(f), method:
 end
 
 function const_prop_function_heuristic(interp::AbstractInterpreter, @nospecialize(f),
-    arginfo::ArgInfo, nargs::Int, all_overridden::Bool, sv::AbsIntState)
+    arginfo::ArgInfo, all_overridden::Bool, sv::AbsIntState)
     argtypes = arginfo.argtypes
-    if nargs > 1
+    if length(argtypes) > 1
         𝕃ᵢ = typeinf_lattice(interp)
         if istopfunction(f, :getindex) || istopfunction(f, :setindex!)
             arrty = argtypes[2]
@@ -1349,20 +1351,6 @@ function matching_cache_argtypes(𝕃::AbstractLattice, mi::MethodInstance,
         end
         given_argtypes[i] = widenslotwrapper(argtype)
     end
-    if condargs !== nothing
-        given_argtypes = let condargs=condargs
-            va_process_argtypes(𝕃, given_argtypes, mi) do isva_given_argtypes::Vector{Any}, last::Int
-                # invalidate `Conditional` imposed on varargs
-                for (slotid, i) in condargs
-                    if slotid ≥ last && (1 ≤ i ≤ length(isva_given_argtypes)) # `Conditional` is already widened to vararg-tuple otherwise
-                        isva_given_argtypes[i] = widenconditional(isva_given_argtypes[i])
-                    end
-                end
-            end
-        end
-    else
-        given_argtypes = va_process_argtypes(𝕃, given_argtypes, mi)
-    end
     return pick_const_args!(𝕃, given_argtypes, cache_argtypes)
 end
 
@@ -1721,7 +1709,7 @@ function abstract_apply(interp::AbstractInterpreter, argtypes::Vector{Any}, si::
     return CallMeta(res, exct, effects, retinfo)
 end
 
-function argtype_by_index(argtypes::Vector{Any}, i::Int)
+function argtype_by_index(argtypes::Vector{Any}, i::Integer)
     n = length(argtypes)
     na = argtypes[n]
     if isvarargtype(na)
@@ -2890,12 +2878,12 @@ end
 struct BestguessInfo{Interp<:AbstractInterpreter}
     interp::Interp
     bestguess
-    nargs::Int
+    nargs::UInt
     slottypes::Vector{Any}
     changes::VarTable
-    function BestguessInfo(interp::Interp, @nospecialize(bestguess), nargs::Int,
+    function BestguessInfo(interp::Interp, @nospecialize(bestguess), nargs::UInt,
         slottypes::Vector{Any}, changes::VarTable) where Interp<:AbstractInterpreter
-        new{Interp}(interp, bestguess, nargs, slottypes, changes)
+        new{Interp}(interp, bestguess, Int(nargs), slottypes, changes)
     end
 end
 
@@ -2970,7 +2958,7 @@ end
         # pick up the first "interesting" slot, convert `rt` to its `Conditional`
         # TODO: ideally we want `Conditional` and `InterConditional` to convey
         # constraints on multiple slots
-        for slot_id = 1:info.nargs
+        for slot_id = 1:Int(info.nargs)
             rt = bool_rt_to_conditional(rt, slot_id, info)
             rt isa InterConditional && break
         end
@@ -2981,6 +2969,9 @@ end
     ⊑ᵢ = ⊑(typeinf_lattice(info.interp))
     old = info.slottypes[slot_id]
     new = widenslotwrapper(info.changes[slot_id].typ) # avoid nested conditional
+    if isvarargtype(old) || isvarargtype(new)
+        return rt
+    end
     if new ⊑ᵢ old && !(old ⊑ᵢ new)
         if isa(rt, Const)
             val = rt.val

base/compiler/inferenceresult.jl

@@ -24,27 +24,53 @@ function matching_cache_argtypes(𝕃::AbstractLattice, mi::MethodInstance,
     for i = 1:length(argtypes)
         given_argtypes[i] = widenslotwrapper(argtypes[i])
     end
-    given_argtypes = va_process_argtypes(𝕃, given_argtypes, mi)
     return pick_const_args!(𝕃, given_argtypes, cache_argtypes)
 end
 
+function pick_const_arg(𝕃::AbstractLattice, @nospecialize(given_argtype), @nospecialize(cache_argtype))
+    if !is_argtype_match(𝕃, given_argtype, cache_argtype, false)
+        # prefer the argtype we were given over the one computed from `mi`
+        if (isa(given_argtype, PartialStruct) && isa(cache_argtype, Type) &&
+            !⊏(𝕃, given_argtype, cache_argtype))
+            # if the type information of this `PartialStruct` is less strict than
+            # declared method signature, narrow it down using `tmeet`
+            given_argtype = tmeet(𝕃, given_argtype, cache_argtype)
+        end
+    else
+        given_argtype = cache_argtype
+    end
+    return given_argtype
+end
+
 function pick_const_args!(𝕃::AbstractLattice, given_argtypes::Vector{Any}, cache_argtypes::Vector{Any})
-    nargtypes = length(given_argtypes)
-    @assert nargtypes == length(cache_argtypes) #= == nargs =# "invalid `given_argtypes` for `mi`"
-    for i = 1:nargtypes
-        given_argtype = given_argtypes[i]
-        cache_argtype = cache_argtypes[i]
-        if !is_argtype_match(𝕃, given_argtype, cache_argtype, false)
-            # prefer the argtype we were given over the one computed from `mi`
-            if (isa(given_argtype, PartialStruct) && isa(cache_argtype, Type) &&
-                !⊏(𝕃, given_argtype, cache_argtype))
-                # if the type information of this `PartialStruct` is less strict than
-                # declared method signature, narrow it down using `tmeet`
-                given_argtypes[i] = tmeet(𝕃, given_argtype, cache_argtype)
-            end
+    if length(given_argtypes) == 0 || length(cache_argtypes) == 0
+        return Any[]
+    end
+    given_va = given_argtypes[end]
+    cache_va = cache_argtypes[end]
+    if isvarargtype(given_va)
+        if isvarargtype(cache_va)
+            # Process the common prefix, then join
+            nprocessargs = max(length(given_argtypes)-1, length(cache_argtypes)-1)
+            resize!(given_argtypes, nprocessargs+1)
+            given_argtypes[end] = Vararg{pick_const_arg(𝕃, unwrapva(given_va), unwrapva(cache_va))}
         else
-            given_argtypes[i] = cache_argtype
+            nprocessargs = length(cache_argtypes)
+            resize!(given_argtypes, nprocessargs)
         end
+    elseif isvarargtype(cache_va)
+        nprocessargs = length(given_argtypes)
+        resize!(given_argtypes, nprocessargs)
+    else
+        @assert length(given_argtypes) == length(cache_argtypes)
+        nprocessargs = length(given_argtypes)
+        resize!(given_argtypes, nprocessargs)
+    end
+    for i = 1:nprocessargs
+        given_argtype = argtype_by_index(given_argtypes, i)
+        cache_argtype = argtype_by_index(cache_argtypes, i)
+        given_argtype = pick_const_arg(𝕃, given_argtype, cache_argtype)
+        given_argtypes[i] = given_argtype
     end
     return given_argtypes
 end
@@ -60,25 +86,33 @@ function is_argtype_match(𝕃::AbstractLattice,
     end
 end
 
-va_process_argtypes(𝕃::AbstractLattice, given_argtypes::Vector{Any}, mi::MethodInstance) =
-    va_process_argtypes(Returns(nothing), 𝕃, given_argtypes, mi)
-function va_process_argtypes(@specialize(va_handler!), 𝕃::AbstractLattice, given_argtypes::Vector{Any}, mi::MethodInstance)
-    def = mi.def::Method
-    isva = def.isva
-    nargs = Int(def.nargs)
-    if isva || isvarargtype(given_argtypes[end])
-        isva_given_argtypes = Vector{Any}(undef, nargs)
+function va_process_argtypes(𝕃::AbstractLattice, given_argtypes::Vector{Any}, nargs::UInt, isva::Bool)
+    if isva || (!isempty(given_argtypes) && isvarargtype(given_argtypes[end]))
+        isva_given_argtypes = Vector{Any}(undef, Int(nargs))
         for i = 1:(nargs-isva)
-            isva_given_argtypes[i] = argtype_by_index(given_argtypes, i)
+            newarg = argtype_by_index(given_argtypes, i)
+            if isva && has_conditional(𝕃) && isa(newarg, Conditional)
+                if newarg.slotid > (nargs-isva)
+                    newarg = widenconditional(newarg)
+                end
+            end
+            isva_given_argtypes[i] = newarg
         end
         if isva
             if length(given_argtypes) < nargs && isvarargtype(given_argtypes[end])
                 last = length(given_argtypes)
             else
                 last = nargs
+                if has_conditional(𝕃)
+                    for i = last:length(given_argtypes)
+                        newarg = given_argtypes[i]
+                        if isa(newarg, Conditional) && newarg.slotid > (nargs-isva)
+                            given_argtypes[i] = widenconditional(newarg)
+                        end
+                    end
+                end
             end
             isva_given_argtypes[nargs] = tuple_tfunc(𝕃, given_argtypes[last:end])
-            va_handler!(isva_given_argtypes, last)
         end
         return isva_given_argtypes
     end
@@ -87,84 +121,44 @@ function va_process_argtypes(@specialize(va_handler!), 𝕃::AbstractLattice, gi
 end
 
 function most_general_argtypes(method::Union{Method,Nothing}, @nospecialize(specTypes))
-    toplevel = method === nothing
-    isva = !toplevel && method.isva
     mi_argtypes = Any[(unwrap_unionall(specTypes)::DataType).parameters...]
-    nargs::Int = toplevel ? 0 : method.nargs
-    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.
-    mi_argtypes_length = length(mi_argtypes)
-    if !toplevel && isva
-        if specTypes::Type == Tuple
-            mi_argtypes = Any[Any for i = 1:nargs]
-            if nargs > 1
-                mi_argtypes[end] = Tuple
-            end
-            vargtype = Tuple
-        else
-            if nargs > mi_argtypes_length
-                va = mi_argtypes[mi_argtypes_length]
-                if isvarargtype(va)
-                    new_va = rewrap_unionall(unconstrain_vararg_length(va), specTypes)
-                    vargtype = Tuple{new_va}
-                else
-                    vargtype = Tuple{}
-                end
-            else
-                vargtype_elements = Any[]
-                for i in nargs:mi_argtypes_length
-                    p = mi_argtypes[i]
-                    p = unwraptv(isvarargtype(p) ? unconstrain_vararg_length(p) : p)
-                    push!(vargtype_elements, elim_free_typevars(rewrap_unionall(p, specTypes)))
-                end
-                for i in 1:length(vargtype_elements)
-                    atyp = vargtype_elements[i]
-                    if issingletontype(atyp)
-                        # replace singleton types with their equivalent Const object
-                        vargtype_elements[i] = Const(atyp.instance)
-                    elseif isconstType(atyp)
-                        vargtype_elements[i] = Const(atyp.parameters[1])
-                    end
-                end
-                vargtype = tuple_tfunc(fallback_lattice, vargtype_elements)
-            end
-        end
-        cache_argtypes[nargs] = vargtype
-        nargs -= 1
+    nargtypes = length(mi_argtypes)
+    nargs = isa(method, Method) ? method.nargs : 0
+    if length(mi_argtypes) < nargs && isvarargtype(mi_argtypes[end])
+        resize!(mi_argtypes, nargs)
     end
     # Now, we propagate type info from `mi_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).
-    if mi_argtypes_length > 0
-        tail_index = nargtypes = min(mi_argtypes_length, nargs)
-        local lastatype
-        for i = 1:nargtypes
-            atyp = mi_argtypes[i]
-            if i == nargtypes && isvarargtype(atyp)
-                atyp = unwrapva(atyp)
-                tail_index -= 1
-            end
-            atyp = unwraptv(atyp)
-            if issingletontype(atyp)
-                # replace singleton types with their equivalent Const object
-                atyp = Const(atyp.instance)
-            elseif isconstType(atyp)
-                atyp = Const(atyp.parameters[1])
-            else
-                atyp = elim_free_typevars(rewrap_unionall(atyp, specTypes))
-            end
-            i == nargtypes && (lastatype = atyp)
-            cache_argtypes[i] = atyp
+    tail_index = min(nargtypes, nargs)
+    local lastatype
+    for i = 1:nargtypes
+        atyp = mi_argtypes[i]
+        wasva = false
+        if i == nargtypes && isvarargtype(atyp)
+            wasva = true
+            atyp = unwrapva(atyp)
         end
-        for i = (tail_index+1):nargs
-            cache_argtypes[i] = lastatype
+        atyp = unwraptv(atyp)
+        if issingletontype(atyp)
+            # replace singleton types with their equivalent Const object
+            atyp = Const(atyp.instance)
+        elseif isconstType(atyp)
+            atyp = Const(atyp.parameters[1])
+        else
+            atyp = elim_free_typevars(rewrap_unionall(atyp, specTypes))
         end
-    else
-        @assert nargs == 0 "invalid specialization of method" # wrong number of arguments
+        mi_argtypes[i] = atyp
+        if wasva
+            lastatype = atyp
+            mi_argtypes[end] = Vararg{atyp}
+        end
+    end
+    for i = (tail_index+1):(nargs-1)
+        mi_argtypes[i] = lastatype
     end
-    return cache_argtypes
+    return mi_argtypes
 end
 
 # eliminate free `TypeVar`s in order to make the life much easier down the road:
@@ -184,7 +178,6 @@ function cache_lookup(𝕃::AbstractLattice, mi::MethodInstance, given_argtypes:
                       cache::Vector{InferenceResult})
     method = mi.def::Method
     nargtypes = length(given_argtypes)
-    @assert nargtypes == Int(method.nargs) "invalid `given_argtypes` for `mi`"
     for cached_result in cache
         cached_result.linfo === mi || @goto next_cache
         cache_argtypes = cached_result.argtypes

base/compiler/inferencestate.jl

@@ -302,6 +302,9 @@ mutable struct InferenceState
         bb_vartables = Union{Nothing,VarTable}[ nothing for i = 1:length(cfg.blocks) ]
         bb_vartable1 = bb_vartables[1] = VarTable(undef, nslots)
         argtypes = result.argtypes
+
+        argtypes = va_process_argtypes(typeinf_lattice(interp), argtypes, src.nargs, src.isva)
+
         nargtypes = length(argtypes)
         for i = 1:nslots
             argtyp = (i > nargtypes) ? Bottom : argtypes[i]
@@ -766,10 +769,9 @@ function print_callstack(sv::InferenceState)
 end
 
 function narguments(sv::InferenceState, include_va::Bool=true)
-    def = sv.linfo.def
-    nargs = length(sv.result.argtypes)
+    nargs = sv.src.nargs
     if !include_va
-        nargs -= isa(def, Method) && def.isva
+        nargs -= sv.src.isva
     end
     return nargs
 end
@@ -831,7 +833,7 @@ function IRInterpretationState(interp::AbstractInterpreter,
     end
     method_info = MethodInfo(src)
     ir = inflate_ir(src, mi)
-    argtypes = va_process_argtypes(optimizer_lattice(interp), argtypes, mi)
+    argtypes = va_process_argtypes(optimizer_lattice(interp), argtypes, src.nargs, src.isva)
     return IRInterpretationState(interp, method_info, ir, mi, argtypes, world,
                                  codeinst.min_world, codeinst.max_world)
 end

base/compiler/optimize.jl

@@ -1264,14 +1264,13 @@ end
 function slot2reg(ir::IRCode, ci::CodeInfo, sv::OptimizationState)
     # need `ci` for the slot metadata, IR for the code
     svdef = sv.linfo.def
-    nargs = isa(svdef, Method) ? Int(svdef.nargs) : 0
     @timeit "domtree 1" domtree = construct_domtree(ir)
-    defuse_insts = scan_slot_def_use(nargs, ci, ir.stmts.stmt)
+    defuse_insts = scan_slot_def_use(ci.nargs, ci, ir.stmts.stmt)
     𝕃ₒ = optimizer_lattice(sv.inlining.interp)
     @timeit "construct_ssa" ir = construct_ssa!(ci, ir, sv, domtree, defuse_insts, 𝕃ₒ) # consumes `ir`
     # NOTE now we have converted `ir` to the SSA form and eliminated slots
     # let's resize `argtypes` now and remove unnecessary types for the eliminated slots
-    resize!(ir.argtypes, nargs)
+    resize!(ir.argtypes, ci.nargs)
     return ir
 end