lattice: Properly percolate down tmerge (#46881)

We were skipping the parent lattice of the partials lattice, which may make external lattice implementation nonfunctional.

Co-authored-by: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>

Author: Keno

base/compiler/typelimits.jl

@@ -433,6 +433,8 @@ function tmerge(lattice::ConditionalsLattice, @nospecialize(typea), @nospecializ
         end
         return Bool
     end
+    typea = widenconditional(typea)
+    typeb = widenconditional(typeb)
     return tmerge(widenlattice(lattice), typea, typeb)
 end
 
@@ -471,8 +473,9 @@ end
 
 function tmerge(lattice::PartialsLattice, @nospecialize(typea), @nospecialize(typeb))
     # type-lattice for Const and PartialStruct wrappers
-    if ((isa(typea, PartialStruct) || isa(typea, Const)) &&
-        (isa(typeb, PartialStruct) || isa(typeb, Const)))
+    acp = isa(typea, Const) || isa(typea, PartialStruct)
+    bcp = isa(typeb, Const) || isa(typeb, PartialStruct)
+    if acp && bcp
         aty = widenconst(typea)
         bty = widenconst(typeb)
         if aty === bty
@@ -521,22 +524,40 @@ function tmerge(lattice::PartialsLattice, @nospecialize(typea), @nospecialize(ty
             return anyrefine ? PartialStruct(aty, fields) : aty
         end
     end
+    # Don't widen const here - external AbstractInterpreter might insert lattice
+    # layers between us and `ConstsLattice`.
+    isa(typea, PartialStruct) && (typea = widenconst(typea))
+    isa(typeb, PartialStruct) && (typeb = widenconst(typeb))
 
     # type-lattice for PartialOpaque wrapper
-    if isa(typea, PartialOpaque) && isa(typeb, PartialOpaque) && widenconst(typea) == widenconst(typeb)
-        if !(typea.source === typeb.source &&
-             typea.parent === typeb.parent)
-            return widenconst(typea)
+    apo = isa(typea, PartialOpaque)
+    bpo = isa(typeb, PartialOpaque)
+    if apo && bpo
+        aty = widenconst(typea)
+        bty = widenconst(typeb)
+        if aty == bty
+            if !(typea.source === typeb.source &&
+                typea.parent === typeb.parent)
+                return widenconst(typea)
+            end
+            return PartialOpaque(typea.typ, tmerge(typea.env, typeb.env),
+                typea.parent, typea.source)
         end
-        return PartialOpaque(typea.typ, tmerge(typea.env, typeb.env),
-            typea.parent, typea.source)
+        typea = aty
+        typeb = bty
+    elseif apo
+        typea = widenconst(typea)
+    elseif bpo
+        typeb = widenconst(typeb)
     end
 
-    # no special type-inference lattice, join the types
-    typea, typeb = widenconst(typea), widenconst(typeb)
-    @assert isa(typea, Type); @assert isa(typeb, Type)
+    return tmerge(widenlattice(lattice), typea, typeb)
+end
 
-    return tmerge(JLTypeLattice(), typea, typeb)
+function tmerge(lattice::ConstsLattice, @nospecialize(typea), @nospecialize(typeb))
+    # the equality of the constants can be checked here, but the equivalent check is usually
+    # done by `tmerge_fast_path` at earlier lattice stage
+    return tmerge(widenlattice(lattice), widenconst(typea), widenconst(typeb))
 end
 
 function tmerge(::JLTypeLattice, @nospecialize(typea::Type), @nospecialize(typeb::Type))

test/compiler/AbstractInterpreter.jl

@@ -127,14 +127,14 @@ import .CC:
 struct TaintLattice{PL<:AbstractLattice} <: CC.AbstractLattice
     parent::PL
 end
-CC.widen(𝕃::TaintLattice) = 𝕃.parent
+CC.widenlattice(𝕃::TaintLattice) = 𝕃.parent
 CC.is_valid_lattice(𝕃::TaintLattice, @nospecialize(elm)) =
     is_valid_lattice(widenlattice(𝕃), elem) || isa(elm, Taint)
 
 struct InterTaintLattice{PL<:AbstractLattice} <: CC.AbstractLattice
     parent::PL
 end
-CC.widen(𝕃::InterTaintLattice) = 𝕃.parent
+CC.widenlattice(𝕃::InterTaintLattice) = 𝕃.parent
 CC.is_valid_lattice(𝕃::InterTaintLattice, @nospecialize(elm)) =
     is_valid_lattice(widenlattice(𝕃), elem) || isa(elm, InterTaint)
 
@@ -156,6 +156,9 @@ struct Taint
     end
 end
 Taint(@nospecialize(typ), id::Int) = Taint(typ, push!(BitSet(), id))
+function Base.:(==)(a::Taint, b::Taint)
+    return a.typ == b.typ && a.slots == b.slots
+end
 
 struct InterTaint
     typ
@@ -169,6 +172,9 @@ struct InterTaint
     end
 end
 InterTaint(@nospecialize(typ), id::Int) = InterTaint(typ, push!(BitSet(), id))
+function Base.:(==)(a::InterTaint, b::InterTaint)
+    return a.typ == b.typ && a.slots == b.slots
+end
 
 const AnyTaint = Union{Taint, InterTaint}
 
@@ -229,4 +235,6 @@ function CC.widenreturn(𝕃::InferenceLattice{<:InterTaintLattice}, @nospeciali
     return CC.widenreturn(widenlattice(𝕃), rt, bestguess, nargs, slottypes, changes)
 end
 
+@test CC.tmerge(typeinf_lattice(TaintInterpreter()), Taint(Int, 1), Taint(Int, 2)) == Taint(Int, BitSet(1:2))
+
 # code_typed(ifelse, (Bool, Int, Int); interp=TaintInterpreter())