From fe21eb83f3b362221e0e4592667f15bc01be0f1b Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <aviatesk@gmail.com>
Date: Wed, 21 Aug 2024 22:03:01 +0900
Subject: [PATCH] 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 JuliaLang/julia#55548
---
 base/compiler/abstractinterpretation.jl |  3 ++
 base/compiler/typelattice.jl            | 48 +++++++++++++++++--------
 test/compiler/inference.jl              | 11 ++++++
 3 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/base/compiler/abstractinterpretation.jl b/base/compiler/abstractinterpretation.jl
index 83abfc952bf8e4..966439f19ef7ce 100644
--- a/base/compiler/abstractinterpretation.jl
+++ b/base/compiler/abstractinterpretation.jl
@@ -3510,6 +3510,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, slot_id(changes.var))
             end
             if refinements isa SlotRefinement
                 apply_refinement!(𝕃ᵢ, refinements.slot, refinements.typ, currstate, changes)
diff --git a/base/compiler/typelattice.jl b/base/compiler/typelattice.jl
index d375d61c0bdd81..24f0d822fc75f5 100644
--- a/base/compiler/typelattice.jl
+++ b/base/compiler/typelattice.jl
@@ -745,15 +745,38 @@ 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)
+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)
         newtyp = @noinline widenwrappedslotwrapper(vt.typ)
-        return VarState(newtyp, vt.undef)
+        state[idx] = VarState(newtyp, vt.undef)
+    end
+end
+
+# remove any lattice elements that wrap the reassigned slot object within `ssavaluetypes`
+function invalidate_ssa_slotwrapper!(ssavaluetypes::Vector{Any}, changeid::Int)
+    for idx = 1:length(ssavaluetypes)
+        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)
+        newtyp = @noinline widenwrappedslotwrapper(typ)
+        ssavaluetypes[idx] = newtyp
     end
-    return nothing
 end
 
 function stupdate!(lattice::AbstractLattice, state::VarTable, changes::VarTable)
@@ -778,13 +801,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
diff --git a/test/compiler/inference.jl b/test/compiler/inference.jl
index 75f33a280e245c..655ff4f2cfe58c 100644
--- a/test/compiler/inference.jl
+++ b/test/compiler/inference.jl
@@ -6044,3 +6044,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,))