inference: allow PartialStruct to represent strict undef field (#57541)

This change allows `PartialStruct` to represent structs with strictly
uninitialized fields.

Now the previous `undef::BitVector` field is changed to
`undefs::Vector{Union{Nothing,Bool}}` to encode defined-ness information
of each field.

Also, this lets us fix the length of `typ::PartialStruct`'s `fields` to
always match the number of fields in `typ.typ`. Instead of the current
design where the length of `fields` changes depending on the number of
initialized fields, it seems simpler to have `PartialStruct`s
representing the same `typ` always have the same `fields` length.
So, I've included that refactoring as well.

- fixes the newly detected error in Oscar.jl

Author: aviatesk

Compiler/src/Compiler.jl

@@ -69,7 +69,7 @@ using Base: @_foldable_meta, @_gc_preserve_begin, @_gc_preserve_end, @nospeciali
     structdiff, tls_world_age, unconstrain_vararg_length, unionlen, uniontype_layout,
     uniontypes, unsafe_convert, unwrap_unionall, unwrapva, vect, widen_diagonal,
     _uncompressed_ir, maybe_add_binding_backedge!, datatype_min_ninitialized,
-    partialstruct_undef_length, partialstruct_init_undef
+    partialstruct_init_undefs, fieldcount_noerror
 using Base.Order
 
 import Base: ==, _topmod, append!, convert, copy, copy!, findall, first, get, get!,
@@ -83,10 +83,6 @@ const modifyproperty! = Core.modifyfield!
 const replaceproperty! = Core.replacefield!
 const _DOCS_ALIASING_WARNING = ""
 
-function _getundef(p::PartialStruct)
-    Base.getproperty(p, :undef)
-end
-
 ccall(:jl_set_istopmod, Cvoid, (Any, Bool), Compiler, false)
 
 eval(x) = Core.eval(Compiler, x)

Compiler/src/abstractinterpretation.jl

@@ -2115,7 +2115,7 @@ function abstract_call_builtin(interp::AbstractInterpreter, f::Builtin, (; fargs
                     end
                     return Conditional(a, thentype, elsetype)
                 else
-                    thentype = form_partially_defined_struct(argtype2, argtypes[3])
+                    thentype = form_partially_defined_struct(𝕃ᵢ, argtype2, argtypes[3])
                     if thentype !== nothing
                         elsetype = argtype2
                         if rt === Const(false)
@@ -2133,22 +2133,32 @@ function abstract_call_builtin(interp::AbstractInterpreter, f::Builtin, (; fargs
     return rt
 end
 
-function form_partially_defined_struct(@nospecialize(obj), @nospecialize(name))
+function form_partially_defined_struct(𝕃ᵢ::AbstractLattice, @nospecialize(obj), @nospecialize(name))
     obj isa Const && return nothing # nothing to refine
     name isa Const || return nothing
     objt0 = widenconst(obj)
     objt = unwrap_unionall(objt0)
     objt isa DataType || return nothing
     isabstracttype(objt) && return nothing
+    objt <: Tuple && return nothing
     fldidx = try_compute_fieldidx(objt, name.val)
     fldidx === nothing && return nothing
-    isa(obj, PartialStruct) && return define_field(obj, fldidx)
+    if isa(obj, PartialStruct)
+        _getundefs(obj)[fldidx] === false && return nothing
+        newundefs = copy(_getundefs(obj))
+        newundefs[fldidx] = false
+        return PartialStruct(𝕃ᵢ, obj.typ, newundefs, copy(obj.fields))
+    end
     nminfld = datatype_min_ninitialized(objt)
-    fldidx > nminfld || return nothing
-    undef = partialstruct_init_undef(objt, fldidx; all_defined = false)
-    undef[fldidx] = false
-    fields = Any[fieldtype(objt0, i) for i = 1:fldidx]
-    return PartialStruct(fallback_lattice, objt0, undef, fields)
+    fldidx ≤ nminfld && return nothing
+    fldcnt = fieldcount_noerror(objt)::Int
+    fields = Any[fieldtype(objt0, i) for i = 1:fldcnt]
+    if fields[fldidx] === Union{}
+        return nothing # `Union{}` field never transitions to be defined
+    end
+    undefs = partialstruct_init_undefs(objt, fldcnt)
+    undefs[fldidx] = false
+    return PartialStruct(𝕃ᵢ, objt0, undefs, fields)
 end
 
 function abstract_call_unionall(interp::AbstractInterpreter, argtypes::Vector{Any}, call::CallMeta)
@@ -2663,7 +2673,7 @@ function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
                 # so we try to encode that information with a `PartialStruct`
                 farg2 = ssa_def_slot(fargs[2], sv)
                 if farg2 isa SlotNumber
-                    refined = form_partially_defined_struct(argtypes[2], argtypes[3])
+                    refined = form_partially_defined_struct(𝕃ᵢ, argtypes[2], argtypes[3])
                     if refined !== nothing
                         refinements = SlotRefinement(farg2, refined)
                     end
@@ -3035,7 +3045,6 @@ function abstract_eval_call(interp::AbstractInterpreter, e::Expr, sstate::Statem
     end
 end
 
-
 function abstract_eval_new(interp::AbstractInterpreter, e::Expr, sstate::StatementState,
                            sv::AbsIntState)
     𝕃ᵢ = typeinf_lattice(interp)
@@ -3093,7 +3102,30 @@ function abstract_eval_new(interp::AbstractInterpreter, e::Expr, sstate::Stateme
                 # - any refinement information is available (`anyrefine`), or when
                 # - `nargs` is greater than `n_initialized` derived from the struct type
                 #   information alone
-                rt = PartialStruct(𝕃ᵢ, rt, ats)
+                undefs = Union{Nothing,Bool}[false for _ in 1:nargs]
+                if nargs < fcount # fill in uninitialized fields
+                    for i = (nargs+1):fcount
+                        ft = fieldtype(rt, i)
+                        push!(ats, ft)
+                        if ft === Union{} # `Union{}`-typed field is never initialized
+                            push!(undefs, true)
+                        elseif isconcretetype(ft) && datatype_pointerfree(ft) # this check is probably incomplete
+                            push!(undefs, false)
+                        # TODO If we can implement the query such that it accurately
+                        #      identifies fields that never be `#undef'd, we can make the
+                        #      following improvements:
+                        # elseif is_field_pointerfree(rt, i)
+                        #     push!(undefs, false)
+                        # elseif ismutable && !isconst(rt, i) # can't constrain this field (as it may be modified later)
+                        #     push!(undefs, nothing)
+                        # else
+                        #     push!(undefs, true)
+                        else
+                            push!(undefs, nothing)
+                        end
+                    end
+                end
+                rt = PartialStruct(𝕃ᵢ, rt, undefs, ats)
             end
         else
             rt = refine_partial_type(rt)
@@ -3122,13 +3154,18 @@ function abstract_eval_splatnew(interp::AbstractInterpreter, e::Expr, sstate::St
             end))
             nothrow = isexact
             rt = Const(ccall(:jl_new_structt, Any, (Any, Any), rt, at.val))
-        elseif (isa(at, PartialStruct) && ⊑(𝕃ᵢ, at, Tuple) && n > 0 &&
-                n == length(at.fields::Vector{Any}) && !isvarargtype(at.fields[end]) &&
-                (let t = rt, at = at
-                    all(i::Int -> ⊑(𝕃ᵢ, (at.fields::Vector{Any})[i], fieldtype(t, i)), 1:n)
-                end))
-            nothrow = isexact
-            rt = PartialStruct(𝕃ᵢ, rt, at.fields::Vector{Any})
+        elseif at isa PartialStruct
+            if ⊑(𝕃ᵢ, at, Tuple) && n > 0
+                fields = at.fields
+                if (n == length(fields) && !isvarargtype(fields[end]) &&
+                    (let t = rt
+                        all(i::Int -> ⊑(𝕃ᵢ, fields[i], fieldtype(t, i)), 1:n)
+                    end))
+                    nothrow = isexact
+                    undefs = Union{Nothing,Bool}[false for _ in 1:n]
+                    rt = PartialStruct(𝕃ᵢ, rt, undefs, fields)
+                end
+            end
         end
     else
         rt = refine_partial_type(rt)
@@ -3713,7 +3750,7 @@ end
 @nospecializeinfer function widenreturn_partials(𝕃ᵢ::PartialsLattice, @nospecialize(rt), info::BestguessInfo)
     if isa(rt, PartialStruct)
         fields = copy(rt.fields)
-        anyrefine = refines_definedness_information(rt)
+        anyrefine = n_initialized(rt) > datatype_min_ninitialized(rt.typ)
         𝕃 = typeinf_lattice(info.interp)
         ⊏ = strictpartialorder(𝕃)
         for i in 1:length(fields)
@@ -3725,7 +3762,7 @@ end
             end
             fields[i] = a
         end
-        anyrefine && return PartialStruct(𝕃ᵢ, rt.typ, rt.undef, fields)
+        anyrefine && return PartialStruct(𝕃ᵢ, rt.typ, _getundefs(rt), fields)
     end
     if isa(rt, PartialOpaque)
         return rt # XXX: this case was missed in #39512

Compiler/src/tfuncs.jl

@@ -439,8 +439,9 @@ end
                 end
             elseif isa(arg1, PartialStruct)
                 if !isvarargtype(arg1.fields[end])
-                    if !is_field_maybe_undef(arg1, idx)
-                        return Const(true)
+                    aundefᵢ = _getundefs(arg1)[idx]
+                    if aundefᵢ isa Bool
+                        return Const(!aundefᵢ)
                     end
                 end
             elseif !isvatuple(a1)
@@ -906,31 +907,6 @@ add_tfunc(<:, 2, 2, subtype_tfunc, 10)
     return lty ⊑ Type && rty ⊑ Type
 end
 
-function fieldcount_noerror(@nospecialize t)
-    if t isa UnionAll || t isa Union
-        t = argument_datatype(t)
-        if t === nothing
-            return nothing
-        end
-    elseif t === Union{}
-        return 0
-    end
-    t isa DataType || return nothing
-    if t.name === _NAMEDTUPLE_NAME
-        names, types = t.parameters
-        if names isa Tuple
-            return length(names)
-        end
-        if types isa DataType && types <: Tuple
-            return fieldcount_noerror(types)
-        end
-        return nothing
-    elseif isabstracttype(t) || (t.name === Tuple.name && isvatuple(t))
-        return nothing
-    end
-    return isdefined(t, :types) ? length(t.types) : length(t.name.names)
-end
-
 function try_compute_fieldidx(@nospecialize(typ), @nospecialize(field))
     typ = argument_datatype(typ)
     typ === nothing && return nothing
@@ -1141,8 +1117,12 @@ end
         sty = unwrap_unionall(s)::DataType
         if isa(name, Const)
             nv = _getfield_fieldindex(sty, name)
-            if isa(nv, Int) && !is_field_maybe_undef(s00, nv)
-                return unwrapva(partialstruct_getfield(s00, nv))
+            if isa(nv, Int)
+                if nv < 1
+                    return Bottom
+                elseif nv ≤ length(s00.fields)
+                    return unwrapva(s00.fields[nv])
+                end
             end
         end
         s00 = s
@@ -1437,7 +1417,7 @@ end
             if TF2 === Bottom
                 RT = Bottom
             elseif isconcretetype(RT) && has_nontrivial_extended_info(𝕃ᵢ, TF2) # isconcrete condition required to form a PartialStruct
-                RT = PartialStruct(fallback_lattice, RT, Any[TF, TF2])
+                RT = PartialStruct(fallback_lattice, RT, Union{Nothing,Bool}[false,false], Any[TF, TF2])
             end
             info = ModifyOpInfo(callinfo.info)
             return CallMeta(RT, Any, Effects(), info)
@@ -2015,15 +1995,15 @@ function tuple_tfunc(𝕃::AbstractLattice, argtypes::Vector{Any})
     typ = Tuple{params...}
     # replace a singleton type with its equivalent Const object
     issingletontype(typ) && return Const(typ.instance)
-    return anyinfo ? PartialStruct(𝕃, typ, argtypes) : typ
+    return anyinfo ? PartialStruct(𝕃, typ, partialstruct_init_undefs(typ, argtypes), argtypes) : typ
 end
 
 @nospecs function memorynew_tfunc(𝕃::AbstractLattice, memtype, memlen)
     hasintersect(widenconst(memlen), Int) || return Bottom
     memt = tmeet(𝕃, instanceof_tfunc(memtype, true)[1], GenericMemory)
     memt == Union{} && return memt
     # PartialStruct so that loads of Const `length` get inferred
-    return PartialStruct(𝕃, memt, Any[memlen, Ptr{Nothing}])
+    return PartialStruct(𝕃, memt, Union{Nothing,Bool}[false,false], Any[memlen, Ptr{Nothing}])
 end
 add_tfunc(Core.memorynew, 2, 2, memorynew_tfunc, 10)
 

Compiler/src/typeinfer.jl

@@ -514,7 +514,7 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter, cycleid::
             rettype_const = result_type.parameters[1]
             const_flags = 0x2
         elseif isa(result_type, PartialStruct)
-            rettype_const = (_getundef(result_type), result_type.fields)
+            rettype_const = (_getundefs(result_type), result_type.fields)
             const_flags = 0x2
         elseif isa(result_type, InterConditional)
             rettype_const = result_type
@@ -958,9 +958,9 @@ function cached_return_type(code::CodeInstance)
     rettype_const = code.rettype_const
     # the second subtyping/egal conditions are necessary to distinguish usual cases
     # from rare cases when `Const` wrapped those extended lattice type objects
-    if isa(rettype_const, Tuple{BitVector, Vector{Any}}) && !(Tuple{BitVector, Vector{Any}} <: rettype)
-        undef, fields = rettype_const
-        return PartialStruct(fallback_lattice, rettype, undef, fields)
+    if isa(rettype_const, Tuple{Vector{Union{Nothing,Bool}}, Vector{Any}}) && !(Tuple{Vector{Union{Nothing,Bool}}, Vector{Any}} <: rettype)
+        undefs, fields = rettype_const
+        return PartialStruct(fallback_lattice, rettype, undefs, fields)
     elseif isa(rettype_const, PartialOpaque) && rettype <: Core.OpaqueClosure
         return rettype_const
     elseif isa(rettype_const, InterConditional) && rettype !== InterConditional

Compiler/src/typelattice.jl

@@ -318,15 +318,15 @@ end
         fields = vartyp.fields
         thenfields = thentype === Bottom ? nothing : copy(fields)
         elsefields = elsetype === Bottom ? nothing : copy(fields)
-        undef = copy(_getundef(vartyp))
+        undefs = copy(_getundefs(vartyp))
         if 1 ≤ fldidx ≤ length(fields)
             thenfields === nothing || (thenfields[fldidx] = thentype)
             elsefields === nothing || (elsefields[fldidx] = elsetype)
-            undef[fldidx] = false
+            undefs[fldidx] = false
         end
         return Conditional(slot,
-            thenfields === nothing ? Bottom : PartialStruct(fallback_lattice, vartyp.typ, undef, thenfields),
-            elsefields === nothing ? Bottom : PartialStruct(fallback_lattice, vartyp.typ, undef, elsefields))
+            thenfields === nothing ? Bottom : PartialStruct(fallback_lattice, vartyp.typ, undefs, thenfields),
+            elsefields === nothing ? Bottom : PartialStruct(fallback_lattice, vartyp.typ, undefs, elsefields))
     else
         vartyp_widened = widenconst(vartyp)
         thenfields = thentype === Bottom ? nothing : Any[]
@@ -424,21 +424,15 @@ end
     if isa(a, PartialStruct)
         if isa(b, PartialStruct)
             a.typ <: b.typ || return false
-            if length(a.fields) ≠ length(b.fields)
-                if !(isvarargtype(a.fields[end]) || isvarargtype(b.fields[end]))
-                    length(a.fields) ≥ length(b.fields) || return false
-                else
+            nflds = length(a.fields)
+            nflds == length(b.fields) || return false
+            for i in 1:nflds
+                if !(_getundefs(b)[i] === nothing || _getundefs(a)[i] === _getundefs(b)[i])
                     return false
                 end
-            end
-            na = length(a.fields)
-            nb = length(b.fields)
-            nmax = max(na, nb)
-            for i in 1:nmax
-                is_field_maybe_undef(a, i) ≤ is_field_maybe_undef(b, i) || return false
-                af = partialstruct_getfield(a, i)
-                bf = partialstruct_getfield(b, i)
-                if i == na || i == nb
+                af = a.fields[i]
+                bf = b.fields[i]
+                if i == nflds
                     if isvarargtype(af)
                         # If `af` is vararg, so must bf by the <: above
                         @assert isvarargtype(bf)
@@ -474,10 +468,13 @@ end
             nf = nfields(a.val)
             for i in 1:nf
                 if !isdefined(a.val, i)
-                    is_field_maybe_undef(b, i) || return false # conflicting defined-ness information
+                    _getundefs(b)[i] === false && return false # conflicting defined-ness information
                     continue # since ∀ T Union{} ⊑ T
                 end
                 i > length(b.fields) && break # `a` has more information than `b` that is partially initialized struct
+                if _getundefs(b)[i] === true
+                    return false # conflicting defined-ness information
+                end
                 bfᵢ = b.fields[i]
                 if i == nf
                     bfᵢ = unwrapva(bfᵢ)
@@ -548,7 +545,7 @@ end
     if isa(a, PartialStruct)
         isa(b, PartialStruct) || return false
         length(a.fields) == length(b.fields) || return false
-        _getundef(a) == _getundef(b) || return false
+        _getundefs(a) == _getundefs(b) || return false
         widenconst(a) == widenconst(b) || return false
         a.fields === b.fields && return true # fast path
         for i in 1:length(a.fields)
@@ -756,14 +753,18 @@ end
 # different instances of the compiler that may share the `Core.PartialStruct`
 # type.
 
-function Core.PartialStruct(𝕃::AbstractLattice, @nospecialize(typ), fields::Vector{Any}; all_defined::Bool = true)
-    undef = partialstruct_init_undef(typ, fields; all_defined)
-    return PartialStruct(𝕃, typ, undef, fields)
+# Legacy constructor
+function Core.PartialStruct(𝕃::AbstractLattice, @nospecialize(typ), fields::Vector{Any})
+    return PartialStruct(𝕃, typ, partialstruct_init_undefs(typ, fields), fields)
 end
 
-function Core.PartialStruct(::AbstractLattice, @nospecialize(typ), undef::BitVector, fields::Vector{Any})
+function Core.PartialStruct(::AbstractLattice, @nospecialize(typ), undefs::Vector{Union{Nothing,Bool}}, fields::Vector{Any})
     for i = 1:length(fields)
         assert_nested_slotwrapper(fields[i])
     end
-    return PartialStruct(typ, undef, fields)
+    return PartialStruct(typ, undefs, fields)
 end
+
+# a special getter for `PartialStruct` to achieve better type stability:
+# `(x::PartialStruct).undefs` will be lowered to `getfield(x, :undefs)::Any` otherwise
+_getundefs(p::PartialStruct) = Base.getproperty(p, :undefs)