inference: invalidate stale slot wrapper types in ssavaluetypes

When updating a state of local variables from an assignment, all stale
slot wrapper types must be invalidated. However, up until now, only
those held in the local variable state were being invalidated. In fact,
those held in `ssavaluetypes` also need to be invalidated, and this
commit addresses that.

Currently all `ssavaluetypes` are iterated over with each assignment to
a local variable, so this could potentially lead to performance issues.
If so it might be necessary to perform invalidation more efficiently
along with flow control.

- fixes #55548

base/compiler/abstractinterpretation.jl

@@ -3357,6 +3357,7 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
 
     states = frame.bb_vartables
     currstate = copy(states[currbb]::VarTable)
+    slotwrapperssas = BitSet()
     while currbb <= nbbs
         delete!(W, currbb)
         bbstart = first(bbs[currbb].stmts)
@@ -3520,6 +3521,9 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
             end
             if changes !== nothing
                 stoverwrite1!(currstate, changes)
+                # widen any slot wrapper types that should be invalidated by this change
+                # just like what's done for `currstate`
+                invalidate_ssa_slotwrapper!(ssavaluetypes, slotwrapperssas, slot_id(changes.var))
             end
             if refinements isa SlotRefinement
                 apply_refinement!(𝕃ᵢ, refinements.slot, refinements.typ, currstate, changes)
@@ -3534,7 +3538,7 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
                 ssavaluetypes[currpc] = Any
                 continue
             end
-            record_ssa_assign!(𝕃ᵢ, currpc, rt, frame)
+            record_ssa_assign!(𝕃ᵢ, currpc, rt, frame, slotwrapperssas)
         end # for currpc in bbstart:bbend
 
         # Case 1: Fallthrough termination
@@ -3612,6 +3616,19 @@ function condition_object_change(currstate::VarTable, condt::Conditional,
     return StateUpdate(condslot, VarState(newcondt, vtype.undef), false)
 end
 
+# remove any lattice elements that wrap the reassigned slot object within `ssavaluetypes`
+function invalidate_ssa_slotwrapper!(ssavaluetypes::Vector{Any}, slotwrapperssas::BitSet, changeid::Int)
+    for idx = slotwrapperssas
+        invalidate_ssa_slotwrapper!(ssavaluetypes, idx, changeid)
+    end
+end
+function invalidate_ssa_slotwrapper!(ssavaluetypes::Vector{Any}, idx::Int, changeid::Int)
+    typ = ssavaluetypes[idx]
+    if should_invalidate(typ, changeid)
+        ssavaluetypes[idx] = @noinline widenwrappedslotwrapper(typ)
+    end
+end
+
 # make as much progress on `frame` as possible (by handling cycles)
 function typeinf_nocycle(interp::AbstractInterpreter, frame::InferenceState)
     typeinf_local(interp, frame)

base/compiler/inferencestate.jl

@@ -734,11 +734,16 @@ end
 
 _topmod(sv::InferenceState) = _topmod(frame_module(sv))
 
-function record_ssa_assign!(𝕃ᵢ::AbstractLattice, ssa_id::Int, @nospecialize(new), frame::InferenceState)
+function record_ssa_assign!(𝕃ᵢ::AbstractLattice, ssa_id::Int, @nospecialize(new),
+                            frame::InferenceState, slotwrapperssas::BitSet)
     ssavaluetypes = frame.ssavaluetypes
     old = ssavaluetypes[ssa_id]
     if old === NOT_FOUND || !is_lattice_equal(𝕃ᵢ, new, old)
         ssavaluetypes[ssa_id] = new
+        wnew = ignorelimited(new)
+        if new isa Conditional || new isa MustAlias
+            push!(slotwrapperssas, ssa_id)
+        end
         W = frame.ip
         for r in frame.ssavalue_uses[ssa_id]
             if was_reached(frame, r)

base/compiler/typelattice.jl

@@ -745,15 +745,23 @@ end
 @nospecializeinfer @inline schanged(lattice::AbstractLattice, @nospecialize(n), @nospecialize(o)) =
     (n !== o) && (o === NOT_FOUND || (n !== NOT_FOUND && !(n.undef <= o.undef && ⊑(lattice, n.typ, o.typ))))
 
-# remove any lattice elements that wrap the reassigned slot object from the vartable
-function invalidate_slotwrapper(vt::VarState, changeid::Int, ignore_conditional::Bool)
-    newtyp = ignorelimited(vt.typ)
-    if (!ignore_conditional && isa(newtyp, Conditional) && newtyp.slot == changeid) ||
-       (isa(newtyp, MustAlias) && newtyp.slot == changeid)
-        newtyp = @noinline widenwrappedslotwrapper(vt.typ)
-        return VarState(newtyp, vt.undef)
+function should_invalidate(@nospecialize(typ), changeid::Int, ignore_conditional::Bool=false)
+    typ = ignorelimited(typ)
+    return ((!ignore_conditional && typ isa Conditional && typ.slot == changeid) ||
+            (typ isa MustAlias && typ.slot == changeid))
+end
+
+# remove any lattice elements that wrap the reassigned slot object within `state`
+function invalidate_slotwrapper!(state::VarTable, changeid::Int, ignore_conditional::Bool)
+    for idx = 1:length(state)
+        invalidate_slotwrapper!(state, idx, changeid, ignore_conditional)
+    end
+end
+function invalidate_slotwrapper!(state::VarTable, idx::Int, changeid::Int, ignore_conditional::Bool)
+    vt = state[idx]
+    if should_invalidate(vt.typ, changeid, ignore_conditional)
+        state[idx] = VarState(@noinline(widenwrappedslotwrapper(vt.typ)), vt.undef)
     end
-    return nothing
 end
 
 function stupdate!(lattice::AbstractLattice, state::VarTable, changes::VarTable)
@@ -778,13 +786,10 @@ end
 
 function stoverwrite1!(state::VarTable, change::StateUpdate)
     changeid = slot_id(change.var)
-    for i = 1:length(state)
-        invalidated = invalidate_slotwrapper(state[i], changeid, change.conditional)
-        if invalidated !== nothing
-            state[i] = invalidated
-        end
-    end
-    # and update the type of it
+    # widen any slot wrapper types that should be invalidated by this change
+    # (unless if this change is made from `Conditional`)
+    invalidate_slotwrapper!(state, changeid, #=ignore_conditional=#change.conditional)
+    # and update the type of the slot
     newtype = change.vtype
     state[changeid] = newtype
     return state

test/compiler/inference.jl

@@ -6060,3 +6060,14 @@ end
 fcondvarargs(a, b, c, d) = isa(d, Int64)
 gcondvarargs(a, x...) = return fcondvarargs(a, x...) ? isa(a, Int64) : !isa(a, Int64)
 @test Core.Compiler.return_type(gcondvarargs, Tuple{Vararg{Any}}) === Bool
+
+# JuliaLang/julia#55548: invalidate stale slot wrapper types in `ssavaluetypes`
+_issue55548_proj1(a, b) = a
+function issue55548(a)
+    a = Base.inferencebarrier(a)::Union{Int64,Float64}
+    if _issue55548_proj1(isa(a, Int64), (a = Base.inferencebarrier(1.0)::Union{Int64,Float64}; true))
+        return a
+    end
+    return 2
+end
+@test Float64 <: Base.infer_return_type(issue55548, (Int,))