inference: propagate partially initialized mutable structs more (#55533)

Author: aviatesk

base/compiler/abstractinterpretation.jl

@@ -1027,7 +1027,7 @@ collect_const_args(arginfo::ArgInfo, start::Int) = collect_const_args(arginfo.ar
 function collect_const_args(argtypes::Vector{Any}, start::Int)
     return Any[ let a = widenslotwrapper(argtypes[i])
                     isa(a, Const) ? a.val :
-                    isconstType(a) ? (a::DataType).parameters[1] :
+                    isconstType(a) ? a.parameters[1] :
                     (a::DataType).instance
                 end for i = start:length(argtypes) ]
 end
@@ -2063,11 +2063,21 @@ function form_partially_defined_struct(@nospecialize(obj), @nospecialize(name))
     fldidx === nothing && return nothing
     nminfld = datatype_min_ninitialized(objt)
     if ismutabletype(objt)
-        fldidx == nminfld+1 || return nothing
+        # 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(objt0, Any[fieldtype(objt0, i) for i = 1:fldidx])
+    return PartialStruct(objt0, Any[obj isa PartialStruct && i≤length(obj.fields) ?
+        obj.fields[i] : fieldtype(objt0,i) for i = 1:fldidx])
 end
 
 function abstract_call_unionall(interp::AbstractInterpreter, argtypes::Vector{Any}, call::CallMeta)

base/compiler/typeutils.jl

@@ -18,7 +18,7 @@ function hasuniquerep(@nospecialize t)
     iskindtype(typeof(t)) || return true # non-types are always compared by egal in the type system
     isconcretetype(t) && return true # these are also interned and pointer comparable
     if isa(t, DataType) && t.name !== Tuple.name && !isvarargtype(t) # invariant DataTypes
-        return _all(hasuniquerep, t.parameters)
+        return all(hasuniquerep, t.parameters)
     end
     return false
 end

test/compiler/inference.jl

@@ -6033,6 +6033,22 @@ end |> Core.Compiler.is_nothrow
         return x.b
     end
 end |> !Core.Compiler.is_nothrow
+@test Base.infer_effects((PartiallyInitialized2,); optimize=false) do x
+    if isdefined(x, :b)
+        if isdefined(x, :c)
+            return x.c
+        end
+        return x.b
+    end
+    return nothing
+end |> Core.Compiler.is_nothrow
+@test Base.infer_effects((Bool,Int,); optimize=false) do c, b
+    x = c ? PartiallyInitialized1(true) : PartiallyInitialized1(true, b)
+    if isdefined(x, :b)
+        return Val(x.a), x.b
+    end
+    return nothing
+end |> Core.Compiler.is_nothrow
 
 # End to end test case for the partially initialized struct with `PartialStruct`
 @noinline broadcast_noescape1(a) = (broadcast(identity, a); nothing)