Extend PartialStruct to represent non-contiguously defined fields

Compiler/src/Compiler.jl

@@ -67,7 +67,7 @@ using Base: @_foldable_meta, @_gc_preserve_begin, @_gc_preserve_end, @nospeciali
     partition_restriction, quoted, rename_unionall, rewrap_unionall, specialize_method,
     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!
+    _uncompressed_ir, maybe_add_binding_backedge!, datatype_min_ninitialized
 using Base.Order
 
 import Base: ==, _topmod, append!, convert, copy, copy!, findall, first, get, get!,

Compiler/src/abstractinterpretation.jl

@@ -2148,23 +2148,14 @@ function form_partially_defined_struct(@nospecialize(obj), @nospecialize(name))
     isabstracttype(objt) && return nothing
     fldidx = try_compute_fieldidx(objt, name.val)
     fldidx === nothing && return nothing
+    isa(obj, PartialStruct) && return define_field(obj, fldidx, fieldtype(objt0, fldidx))
     nminfld = datatype_min_ninitialized(objt)
-    if ismutabletype(objt)
-        # A mutable struct can have non-contiguous undefined fields, but `PartialStruct` cannot
-        # model such a state. So here `PartialStruct` can be used to represent only the
-        # objects where the field following the minimum initialized fields is also defined.
-        if fldidx ≠ nminfld+1
-            # if it is already represented as a `PartialStruct`, we can add one more
-            # `isdefined`-field information on top of those implied by its `fields`
-            if !(obj isa PartialStruct && fldidx == length(obj.fields)+1)
-                return nothing
-            end
-        end
-    else
-        fldidx > nminfld || return nothing
-    end
-    return PartialStruct(fallback_lattice, objt0, Any[obj isa PartialStruct && i≤length(obj.fields) ?
-        obj.fields[i] : fieldtype(objt0,i) for i = 1:fldidx])
+    fldidx > nminfld || return nothing
+    fields = collect(Any, fieldtypes(objt0))
+    nmaxfld = something(datatype_fieldcount(objt), fldidx)
+    undef = trues(nmaxfld)
+    undef[fldidx] = false
+    return PartialStruct(fallback_lattice, objt0, undef, fields)
 end
 
 function abstract_call_unionall(interp::AbstractInterpreter, argtypes::Vector{Any}, call::CallMeta)
@@ -3144,7 +3135,7 @@ function abstract_eval_splatnew(interp::AbstractInterpreter, e::Expr, sstate::St
                     all(i::Int -> ⊑(𝕃ᵢ, (at.fields::Vector{Any})[i], fieldtype(t, i)), 1:n)
                 end))
             nothrow = isexact
-            rt = PartialStruct(𝕃ᵢ, rt, at.fields::Vector{Any})
+            rt = PartialStruct(𝕃ᵢ, rt, at.undef, at.fields::Vector{Any})
         end
     else
         rt = refine_partial_type(rt)
@@ -3725,8 +3716,7 @@ end
 @nospecializeinfer function widenreturn_partials(𝕃ᵢ::PartialsLattice, @nospecialize(rt), info::BestguessInfo)
     if isa(rt, PartialStruct)
         fields = copy(rt.fields)
-        anyrefine = !isvarargtype(rt.fields[end]) &&
-            length(rt.fields) > datatype_min_ninitialized(rt.typ)
+        anyrefine = refines_definedness_information(rt)
         𝕃 = typeinf_lattice(info.interp)
         ⊏ = strictpartialorder(𝕃)
         for i in 1:length(fields)
@@ -3738,7 +3728,7 @@ end
             end
             fields[i] = a
         end
-        anyrefine && return PartialStruct(𝕃ᵢ, rt.typ, fields)
+        anyrefine && return PartialStruct(𝕃ᵢ, rt.typ, rt.undef, fields)
     end
     if isa(rt, PartialOpaque)
         return rt # XXX: this case was missed in #39512

Compiler/src/tfuncs.jl

@@ -439,7 +439,7 @@ end
                 end
             elseif isa(arg1, PartialStruct)
                 if !isvarargtype(arg1.fields[end])
