lattice: Allow external lattice to have a say in how to widen (#47307)

Currently, when e.g. a PartialStruct is not of interest to the
lattice code, it just calls `widenconst` on it to pass it to the
next lattice layer. This might be insufficient if an intermediate
lattice layer has some other representation that is wider that
the `PartialStruct`, but narrower than the corresponding `widenconst`.
This adds `widenconst(::AbstractLatice, ::Any)` to allow the lattices
to insert custom widening code. By default, it ends up calling
`widenconst`, so there's no functional change in base, but custom
lattices can make use of the extra hook.

I'm not entirely sure that this is what we want the final interface
to look like (I think it probably does too many type checks), but
it works reasonably well and I think is good enough to experiment
with.

base/compiler/abstractlattice.jl

@@ -9,7 +9,8 @@ extensions.
 """
 struct JLTypeLattice <: AbstractLattice; end
 widenlattice(::JLTypeLattice) = error("Type lattice is the least-precise lattice available")
-is_valid_lattice(::JLTypeLattice, @nospecialize(elem)) = isa(elem, Type)
+is_valid_lattice(lattice::JLTypeLattice, @nospecialize(elem)) = is_valid_lattice_norec(lattice, elem)
+is_valid_lattice_norec(::JLTypeLattice, @nospecialize(elem)) = isa(elem, Type)
 
 """
     struct ConstsLattice
@@ -18,8 +19,7 @@ A lattice extending `JLTypeLattice` and adjoining `Const` and `PartialTypeVar`.
 """
 struct ConstsLattice <: AbstractLattice; end
 widenlattice(::ConstsLattice) = JLTypeLattice()
-is_valid_lattice(lattice::ConstsLattice, @nospecialize(elem)) =
-    is_valid_lattice(widenlattice(lattice), elem) || isa(elem, Const) || isa(elem, PartialTypeVar)
+is_valid_lattice_norec(lattice::ConstsLattice, @nospecialize(elem)) = isa(elem, Const) || isa(elem, PartialTypeVar)
 
 """
     struct PartialsLattice{L}
@@ -30,9 +30,7 @@ struct PartialsLattice{L <: AbstractLattice} <: AbstractLattice
     parent::L
 end
 widenlattice(L::PartialsLattice) = L.parent
-is_valid_lattice(lattice::PartialsLattice, @nospecialize(elem)) =
-    is_valid_lattice(widenlattice(lattice), elem) ||
-    isa(elem, PartialStruct) || isa(elem, PartialOpaque)
+is_valid_lattice_norec(lattice::PartialsLattice, @nospecialize(elem)) = isa(elem, PartialStruct) || isa(elem, PartialOpaque)
 
 """
     struct ConditionalsLattice{L}
@@ -43,15 +41,13 @@ struct ConditionalsLattice{L <: AbstractLattice} <: AbstractLattice
     parent::L
 end
 widenlattice(L::ConditionalsLattice) = L.parent
-is_valid_lattice(lattice::ConditionalsLattice, @nospecialize(elem)) =
-    is_valid_lattice(widenlattice(lattice), elem) || isa(elem, Conditional)
+is_valid_lattice_norec(lattice::ConditionalsLattice, @nospecialize(elem)) = isa(elem, Conditional)
 
 struct InterConditionalsLattice{L <: AbstractLattice} <: AbstractLattice
     parent::L
 end
 widenlattice(L::InterConditionalsLattice) = L.parent
-is_valid_lattice(lattice::InterConditionalsLattice, @nospecialize(elem)) =
-    is_valid_lattice(widenlattice(lattice), elem) || isa(elem, InterConditional)
+is_valid_lattice_norec(lattice::InterConditionalsLattice, @nospecialize(elem)) = isa(elem, InterConditional)
 
 const AnyConditionalsLattice{L} = Union{ConditionalsLattice{L}, InterConditionalsLattice{L}}
 const BaseInferenceLattice = typeof(ConditionalsLattice(PartialsLattice(ConstsLattice())))
@@ -67,8 +63,7 @@ struct InferenceLattice{L} <: AbstractLattice
     parent::L
 end
 widenlattice(L::InferenceLattice) = L.parent
-is_valid_lattice(lattice::InferenceLattice, @nospecialize(elem)) =
-    is_valid_lattice(widenlattice(lattice), elem) || isa(elem, LimitedAccuracy)
+is_valid_lattice_norec(lattice::InferenceLattice, @nospecialize(elem)) = isa(elem, LimitedAccuracy)
 
 """
     struct OptimizerLattice
@@ -81,8 +76,7 @@ struct OptimizerLattice{L} <: AbstractLattice
 end
 OptimizerLattice() = OptimizerLattice(BaseInferenceLattice.instance)
 widenlattice(L::OptimizerLattice) = L.parent
-is_valid_lattice(lattice::OptimizerLattice, @nospecialize(elem)) =
-    is_valid_lattice(widenlattice(lattice), elem) || isa(elem, MaybeUndef)
+is_valid_lattice_norec(lattice::OptimizerLattice, @nospecialize(elem)) = isa(elem, MaybeUndef)
 
 """
     tmeet(lattice, a, b::Type)
@@ -174,3 +168,13 @@ tmerge(@nospecialize(a), @nospecialize(b)) = tmerge(fallback_lattice, a, b)
 ⊏(@nospecialize(a), @nospecialize(b)) = ⊏(fallback_lattice, a, b)
 ⋤(@nospecialize(a), @nospecialize(b)) = ⋤(fallback_lattice, a, b)
 is_lattice_equal(@nospecialize(a), @nospecialize(b)) = is_lattice_equal(fallback_lattice, a, b)
+
+is_valid_lattice(lattice::AbstractLattice, @nospecialize(elem)) = is_valid_lattice_norec(lattice, elem) &&
+    is_valid_lattice(widenlattice(lattice), elem)
+
+# Widenlattice with argument
+widenlattice(::JLTypeLattice, @nospecialize(t)) = widenconst(t)
+function widenlattice(lattice::AbstractLattice, @nospecialize(t))
+    is_valid_lattice_norec(lattice, t) && return t
+    widenlattice(widenlattice(lattice), t)
+end

base/compiler/typelimits.jl

@@ -468,6 +468,8 @@ function tmerge(lattice::InterConditionalsLattice, @nospecialize(typea), @nospec
         end
         return Bool
     end
+    typea = widenconditional(typea)
+    typeb = widenconditional(typeb)
     return tmerge(widenlattice(lattice), typea, typeb)
 end
 
@@ -524,10 +526,13 @@ 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))
+    wl = widenlattice(lattice)
+    isa(typea, PartialStruct) && (typea = widenlattice(wl, typea))
+    isa(typeb, PartialStruct) && (typeb = widenlattice(wl, typeb))
 
     # type-lattice for PartialOpaque wrapper
     apo = isa(typea, PartialOpaque)
@@ -540,24 +545,27 @@ function tmerge(lattice::PartialsLattice, @nospecialize(typea), @nospecialize(ty
                 typea.parent === typeb.parent)
                 return widenconst(typea)
             end
-            return PartialOpaque(typea.typ, tmerge(typea.env, typeb.env),
+            return PartialOpaque(typea.typ, tmerge(lattice, typea.env, typeb.env),
                 typea.parent, typea.source)
         end
         typea = aty
         typeb = bty
     elseif apo
-        typea = widenconst(typea)
+        typea = widenlattice(wl, typea)
     elseif bpo
-        typeb = widenconst(typeb)
+        typeb = widenlattice(wl, typeb)
     end
 
-    return tmerge(widenlattice(lattice), typea, typeb)
+    return tmerge(wl, typea, typeb)
 end
 
 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))
+    wl = widenlattice(lattice)
+    (isa(typea, Const) || isa(typea, PartialTypeVar)) && (typea = widenlattice(wl, typea))
+    (isa(typeb, Const) || isa(typeb, PartialTypeVar)) && (typeb = widenlattice(wl, typeb))
+    return tmerge(wl, typea, typeb)
 end
 
 function tmerge(::JLTypeLattice, @nospecialize(typea::Type), @nospecialize(typeb::Type))

test/compiler/AbstractInterpreter.jl

@@ -195,7 +195,7 @@ function CC.tmerge(𝕃::AnyTaintLattice, @nospecialize(typea), @nospecialize(ty
     if isa(typea, T)
         if isa(typeb, T)
             return T(
-                tmerge(widenlattice(𝕃), typea.typ, typeb),
+                tmerge(widenlattice(𝕃), typea.typ, typeb.typ),
                 typea.slots ∪ typeb.slots)
         else
             typea = typea.typ