Inference: propagate struct initialization info on setfield! (JuliaLang#57222)

When a variable has a field set with `setfield!(var, field, value)`,
inference now assumes that this specific field is defined and may for
example constant-propagate `isdefined(var, field)` as `true`.
`PartialStruct`, the lattice element used to encode this information,
still has a few limitations in terms of what it may represent (it cannot
represent mutable structs with non-contiguously defined fields yet),
further work on extending it would increase the impact of this change.

Consider the following function:
```julia
julia> function f()
           a = A(1)
           setfield!(a, :y, 2)
           invokelatest(identity, a)
           isdefined(a, :y) && return 1.0
           a
       end
f (generic function with 1 method)
```

Here is before on `master`:
```julia
julia> @code_typed f()
CodeInfo(
1 ─ %1 = %new(Main.A, 1)::A
│          builtin Main.setfield!(%1, :y, 2)::Int64
│        dynamic builtin (Core._call_latest)(identity, %1)::Any
│   %4 =   builtin Main.isdefined(%1, :y)::Bool
└──      goto JuliaLang#3 if not %4
2 ─      return 1.0
3 ─      return %1
) => Union{Float64, A}
```

And after this PR:
```julia
julia> @code_typed f()
CodeInfo(
1 ─ %1 = %new(Main.A, 1)::A
│          builtin Main.setfield!(%1, :y, 2)::Int64
│        dynamic builtin (Core._call_latest)(identity, %1)::Any
└──      return 1.0
) => Float64
```

---------

Co-authored-by: Cédric Belmant <[email protected]>

Author: serenity4

Compiler/src/abstractinterpretation.jl

@@ -2665,12 +2665,24 @@ function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
         end
         pushfirst!(argtypes, ft)
         refinements = nothing
-        if sv isa InferenceState && f === typeassert
-            # perform very limited back-propagation of invariants after this type assertion
-            if rt !== Bottom && isa(fargs, Vector{Any})
+        if sv isa InferenceState
+            if f === typeassert
+                # perform very limited back-propagation of invariants after this type assertion
+                if rt !== Bottom && isa(fargs, Vector{Any})
+                    farg2 = ssa_def_slot(fargs[2], sv)
+                    if farg2 isa SlotNumber
+                        refinements = SlotRefinement(farg2, rt)
+                    end
+                end
+            elseif f === setfield! && length(argtypes) == 4 && isa(argtypes[3], Const)
+                # from there on we know that the struct field will never be undefined,
+                # so we try to encode that information with a `PartialStruct`
                 farg2 = ssa_def_slot(fargs[2], sv)
                 if farg2 isa SlotNumber
-                    refinements = SlotRefinement(farg2, rt)
+                    refined = form_partially_defined_struct(argtypes[2], argtypes[3])
+                    if refined !== nothing
+                        refinements = SlotRefinement(farg2, refined)
+                    end
                 end
             end
         end

Compiler/test/irpasses.jl

@@ -2042,3 +2042,14 @@ let src = code_typed1(()) do
     end
     @test count(iscall((src, setfield!)), src.code) == 1
 end
+
+# optimize `isdefined` away in the presence of a dominating `setfield!`
+let src = code_typed1(()) do
+        a = Ref{Any}()
+        setfield!(a, :x, 2)
+        invokelatest(identity, a)
+        isdefined(a, :x) && return 1.0
+        a[]
+    end
+    @test count(iscall((src, isdefined)), src.code) == 0
+end