From e976206563b86734a9e5a07d2c58cafe94e3493c Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>
Date: Tue, 20 Dec 2022 16:33:09 +0900
Subject: [PATCH] allow external lattice to provide its own field-merge
 implementation (#47910)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When merging `PartialStruct`, we are currently merging their fields a
bit aggressively (#44404) in order to accelerate (or rather, guarantee) convergence.
However, when `PartialStruct` wraps external lattice elements, this can
be too aggressive since it does not use `tmerge(𝕃, fields...)` recursively
and thus the external lattice elements are not merged as expected.

This commit adds an additional lattice hook, `tmerge_field`, inside
`tmerge(::PartialsLattice)` so that external lattice implementation
can provide its own field-merge strategies.
---
 base/compiler/abstractlattice.jl | 17 +++++++++++++++++
 base/compiler/typelimits.jl      |  6 +++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/base/compiler/abstractlattice.jl b/base/compiler/abstractlattice.jl
index d98c2b818b649..0302b3f83e373 100644
--- a/base/compiler/abstractlattice.jl
+++ b/base/compiler/abstractlattice.jl
@@ -131,6 +131,23 @@ remaining mostly associative and commutative.
 """
 function tmerge end
 
+"""
+    tmerge_field(𝕃::AbstractLattice, a, b) -> nothing or lattice element
+
+Compute a lattice join of elements `a` and `b` over the lattice `𝕃`,
+where `a` and `b` are fields of `PartialStruct` or `Const`.
+This is an opt-in interface to allow external lattice implementation to provide its own
+field-merge strategy. If it returns `nothing`, `tmerge(::PartialsLattice, ...)`
+will use the default aggressive type merge implementation that does not use `tmerge`
+recursively to reach convergence.
+"""
+function tmerge_field end
+
+function tmerge_field(𝕃::AbstractLattice, @nospecialize(a), @nospecialize(b))
+    return tmerge_field(widenlattice(𝕃), a, b)
+end
+tmerge_field(::JLTypeLattice, @nospecialize(a), @nospecialize(b)) = nothing
+
 """
     ⊑(𝕃::AbstractLattice, a, b)
 
diff --git a/base/compiler/typelimits.jl b/base/compiler/typelimits.jl
index e470c4711110c..2501dd51edf6d 100644
--- a/base/compiler/typelimits.jl
+++ b/base/compiler/typelimits.jl
@@ -515,8 +515,12 @@ function tmerge(lattice::PartialsLattice, @nospecialize(typea), @nospecialize(ty
                     tyi = ai
                 elseif is_lattice_equal(lattice, bi, ft)
                     tyi = bi
+                elseif (tyi′ = tmerge_field(lattice, ai, bi); tyi′ !== nothing)
+                    # allow external lattice implementation to provide a custom field-merge strategy
+                    tyi = tyi′
                 else
-                    # Otherwise choose between using the fieldtype or some other simple merged type.
+                    # Otherwise use the default aggressive field-merge implementation, and
+                    # choose between using the fieldtype or some other simple merged type.
                     # The wrapper type never has restrictions on complexity,
                     # so try to use that to refine the estimated type too.
                     tni = _typename(widenconst(ai))