lattice: Thread through lattice argument for getfield_tfunc (#47097)

Like `tuple`, `getfield` needs some lattice awareness to give
the correct answer in the presence of extended lattices. Refactor
to split and thread through the lattice argument through
_getfield_tfunc so external lattices can provide `getfield` tfuncs
for their custom elements.

Author: Keno

base/compiler/abstractinterpretation.jl

@@ -1326,13 +1326,13 @@ function abstract_iteration(interp::AbstractInterpreter, @nospecialize(itft), @n
         if !isa(stateordonet_widened, DataType) || !(stateordonet_widened <: Tuple) || isvatuple(stateordonet_widened) || length(stateordonet_widened.parameters) != 2
             break
         end
-        nstatetype = getfield_tfunc(stateordonet, Const(2))
+        nstatetype = getfield_tfunc(typeinf_lattice(interp), stateordonet, Const(2))
         # If there's no new information in this statetype, don't bother continuing,
         # the iterator won't be finite.
         if ⊑(typeinf_lattice(interp), nstatetype, statetype)
             return Any[Bottom], nothing
         end
-        valtype = getfield_tfunc(stateordonet, Const(1))
+        valtype = getfield_tfunc(typeinf_lattice(interp), stateordonet, Const(1))
         push!(ret, valtype)
         statetype = nstatetype
         call = abstract_call_known(interp, iteratef, ArgInfo(nothing, Any[Const(iteratef), itertype, statetype]), StmtInfo(true), sv)

base/compiler/tfuncs.jl

@@ -854,25 +854,33 @@ function getfield_nothrow(@nospecialize(s00), @nospecialize(name), boundscheck::
     return false
 end
 
-function getfield_tfunc(s00, name, boundscheck_or_order)
-    @nospecialize
+function getfield_tfunc(@specialize(lattice::AbstractLattice), @nospecialize(s00),
+        @nospecialize(name), @nospecialize(boundscheck_or_order))
     t = isvarargtype(boundscheck_or_order) ? unwrapva(boundscheck_or_order) :
         widenconst(boundscheck_or_order)
     hasintersect(t, Symbol) || hasintersect(t, Bool) || return Bottom
-    return getfield_tfunc(s00, name)
+    return getfield_tfunc(lattice, s00, name)
 end
-function getfield_tfunc(s00, name, order, boundscheck)
-    @nospecialize
+function getfield_tfunc(@nospecialize(s00), name, boundscheck_or_order)
+    return getfield_tfunc(fallback_lattice, s00, name, boundscheck_or_order)
+end
+function getfield_tfunc(@specialize(lattice::AbstractLattice), @nospecialize(s00),
+        @nospecialize(name), @nospecialize(order), @nospecialize(boundscheck))
     hasintersect(widenconst(order), Symbol) || return Bottom
     if isvarargtype(boundscheck)
         t = unwrapva(boundscheck)
         hasintersect(t, Symbol) || hasintersect(t, Bool) || return Bottom
     else
         hasintersect(widenconst(boundscheck), Bool) || return Bottom
     end
-    return getfield_tfunc(s00, name)
+    return getfield_tfunc(lattice, s00, name)
 end
-getfield_tfunc(@nospecialize(s00), @nospecialize(name)) = _getfield_tfunc(s00, name, false)
+function getfield_tfunc(@nospecialize(s00), @nospecialize(name), @nospecialize(order), @nospecialize(boundscheck))
+    return getfield_tfunc(fallback_lattice, s00, name, order, boundscheck)
+end
+getfield_tfunc(@nospecialize(s00), @nospecialize(name)) = _getfield_tfunc(fallback_lattice, s00, name, false)
+getfield_tfunc(@specialize(lattice::AbstractLattice), @nospecialize(s00), @nospecialize(name)) = _getfield_tfunc(lattice, s00, name, false)
+
 
 function _getfield_fieldindex(@nospecialize(s), name::Const)
     nv = name.val
@@ -902,10 +910,46 @@ function _getfield_tfunc_const(@nospecialize(sv), name::Const, setfield::Bool)
     return nothing
 end
 
-function _getfield_tfunc(@nospecialize(s00), @nospecialize(name), setfield::Bool)
-    if isa(s00, Conditional)
+function _getfield_tfunc(@specialize(lattice::InferenceLattice), @nospecialize(s00), @nospecialize(name), setfield::Bool)
+    if isa(s00, LimitedAccuracy)
+        # This will error, but it's better than duplicating the error here
+        s00 = widenconst(s00)
+    end
+    return _getfield_tfunc(widenlattice(lattice), s00, name, setfield)
+end
+
+function _getfield_tfunc(@specialize(lattice::OptimizerLattice), @nospecialize(s00), @nospecialize(name), setfield::Bool)
+    # If undef, that's a Union, but that doesn't affect the rt when tmerged
+    # into the unwrapped result.
+    isa(s00, MaybeUndef) && (s00 = s00.typ)
+    return _getfield_tfunc(widenlattice(lattice), s00, name, setfield)
+end
+
+function _getfield_tfunc(@specialize(lattice::AnyConditionalsLattice), @nospecialize(s00), @nospecialize(name), setfield::Bool)
+    if isa(s00, AnyConditional)
         return Bottom # Bool has no fields
-    elseif isa(s00, Const)
+    end
+    return _getfield_tfunc(widenlattice(lattice), s00, name, setfield)
+end
+
+function _getfield_tfunc(@specialize(lattice::PartialsLattice), @nospecialize(s00), @nospecialize(name), setfield::Bool)
+    if isa(s00, PartialStruct)
+        s = widenconst(s00)
+        sty = unwrap_unionall(s)::DataType
+        if isa(name, Const)
+            nv = _getfield_fieldindex(sty, name)
+            if isa(nv, Int) && 1 <= nv <= length(s00.fields)
+                return unwrapva(s00.fields[nv])
+            end
+        end
+        s00 = s
+    end
+
+    return _getfield_tfunc(widenlattice(lattice), s00, name, setfield)
+end
+
+function _getfield_tfunc(lattice::ConstsLattice, @nospecialize(s00), @nospecialize(name), setfield::Bool)
+    if isa(s00, Const)
         sv = s00.val
         if isa(name, Const)
             nv = name.val
@@ -919,30 +963,24 @@ function _getfield_tfunc(@nospecialize(s00), @nospecialize(name), setfield::Bool
             r = _getfield_tfunc_const(sv, name, setfield)
             r !== nothing && return r
         end
-        s = typeof(sv)
-    elseif isa(s00, PartialStruct)
-        s = widenconst(s00)
-        sty = unwrap_unionall(s)::DataType
-        if isa(name, Const)
-            nv = _getfield_fieldindex(sty, name)
-            if isa(nv, Int) && 1 <= nv <= length(s00.fields)
-                return unwrapva(s00.fields[nv])
-            end
-        end
-    else
-        s = unwrap_unionall(s00)
+        s00 = widenconst(s00)
     end
+    return _getfield_tfunc(widenlattice(lattice), s00, name, setfield)
+end
+
+function _getfield_tfunc(lattice::JLTypeLattice, @nospecialize(s00), @nospecialize(name), setfield::Bool)
+    s = unwrap_unionall(s00)
     if isa(s, Union)
-        return tmerge(_getfield_tfunc(rewrap_unionall(s.a, s00), name, setfield),
-                      _getfield_tfunc(rewrap_unionall(s.b, s00), name, setfield))
+        return tmerge(_getfield_tfunc(lattice, rewrap_unionall(s.a, s00), name, setfield),
+                      _getfield_tfunc(lattice, rewrap_unionall(s.b, s00), name, setfield))
     end
     if isType(s)
         if isconstType(s)
             sv = s00.parameters[1]
-	    if isa(name, Const)
+            if isa(name, Const)
                 r = _getfield_tfunc_const(sv, name, setfield)
                 r !== nothing && return r
-	    end
+            end
             s = typeof(sv)
         else
             sv = s.parameters[1]
@@ -982,7 +1020,7 @@ function _getfield_tfunc(@nospecialize(s00), @nospecialize(name), setfield::Bool
         if !(_ts <: Tuple)
             return Any
         end
-        return _getfield_tfunc(_ts, name, setfield)
+        return _getfield_tfunc(lattice, _ts, name, setfield)
     end
     ftypes = datatype_fieldtypes(s)
     nf = length(ftypes)
@@ -1090,7 +1128,7 @@ end
 function setfield!_tfunc(o, f, v)
     @nospecialize
     mutability_errorcheck(o) || return Bottom
-    ft = _getfield_tfunc(o, f, true)
+    ft = _getfield_tfunc(fallback_lattice, o, f, true)
     ft === Bottom && return Bottom
     hasintersect(widenconst(v), widenconst(ft)) || return Bottom
     return v
@@ -1168,7 +1206,7 @@ function abstract_modifyfield!(interp::AbstractInterpreter, argtypes::Vector{Any
         # as well as compute the info for the method matches
         op = unwrapva(argtypes[4])
         v = unwrapva(argtypes[5])
-        TF = getfield_tfunc(o, f)
+        TF = getfield_tfunc(typeinf_lattice(interp), o, f)
         callinfo = abstract_call(interp, ArgInfo(nothing, Any[op, TF, v]), StmtInfo(true), sv, #=max_methods=# 1)
         TF2 = tmeet(callinfo.rt, widenconst(TF))
         if TF2 === Bottom
@@ -2118,6 +2156,9 @@ function builtin_tfunction(interp::AbstractInterpreter, @nospecialize(f), argtyp
         # wrong # of args
         return Bottom
     end
+    if f === getfield
+        return getfield_tfunc(typeinf_lattice(interp), argtypes...)
+    end
     return tf[3](argtypes...)
 end
 

base/compiler/typelattice.jl

@@ -419,7 +419,7 @@ function tmeet(lattice::PartialsLattice, @nospecialize(v), @nospecialize(t::Type
             if isvarargtype(vfi)
                 new_fields[i] = vfi
             else
-                new_fields[i] = tmeet(lattice, vfi, widenconst(getfield_tfunc(t, Const(i))))
+                new_fields[i] = tmeet(lattice, vfi, widenconst(getfield_tfunc(lattice, t, Const(i))))
                 if new_fields[i] === Bottom
                     return Bottom
                 end

base/compiler/typelimits.jl

@@ -321,7 +321,7 @@ function issimplertype(lattice::AbstractLattice, @nospecialize(typea), @nospecia
                 bi = (tni.val::Core.TypeName).wrapper
                 is_lattice_equal(lattice, ai, bi) && continue
             end
-            bi = getfield_tfunc(typeb, Const(i))
+            bi = getfield_tfunc(lattice, typeb, Const(i))
             is_lattice_equal(lattice, ai, bi) && continue
             # It is not enough for ai to be simpler than bi: it must exactly equal
             # (for this, an invariant struct field, by contrast to
@@ -490,8 +490,8 @@ function tmerge(lattice::PartialsLattice, @nospecialize(typea), @nospecialize(ty
             fields = Vector{Any}(undef, type_nfields)
             anyrefine = false
             for i = 1:type_nfields
-                ai = getfield_tfunc(typea, Const(i))
-                bi = getfield_tfunc(typeb, Const(i))
+                ai = getfield_tfunc(lattice, typea, Const(i))
+                bi = getfield_tfunc(lattice, typeb, Const(i))
                 ft = fieldtype(aty, i)
                 if is_lattice_equal(lattice, ai, bi) || is_lattice_equal(lattice, ai, ft)
                     # Since ai===bi, the given type has no restrictions on complexity.