lattice: Allow external lattice to have a say in how to widen (JuliaL…

…ang#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