-                    if 1 ≤ idx ≤ length(arg1.fields)
+                    if is_field_defined(arg1, idx)
                         return Const(true)
                     end
                 end
@@ -1141,7 +1141,7 @@ end
         sty = unwrap_unionall(s)::DataType
         if isa(name, Const)
             nv = _getfield_fieldindex(sty, name)
-            if isa(nv, Int) && 1 <= nv <= length(s00.fields)
+            if isa(nv, Int) && is_field_defined(s00, nv)
                 return unwrapva(s00.fields[nv])
             end
         end

Compiler/src/typeinfer.jl

@@ -513,7 +513,7 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter)
             rettype_const = result_type.parameters[1]
             const_flags = 0x2
         elseif isa(result_type, PartialStruct)
-            rettype_const = result_type.fields
+            rettype_const = (result_type.undef, result_type.fields)
             const_flags = 0x2
         elseif isa(result_type, InterConditional)
             rettype_const = result_type
@@ -957,8 +957,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, Vector{Any}) && !(Vector{Any} <: rettype)
-        return PartialStruct(fallback_lattice, rettype, rettype_const)
+    if isa(rettype_const, Tuple{BitVector, Vector{Any}}) && !(Tuple{BitVector, Vector{Any}} <: rettype)
+        undef, fields = rettype_const
+        return PartialStruct(fallback_lattice, rettype, undef, fields)
     elseif isa(rettype_const, PartialOpaque) && rettype <: Core.OpaqueClosure
         return rettype_const
     elseif isa(rettype_const, InterConditional) && rettype !== InterConditional

Compiler/src/typelattice.jl

@@ -318,20 +318,23 @@ end
         fields = vartyp.fields
         thenfields = thentype === Bottom ? nothing : copy(fields)
         elsefields = elsetype === Bottom ? nothing : copy(fields)
+        undef = copy(vartyp.undef)
         for i in 1:length(fields)
             if i == fldidx
                 thenfields === nothing || (thenfields[i] = thentype)
                 elsefields === nothing || (elsefields[i] = elsetype)
+                undef[i] = false
             end
         end
         return Conditional(slot,
-            thenfields === nothing ? Bottom : PartialStruct(fallback_lattice, vartyp.typ, thenfields),
-            elsefields === nothing ? Bottom : PartialStruct(fallback_lattice, vartyp.typ, elsefields))
+            thenfields === nothing ? Bottom : PartialStruct(fallback_lattice, vartyp.typ, undef, thenfields),
+            elsefields === nothing ? Bottom : PartialStruct(fallback_lattice, vartyp.typ, undef, elsefields))
     else
         vartyp_widened = widenconst(vartyp)
         thenfields = thentype === Bottom ? nothing : Any[]
         elsefields = elsetype === Bottom ? nothing : Any[]
-        for i in 1:fieldcount(vartyp_widened)
+        nf = fieldcount(vartyp_widened)
+        for i in 1:nf
             if i == fldidx
                 thenfields === nothing || push!(thenfields, thentype)
                 elsefields === nothing || push!(elsefields, elsetype)
@@ -431,7 +434,9 @@ end
                     return false
                 end
             end
+            length(a.undef) ≥ length(b.undef) || return false
             for i in 1:length(b.fields)
+                !is_field_defined(a, i) && is_field_defined(b, i) && return false
                 af = a.fields[i]
                 bf = b.fields[i]
                 if i == length(b.fields)
@@ -465,12 +470,15 @@ end
             else
                 widea <: wideb || return false
                 # for structs we need to check that `a` has more information than `b` that may be partially initialized
-                n_initialized(a) ≥ length(b.fields) || return false
+                # it may happen that `b` has more information beyond the first undefined field
+                # but in this case we choose `Const` nonetheless.
+                n_initialized(a) ≥ n_initialized(b) || return false
             end
             nf = nfields(a.val)
             for i in 1:nf
                 isdefined(a.val, i) || continue # since ∀ T Union{} ⊑ T
                 i > length(b.fields) && break # `a` has more information than `b` that is partially initialized struct
+                is_field_defined(b, i) || continue # `a` gives a decisive answer as to whether the field is defined or undefined
                 bfᵢ = b.fields[i]
                 if i == nf
                     bfᵢ = unwrapva(bfᵢ)
@@ -541,6 +549,7 @@ end
     if isa(a, PartialStruct)
         isa(b, PartialStruct) || return false
         length(a.fields) == length(b.fields) || return false
+        a.undef == b.undef || return false
         widenconst(a) == widenconst(b) || return false
         a.fields === b.fields && return true # fast path
         for i in 1:length(a.fields)
@@ -751,5 +760,12 @@ function Core.PartialStruct(::AbstractLattice, @nospecialize(typ), fields::Vecto
     for i = 1:length(fields)
         assert_nested_slotwrapper(fields[i])
     end
-    return Core._PartialStruct(typ, fields)
+    return PartialStruct(typ, fields)
+end
+
+function Core.PartialStruct(::AbstractLattice, @nospecialize(typ), undef::BitVector, fields::Vector{Any})
+    for i = 1:length(fields)
+        assert_nested_slotwrapper(fields[i])
+    end
+    return PartialStruct(typ, undef, fields)
 end

Compiler/src/typelimits.jl

@@ -326,6 +326,64 @@ function n_initialized(t::Const)
     return something(findfirst(i::Int->!isdefined(t.val,i), 1:nf), nf+1)-1
 end
 
+function n_initialized(pstruct::PartialStruct)
+    i = findfirst(pstruct.undef)
+    i !== nothing && return i - 1
+    length(pstruct.undef)
+end
+
+maybeundef_fields(pstruct::PartialStruct) = pstruct.undef
+is_field_defined(pstruct::PartialStruct, fi) = 1 ≤ fi ≤ length(pstruct.undef) && !pstruct.undef[fi]
+
+refines_definedness_information(pstruct::PartialStruct) = !isvarargtype(pstruct.fields[end]) && refines_definedness_information(pstruct.typ, pstruct.undef)
+function refines_definedness_information(@nospecialize(typ), undef)
+    nflds = length(undef)
+    something(findfirst(undef), nflds + 1) > datatype_min_ninitialized(typ) + 1
+end
+
+# Returns an iterator over contiguously defined fields
+function defined_fields(pstruct::PartialStruct)
+    i = findfirst(pstruct.undef)
+    i === nothing && return pstruct.fields
+    Any[pstruct.fields[i] for i in 1:(i - 1)]
+end
+
+function define_field(pstruct::PartialStruct, fi, @nospecialize(ft))
+    !pstruct.undef[fi] && return nothing # no new information to be gained
+    undef = copy(pstruct.undef)
+    fields = copy(pstruct.fields)
+    undef[fi] = false
+    fields[fi] = ft
+    PartialStruct(fallback_lattice, pstruct.typ, undef, fields)
+end
+
+# needed while we are missing functions such as broadcasting or ranges
+
+function _bitvector(nt::NTuple)
+    bv = BitVector(undef, length(nt))
+    i = 1
+    while i ≤ length(nt)
+        bv[i] = nt[i]
+        i += 1
+    end
+    bv
+end
+
+#-
+
+maybeundef_fields(t::Const) = undefined_fields(t)
+function undefined_fields(t::Const)
+    nf = nfields(t.val)
+    _bitvector(ntuple(i -> !isdefined(t.val, i), nf))
+end
+
+function maybeundef_fields(x, y)
+    xdef = maybeundef_fields(x)
+    ydef = maybeundef_fields(y)
+    n = min(length(xdef), length(ydef))
+    _bitvector(ntuple(i -> xdef[i] | ydef[i], n))
+end
+
 # A simplified type_more_complex query over the extended lattice
 # (assumes typeb ⊑ typea)
 @nospecializeinfer function issimplertype(𝕃::AbstractLattice, @nospecialize(typea), @nospecialize(typeb))
@@ -334,9 +392,10 @@ end
     if typea isa PartialStruct
         aty = widenconst(typea)
         if typeb isa Const
-            @assert length(typea.fields) ≤ n_initialized(typeb) "typeb ⊑ typea is assumed"
+            @assert n_initialized(typea) ≤ n_initialized(typeb) "typeb ⊑ typea is assumed"
         elseif typeb isa PartialStruct
-            @assert length(typea.fields) ≤ length(typeb.fields) "typeb ⊑ typea is assumed"
+            @assert length(typea.fields) ≤ length(typeb.fields) &&
+                all(!b | a for (a, b) in zip(typea.undef, typeb.undef)) "typeb ⊑ typea is assumed"
         else
             return false
         end
@@ -588,20 +647,17 @@ end
     aty = widenconst(typea)
     bty = widenconst(typeb)
     if aty === bty && !isType(aty)
-        if typea isa PartialStruct
-            if typeb isa PartialStruct
-                nflds = min(length(typea.fields), length(typeb.fields))
-            else
-                nflds = min(length(typea.fields), n_initialized(typeb::Const))
-            end
-        elseif typeb isa PartialStruct
-            nflds = min(n_initialized(typea::Const), length(typeb.fields))
-        else
-            nflds = min(n_initialized(typea::Const), n_initialized(typeb::Const))
-        end
-        nflds == 0 && return nothing
+        typea::Union{PartialStruct, Const}
+        typeb::Union{PartialStruct, Const}
+        maybeundef = maybeundef_fields(typea, typeb)
+        if all(maybeundef)
+            # We could also preserve information about refined field types
+            # (e.g. to better infer non-throwing `getfield` branches).
+            return nothing
+        end
+        nflds = length(maybeundef)
         fields = Vector{Any}(undef, nflds)
-        anyrefine = nflds > datatype_min_ninitialized(aty)
+        anyrefine = refines_definedness_information(aty, maybeundef)
         for i = 1:nflds
             ai = getfield_tfunc(𝕃, typea, Const(i))
             bi = getfield_tfunc(𝕃, typeb, Const(i))
@@ -638,7 +694,13 @@ end
                             ⋤(𝕃, tyi, ft) # just a type-level information, but more precise than the declared type
             end
         end
-        anyrefine && return PartialStruct(𝕃, aty, fields)
+        if isa(typea, PartialStruct) && isa(typeb, PartialStruct) &&
+            isvarargtype(typea.fields[end]) && isvarargtype(typeb.fields[end])
+            # XXX: If it may be more precise than `Vararg` (e.g. `Vararg{T}`),
+            # handle that in the main loop above to get a more accurate type.
+            push!(fields, Vararg)
+        end
+        anyrefine && return PartialStruct(𝕃, aty, maybeundef, fields)
     end
     return nothing
 end

Compiler/src/typeutils.jl

@@ -61,39 +61,6 @@ function isknownlength(t::DataType)
     return isdefined(va, :N) && va.N isa Int
 end
 
-# Compute the minimum number of initialized fields for a particular datatype
-# (therefore also a lower bound on the number of fields)
-function datatype_min_ninitialized(@nospecialize t0)
-    t = unwrap_unionall(t0)
-    t isa DataType || return 0
-    isabstracttype(t) && return 0
-    if t.name === _NAMEDTUPLE_NAME
-        names, types = t.parameters[1], t.parameters[2]
-        if names isa Tuple
-            return length(names)
-        end
-        t = argument_datatype(types)
-        t isa DataType || return 0
-        t.name === Tuple.name || return 0
-    end
-    if t.name === Tuple.name
-        n = length(t.parameters)
-        n == 0 && return 0
-        va = t.parameters[n]
-        if isvarargtype(va)
-            n -= 1
-            if isdefined(va, :N)
-                va = va.N
-                if va isa Int
-                    n += va
-                end
-            end
-        end
-        return n
-    end
-    return length(t.name.names) - t.name.n_uninitialized
-end
-
 has_concrete_subtype(d::DataType) = d.flags & 0x0020 == 0x0020 # n.b. often computed only after setting the type and layout fields
 
 # determine whether x is a valid lattice element