Cold-Air-Networks/mirror-llvm-project
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 7

[ADT] Add single-pass merge for ImmutableSet/ImmutableMap (#209807)

Browse source 
The immutable-set/map dataflow joins merged two containers by inserting
the
elements of one into the other one at a time, costing O(|B| * log|A|)
and
re-copying shared spine nodes on every insert.

Add a single-pass, structure-sharing tree merge on ImutAVLFactory
(mergeTrees), exposed as ImmutableSet::Factory::unionSets and
ImmutableMap::Factory::mergeWith. It recurses over the larger operand
and
splits the smaller at each key, returning non-overlapping subtrees
unchanged,
so each spine node is copied at most once: O(|B| * log(|A|/|B| + 1)).

Two flags tune it for the different joins:
* KeepUnmatched - keep keys unique to one side (set union / a lattice
join
with an identity) vs. pass them through the combiner (a symmetric join,
   e.g. liveness Must->Maybe).
* SkipShared - return a pointer-identical subtree unchanged in O(1).
Valid
for an idempotent merge (Combine(a, a) == a), which holds for set union
and
every lattice join. Since dataflow states are path-copied from one
another,
a join's operands usually share most subtrees by pointer, so a join that
   touches only a few keys becomes nearly O(diff).

Wire it into the Clang lifetime-safety analysis (set and map joins) and
the
LiveVariables merge, and add ImmutableMergeBM to track it.

Speedups on the lifetime-safety dataflow benchmark (LoanPropagation
phase):
      pointer cycle in a loop  ~2.6x   (1.59 s  -> 0.61 s)
      CFG merges               ~190x   (1.83 s  -> 9.5 ms)
switch fan-out ~1.9x (LiveOrigins; disjoint states, no skip)
LiveVariables is ~2-3x faster on liveness-heavy inputs and neutral on
typical
code (the join is a small fraction of real-code analysis time).

ImmutableMergeBM confirms the merge beats the add loop for small and
near-identical operands, except for independently-built sets differing
by
exactly one element; when one operand is derived from the other -- the
real
dataflow case -- the pointer-skip makes even that a large win. New
randomized
ADT stress tests check the merge against std::set/std::map oracles for
both
canonicalizing and non-canonicalizing factories.

Assisted by: Claude Opus 4.8

Co-authored-by: Gabor Horvath <gaborh@apple.com>
This commit is contained in:
Gábor Horváth 2026-07-21 12:33:28 +01:00 committed by GitHub
parent df4bb1f310
commit ad3f09c7de
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 735 additions and 67 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
clang/include/clang/Analysis/Analyses/LifetimeSafety/Utils.h
View file

@ -47,13 +47,7 @@ using MapTy = llvm::ImmutableMap<KeyT, ValT, llvm::ImutKeyValueInfo<KeyT, ValT>,
/// Computes the union of two ImmutableSets.
template <typename T>
SetTy<T> join(SetTy<T> A, SetTy<T> B, typename SetTy<T>::Factory &F) {
if (A.getRootWithoutRetain() == B.getRootWithoutRetain())
return A;
if (A.getHeight() < B.getHeight())
std::swap(A, B);
for (const auto &E : B)
A = F.add(A, E);
return A;
return F.unionSets(A, B);
}
/// Describes the strategy for joining two `ImmutableMap` instances, primarily
@ -72,30 +66,36 @@ enum class JoinKind {
Asymmetric,
};
/// Computes the key-wise union of two ImmutableMaps.
// TODO(opt): This key-wise join is a performance bottleneck. A more
// efficient merge could be implemented using a Patricia Trie or HAMT
// instead of the current AVL-tree-based ImmutableMap.
/// Computes the key-wise union of two ImmutableMaps in a single traversal
/// (see ImmutableMap::Factory::mergeWith), sharing subtrees the two maps do
/// not overlap. This assumes -- as the swap below already does -- that
/// JoinValues is commutative with a left identity, which holds for the
/// lifetime lattices.
template <typename KeyT, typename ValT, typename Joiner>
MapTy<KeyT, ValT> join(const MapTy<KeyT, ValT> &A, const MapTy<KeyT, ValT> &B,
MapTy<KeyT, ValT> join(MapTy<KeyT, ValT> A, MapTy<KeyT, ValT> B,
typename MapTy<KeyT, ValT>::Factory &F,
Joiner JoinValues, JoinKind Kind) {
if (A.getRootWithoutRetain() == B.getRootWithoutRetain())
return A;
// Drive the merge with the taller map so the shorter one is the one split.
if (A.getHeight() < B.getHeight())
return join(B, A, F, JoinValues, Kind);
std::swap(A, B);
// For each element in B, join it with the corresponding element in A
// (or with an empty value if it doesn't exist in A).
MapTy<KeyT, ValT> Res = A;
for (const auto &[Key, ValB] : B)
Res = F.add(Res, Key, JoinValues(A.lookup(Key), &ValB));
if (Kind == JoinKind::Symmetric) {
for (const auto &[Key, ValA] : A)
if (!B.contains(Key))
Res = F.add(Res, Key, JoinValues(&ValA, nullptr));
}
return Res;
using ValueTy = typename MapTy<KeyT, ValT>::value_type;
auto Combine = [&JoinValues](const ValueTy *AElem,
const ValueTy *BElem) -> std::pair<KeyT, ValT> {
const KeyT &Key = AElem ? AElem->first : BElem->first;
return std::pair<KeyT, ValT>(Key,
JoinValues(AElem ? &AElem->second : nullptr,
BElem ? &BElem->second : nullptr));
};
// Asymmetric keeps keys unique to either map as-is (valid because JoinValues
// has a left identity); symmetric passes unmatched keys through JoinValues.
// The lifetime joins are idempotent lattice joins, so pointer-identical
// subtrees (common once one state is derived from the other) can be shared.
return F.mergeWith(A, B, Combine,
/*KeepUnmatched=*/Kind == JoinKind::Asymmetric,
/*SkipShared=*/true);
}
} // namespace clang::lifetimes::internal::utils

50
clang/lib/Analysis/LiveVariables.cpp
View file

@ -30,9 +30,6 @@ namespace {
class LiveVariablesImpl {
public:
template <typename T> using SetTy = LiveVariables::SetTy<T>;
template <typename T>
using SetRefTy = llvm::ImmutableSetRef<T, llvm::ImutContainerInfo<T>,
/*Canonicalize=*/false>;
AnalysisDeclContext &analysisContext;
SetTy<const Expr *>::Factory ESetFact;
@ -89,51 +86,18 @@ bool LiveVariables::LivenessValues::isLive(const VarDecl *D) const {
return liveDecls.contains(D);
}
namespace {
template <typename SET> SET mergeSets(SET A, SET B) {
if (A.getRootWithoutRetain() == B.getRootWithoutRetain())
return A;
if (A.getHeight() < B.getHeight())
std::swap(A, B);
for (const auto *Elem : B) {
A = A.add(Elem);
}
return A;
}
} // namespace
void LiveVariables::Observer::anchor() { }
LiveVariables::LivenessValues
LiveVariablesImpl::merge(LiveVariables::LivenessValues valsA,
LiveVariables::LivenessValues valsB) {
SetRefTy<const Expr *> SSetRefA(valsA.liveExprs.getRootWithoutRetain(),
ESetFact.getTreeFactory()),
SSetRefB(valsB.liveExprs.getRootWithoutRetain(),
ESetFact.getTreeFactory());
SetRefTy<const VarDecl *> DSetRefA(valsA.liveDecls.getRootWithoutRetain(),
DSetFact.getTreeFactory()),
DSetRefB(valsB.liveDecls.getRootWithoutRetain(),
DSetFact.getTreeFactory());
SetRefTy<const BindingDecl *> BSetRefA(
valsA.liveBindings.getRootWithoutRetain(), BSetFact.getTreeFactory()),
BSetRefB(valsB.liveBindings.getRootWithoutRetain(),
BSetFact.getTreeFactory());
SSetRefA = mergeSets(SSetRefA, SSetRefB);
DSetRefA = mergeSets(DSetRefA, DSetRefB);
BSetRefA = mergeSets(BSetRefA, BSetRefB);
// Finalize the merged builder refs into immutable sets. These are not
// canonicalized; LivenessValues::operator== compares them structurally.
return LiveVariables::LivenessValues(SSetRefA.asImmutableSet(),
DSetRefA.asImmutableSet(),
BSetRefA.asImmutableSet());
// Liveness at a merge point is the union of the successors' live sets. These
// sets are not canonicalized; LivenessValues::operator== compares them
// structurally.
return LiveVariables::LivenessValues(
ESetFact.unionSets(valsA.liveExprs, valsB.liveExprs),
DSetFact.unionSets(valsA.liveDecls, valsB.liveDecls),
BSetFact.unionSets(valsA.liveBindings, valsB.liveBindings));
}
bool LiveVariables::LivenessValues::operator==(const LivenessValues &V) const {

1
llvm/benchmarks/CMakeLists.txt
View file

@ -17,6 +17,7 @@ add_benchmark(DWARFVerifierBM DWARFVerifierBM.cpp PARTIAL_SOURCES_INTENDED)
add_benchmark(PointerUnionBM PointerUnionBM.cpp PARTIAL_SOURCES_INTENDED)
add_benchmark(ImmutableSetIteratorBM ImmutableSetIteratorBM.cpp PARTIAL_SOURCES_INTENDED)
add_benchmark(ImmutableSetBuildBM ImmutableSetBuildBM.cpp PARTIAL_SOURCES_INTENDED)
add_benchmark(ImmutableMergeBM ImmutableMergeBM.cpp PARTIAL_SOURCES_INTENDED)
add_benchmark(RuntimeLibcallsBench RuntimeLibcalls.cpp PARTIAL_SOURCES_INTENDED)
add_benchmark(writeToOutputInParallelBench writeToOutputInParallel.cpp PARTIAL_SOURCES_INTENDED)

379
llvm/benchmarks/ImmutableMergeBM.cpp Normal file
View file

@ -0,0 +1,379 @@
//===- ImmutableMergeBM.cpp - Benchmark ImmutableSet/Map merges -----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Compares the single-pass, structure-sharing merge (ImmutableSet::unionSets /
// ImmutableMap::mergeWith) against the per-element add loop it replaced, on the
// operand shapes that dominate dataflow joins:
//
// * Small - both operands tiny (constant-factor sensitive).
// * NearId - both large but differing by only a few keys (the fixpoint
// case, where most subtrees are shared).
//
// It covers a plain set, a scalar-valued map, and the map-of-sets shape used by
// the Clang lifetime analysis (ImmutableMap<Origin, ImmutableSet<Loan>>, whose
// join key-wise unions the inner loan sets). The point of these benchmarks is
// to confirm the bulk merge is not slower than the add loop for small or nearly
// identical inputs. Run before/after and compare (e.g. llvm/utils/compare.py).
//
//===----------------------------------------------------------------------===//
#include "benchmark/benchmark.h"
#include "llvm/ADT/ImmutableMap.h"
#include "llvm/ADT/ImmutableSet.h"
#include <algorithm>
#include <memory>
#include <numeric>
#include <vector>
using namespace llvm;
namespace {
using IntSet =
ImmutableSet<int, ImutContainerInfo<int>, /*Canonicalize=*/false>;
using IntMap =
ImmutableMap<int, int, ImutKeyValueInfo<int, int>, /*Canonicalize=*/false>;
// The lifetime analysis' OriginLoanMap shape: a map whose values are sets.
using SetMap = ImmutableMap<int, IntSet, ImutKeyValueInfo<int, IntSet>,
/*Canonicalize=*/false>;
//===----------------------------------------------------------------------===//
// The per-element merges the bulk path replaced.
//===----------------------------------------------------------------------===//
static IntSet setUnionAdd(IntSet::Factory &F, IntSet A, IntSet B) {
if (A.getRootWithoutRetain() == B.getRootWithoutRetain() || B.isEmpty())
return A;
if (A.isEmpty())
return B;
if (A.getHeight() < B.getHeight())
std::swap(A, B);
for (int E : B)
A = F.add(A, E);
return A;
}
static IntMap mapUnionAdd(IntMap::Factory &F, IntMap A, IntMap B) {
if (A.getRootWithoutRetain() == B.getRootWithoutRetain())
return A;
if (A.getHeight() < B.getHeight())
std::swap(A, B);
IntMap Res = A;
for (IntMap::iterator I = B.begin(), E = B.end(); I != E; ++I) {
const int *AV = A.lookup(I.getKey());
Res = F.add(Res, I.getKey(), AV ? std::max(*AV, I.getData()) : I.getData());
}
return Res;
}
static SetMap setMapUnionAdd(SetMap::Factory &MF, IntSet::Factory &SF, SetMap A,
SetMap B) {
if (A.getRootWithoutRetain() == B.getRootWithoutRetain())
return A;
if (A.getHeight() < B.getHeight())
std::swap(A, B);
SetMap Res = A;
for (SetMap::iterator I = B.begin(), E = B.end(); I != E; ++I) {
const IntSet *AV = A.lookup(I.getKey());
IntSet Merged = AV ? setUnionAdd(SF, *AV, I.getData()) : I.getData();
Res = MF.add(Res, I.getKey(), Merged);
}
return Res;
}
//===----------------------------------------------------------------------===//
// The bulk merges (mirroring the caller-side short-circuit/swap the analyses
// perform around them, so both sides are compared on equal footing).
//===----------------------------------------------------------------------===//
static IntSet setUnionNew(IntSet::Factory &F, IntSet A, IntSet B) {
return F.unionSets(A, B);
}
static IntMap mapUnionNew(IntMap::Factory &F, IntMap A, IntMap B) {
if (A.getRootWithoutRetain() == B.getRootWithoutRetain())
return A;
if (A.getHeight() < B.getHeight())
std::swap(A, B);
auto Combine = [](const IntMap::value_type *AE,
const IntMap::value_type *BE) -> std::pair<int, int> {
int Key = AE ? AE->first : BE->first;
return {Key, std::max(AE ? AE->second : BE->second,
BE ? BE->second : AE->second)};
};
return F.mergeWith(A, B, Combine, /*KeepUnmatched=*/true);
}
static SetMap setMapUnionNew(SetMap::Factory &MF, IntSet::Factory &SF, SetMap A,
SetMap B) {
if (A.getRootWithoutRetain() == B.getRootWithoutRetain())
return A;
if (A.getHeight() < B.getHeight())
std::swap(A, B);
auto Combine = [&SF](const SetMap::value_type *AE,
const SetMap::value_type *BE) -> std::pair<int, IntSet> {
int Key = AE ? AE->first : BE->first;
IntSet AS = AE ? AE->second : SF.getEmptySet();
IntSet BS = BE ? BE->second : SF.getEmptySet();
return {Key, SF.unionSets(AS, BS)};
};
return MF.mergeWith(A, B, Combine, /*KeepUnmatched=*/true);
}
//===----------------------------------------------------------------------===//
// Key generators. Small: two tiny half-overlapping sets. NearId: two size-N
// sets sharing all but D keys (D unique to each side).
//===----------------------------------------------------------------------===//
static void makeKeys(bool NearId, int N, int D, std::vector<int> &KA,
std::vector<int> &KB) {
KA.resize(N);
std::iota(KA.begin(), KA.end(), 0);
KB = KA;
if (NearId) {
// Replace the first D keys of B with fresh keys, so A and B share N-D keys
// and each has D unique ones.
for (int I = 0; I < D && I < N; ++I)
KB[I] = N + I;
} else {
// Half-overlap: shift B up by N/2.
for (int &V : KB)
V += N / 2;
}
}
// Bound allocator growth: results accumulate in the factory, so periodically
// rebuild it (excluded from timing) rather than let it grow unboundedly.
static constexpr int ResetEvery = 4096;
static IntSet buildSet(IntSet::Factory &F, const std::vector<int> &Keys) {
IntSet S = F.getEmptySet();
for (int V : Keys)
S = F.add(S, V);
return S;
}
//===----------------------------------------------------------------------===//
// Benchmarks.
//===----------------------------------------------------------------------===//
static void BM_Set(benchmark::State &State, bool UseNew, bool NearId, int N,
int D) {
std::vector<int> KA, KB;
makeKeys(NearId, N, D, KA, KB);
std::unique_ptr<IntSet::Factory> F;
IntSet A{nullptr}, B{nullptr};
auto reset = [&] {
A = IntSet(nullptr);
B = IntSet(nullptr);
F = std::make_unique<IntSet::Factory>();
A = buildSet(*F, KA);
B = buildSet(*F, KB);
};
reset();
int Since = 0;
for (auto _ : State) {
// Rebuild before creating this iteration's result, so no live result
// references the factory being torn down.
if (++Since == ResetEvery) {
State.PauseTiming();
reset();
Since = 0;
State.ResumeTiming();
}
IntSet U = UseNew ? setUnionNew(*F, A, B) : setUnionAdd(*F, A, B);
benchmark::DoNotOptimize(U.getRootWithoutRetain());
}
}
// Like BM_Set, but B is derived from A by adding D fresh keys, so B shares all
// of A's untouched subtrees by pointer -- the realistic "state after a small
// edit" shape that the pointer-equality skip in unionSets is meant to exploit.
static void BM_SetShared(benchmark::State &State, bool UseNew, int N, int D) {
std::vector<int> KA(N);
std::iota(KA.begin(), KA.end(), 0);
std::unique_ptr<IntSet::Factory> F;
IntSet A{nullptr}, B{nullptr};
auto reset = [&] {
A = IntSet(nullptr);
B = IntSet(nullptr);
F = std::make_unique<IntSet::Factory>();
A = buildSet(*F, KA);
B = A;
for (int I = 0; I < D; ++I)
B = F->add(B, N + I);
};
reset();
int Since = 0;
for (auto _ : State) {
if (++Since == ResetEvery) {
State.PauseTiming();
reset();
Since = 0;
State.ResumeTiming();
}
IntSet U = UseNew ? setUnionNew(*F, A, B) : setUnionAdd(*F, A, B);
benchmark::DoNotOptimize(U.getRootWithoutRetain());
}
}
static void BM_Map(benchmark::State &State, bool UseNew, bool NearId, int N,
int D) {
std::vector<int> KA, KB;
makeKeys(NearId, N, D, KA, KB);
std::unique_ptr<IntMap::Factory> F;
IntMap A{nullptr}, B{nullptr};
auto reset = [&] {
A = IntMap(nullptr);
B = IntMap(nullptr);
F = std::make_unique<IntMap::Factory>();
IntMap MA = F->getEmptyMap(), MB = F->getEmptyMap();
for (int V : KA)
MA = F->add(MA, V, V);
for (int V : KB)
MB = F->add(MB, V, V);
A = MA;
B = MB;
};
reset();
int Since = 0;
for (auto _ : State) {
if (++Since == ResetEvery) {
State.PauseTiming();
reset();
Since = 0;
State.ResumeTiming();
}
IntMap U = UseNew ? mapUnionNew(*F, A, B) : mapUnionAdd(*F, A, B);
benchmark::DoNotOptimize(U.getRootWithoutRetain());
}
}
// Map<int, Set<int>>: each key maps to a small set; the merge key-wise unions
// the value sets (the lifetime OriginLoanMap join).
static void BM_MapOfSets(benchmark::State &State, bool UseNew, bool NearId,
int N, int D, int SetSize) {
std::vector<int> KA, KB;
makeKeys(NearId, N, D, KA, KB);
std::unique_ptr<SetMap::Factory> MF;
std::unique_ptr<IntSet::Factory> SF;
SetMap A{nullptr}, B{nullptr};
auto valueSet = [&](int Key) {
// A deterministic small set derived from the key.
IntSet S = SF->getEmptySet();
for (int I = 0; I < SetSize; ++I)
S = SF->add(S, Key * 31 + I);
return S;
};
auto reset = [&] {
A = SetMap(nullptr);
B = SetMap(nullptr);
MF = std::make_unique<SetMap::Factory>();
SF = std::make_unique<IntSet::Factory>();
SetMap MA = MF->getEmptyMap(), MB = MF->getEmptyMap();
for (int V : KA)
MA = MF->add(MA, V, valueSet(V));
for (int V : KB)
MB = MF->add(MB, V, valueSet(V));
A = MA;
B = MB;
};
reset();
int Since = 0;
for (auto _ : State) {
if (++Since == ResetEvery) {
State.PauseTiming();
reset();
Since = 0;
State.ResumeTiming();
}
SetMap U = UseNew ? setMapUnionNew(*MF, *SF, A, B)
: setMapUnionAdd(*MF, *SF, A, B);
benchmark::DoNotOptimize(U.getRootWithoutRetain());
}
}
} // namespace
// Small: both operands tiny. NearId: size N, differ by D keys each side.
BENCHMARK_CAPTURE(BM_Set, "Set/Small/Add", false, false, 8, 0);
BENCHMARK_CAPTURE(BM_Set, "Set/Small/Union", true, false, 8, 0);
BENCHMARK_CAPTURE(BM_Set, "Set/Small16/Add", false, false, 16, 0);
BENCHMARK_CAPTURE(BM_Set, "Set/Small16/Union", true, false, 16, 0);
// Small and nearly identical (share all but one key each side) -- the common
// dataflow live-set shape, and the case most exposed to split overhead.
BENCHMARK_CAPTURE(BM_Set, "Set/SmallNearId8/Add", false, true, 8, 1);
BENCHMARK_CAPTURE(BM_Set, "Set/SmallNearId8/Union", true, true, 8, 1);
// Very small sets, delta 2 (and a couple of neighbours) to see if the Δ1
// regression persists at tiny sizes once the diff is at least 2.
BENCHMARK_CAPTURE(BM_Set, "Set/N8D2/Add", false, true, 8, 2);
BENCHMARK_CAPTURE(BM_Set, "Set/N8D2/Union", true, true, 8, 2);
BENCHMARK_CAPTURE(BM_Set, "Set/N16D2/Add", false, true, 16, 2);
BENCHMARK_CAPTURE(BM_Set, "Set/N16D2/Union", true, true, 16, 2);
BENCHMARK_CAPTURE(BM_Set, "Set/N32D2/Add", false, true, 32, 2);
BENCHMARK_CAPTURE(BM_Set, "Set/N32D2/Union", true, true, 32, 2);
BENCHMARK_CAPTURE(BM_Set, "Set/SmallNearId64/Add", false, true, 64, 1);
BENCHMARK_CAPTURE(BM_Set, "Set/SmallNearId64/Union", true, true, 64, 1);
// Delta sweep at N=64 (where the near-identical regression peaked): how does
// the gap behave as the two sets differ by more keys?
BENCHMARK_CAPTURE(BM_Set, "Set/N64D2/Add", false, true, 64, 2);
BENCHMARK_CAPTURE(BM_Set, "Set/N64D2/Union", true, true, 64, 2);
BENCHMARK_CAPTURE(BM_Set, "Set/N64D4/Add", false, true, 64, 4);
BENCHMARK_CAPTURE(BM_Set, "Set/N64D4/Union", true, true, 64, 4);
BENCHMARK_CAPTURE(BM_Set, "Set/N64D8/Add", false, true, 64, 8);
BENCHMARK_CAPTURE(BM_Set, "Set/N64D8/Union", true, true, 64, 8);
BENCHMARK_CAPTURE(BM_Set, "Set/N64D16/Add", false, true, 64, 16);
BENCHMARK_CAPTURE(BM_Set, "Set/N64D16/Union", true, true, 64, 16);
BENCHMARK_CAPTURE(BM_Set, "Set/N64D32/Add", false, true, 64, 32);
BENCHMARK_CAPTURE(BM_Set, "Set/N64D32/Union", true, true, 64, 32);
// Δ1 but B derived from A (shares subtrees by pointer): the case the
// pointer-equality skip targets. Compare against the independent-build
// Set/*NearId* numbers above.
BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N8D1/Add", false, 8, 1);
BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N8D1/Union", true, 8, 1);
BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N64D1/Add", false, 64, 1);
BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N64D1/Union", true, 64, 1);
BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N1024D1/Add", false, 1024, 1);
BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N1024D1/Union", true, 1024, 1);
BENCHMARK_CAPTURE(BM_Set, "Set/NearId128/Add", false, true, 128, 1);
BENCHMARK_CAPTURE(BM_Set, "Set/NearId128/Union", true, true, 128, 1);
BENCHMARK_CAPTURE(BM_Set, "Set/NearId256/Add", false, true, 256, 1);
BENCHMARK_CAPTURE(BM_Set, "Set/NearId256/Union", true, true, 256, 1);
BENCHMARK_CAPTURE(BM_Set, "Set/NearId512/Add", false, true, 512, 1);
BENCHMARK_CAPTURE(BM_Set, "Set/NearId512/Union", true, true, 512, 1);
BENCHMARK_CAPTURE(BM_Set, "Set/NearId/Add", false, true, 1024, 1);
BENCHMARK_CAPTURE(BM_Set, "Set/NearId/Union", true, true, 1024, 1);
BENCHMARK_CAPTURE(BM_Map, "Map/Small/Add", false, false, 8, 0);
BENCHMARK_CAPTURE(BM_Map, "Map/Small/Union", true, false, 8, 0);
BENCHMARK_CAPTURE(BM_Map, "Map/SmallNearId8/Add", false, true, 8, 1);
BENCHMARK_CAPTURE(BM_Map, "Map/SmallNearId8/Union", true, true, 8, 1);
BENCHMARK_CAPTURE(BM_Map, "Map/SmallNearId64/Add", false, true, 64, 1);
BENCHMARK_CAPTURE(BM_Map, "Map/SmallNearId64/Union", true, true, 64, 1);
BENCHMARK_CAPTURE(BM_Map, "Map/NearId/Add", false, true, 1024, 1);
BENCHMARK_CAPTURE(BM_Map, "Map/NearId/Union", true, true, 1024, 1);
BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/Small/Add", false, false, 8, 0, 4);
BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/Small/Union", true, false, 8, 0, 4);
BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/SmallNearId8/Add", false, true, 8, 1,
4);
BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/SmallNearId8/Union", true, true, 8,
1, 4);
BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/NearId/Add", false, true, 256, 1, 8);
BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/NearId/Union", true, true, 256, 1,
8);
BENCHMARK_MAIN();

21
llvm/include/llvm/ADT/ImmutableMap.h
View file

@ -103,6 +103,27 @@ public:
return ImmutableMap(T);
}
/// Merges \p A and \p B in a single traversal (see
/// ImutAVLFactory::mergeTrees), sharing subtrees the two maps do not
/// overlap. \p Combine(AElem, BElem) yields the stored element for a key;
/// \p KeepUnmatched governs keys unique to one side. \p SkipShared returns
/// a pointer-identical subtree unchanged in O(1); only pass true when the
/// merge is idempotent (Combine(a, a) == a, e.g. a lattice join). This is
/// more efficient than repeatedly adding \p B's entries to \p A when \p B
/// is large. Only valid for non-canonicalizing factories (the bulk path
/// does not canonicalize the nodes it creates). Does not short-circuit
/// equal/empty operands or reorder by size; callers wanting those apply
/// them first.
template <typename CombineFn>
[[nodiscard]] ImmutableMap mergeWith(ImmutableMap A, ImmutableMap B,
CombineFn Combine, bool KeepUnmatched,
bool SkipShared = false) {
static_assert(!Canonicalize,
"mergeWith does not canonicalize the nodes it creates");
return ImmutableMap(F.mergeTrees(A.Root.get(), B.Root.get(), Combine,
KeepUnmatched, SkipShared));
}
[[nodiscard]] ImmutableMap remove(ImmutableMap Old, key_type_ref K) {
TreeTy *T = F.remove(Old.Root.get(), K);
if constexpr (Canonicalize)

167
llvm/include/llvm/ADT/ImmutableSet.h
View file

@ -404,6 +404,7 @@ class ImutAVLFactory
friend class ImutAVLTree<ImutInfo, Canonicalize>;
using TreeTy = ImutAVLTree<ImutInfo, Canonicalize>;
using value_type = typename TreeTy::value_type;
using value_type_ref = typename TreeTy::value_type_ref;
using key_type_ref = typename TreeTy::key_type_ref;
@ -440,6 +441,39 @@ public:
return T;
}
/// Merges \p A and \p B in a single traversal, sharing every subtree that the
/// two operands do not overlap. \p Combine(AElem, BElem) produces the element
/// stored for a key present in both; \p KeepUnmatched governs keys unique to
/// one side (see merge_internal). For merging |B| entries into |A|
/// (|B| <= |A|) this costs O(|B| * log(|A|/|B| + 1)) and copies each spine
/// node at most once, versus O(|B| * log|A|) repeated \ref add descents.
/// \p A and \p B must be immutable. This does not short-circuit equal or
/// empty operands (merge_internal handles them correctly but not specially);
/// callers that want those fast paths, or size-driven operand ordering,
/// should apply them first (see ImmutableSet::Factory::unionSets).
template <typename CombineFn>
TreeTy *mergeTrees(TreeTy *A, TreeTy *B, CombineFn Combine,
bool KeepUnmatched, bool SkipShared = false) {
TreeTy *T = merge_internal(A, B, Combine, KeepUnmatched, SkipShared);
recoverNodes(T);
return T;
}
/// Returns the set union of \p A and \p B (keeping \p A's element on matching
/// keys). Shorthand for the fully sharing \ref mergeTrees.
TreeTy *unionTrees(TreeTy *A, TreeTy *B) {
// With KeepUnmatched=true, unmatched elements are shared as-is and Combine
// is invoked only for keys present in both, where it keeps A's element.
auto KeepFirst = [](const value_type *L,
const value_type *R) -> const value_type & {
return L ? *L : *R;
};
// Set union is idempotent, so identical (pointer-equal) subtrees -- common
// once one operand is derived from the other -- can be shared in O(1).
return mergeTrees(A, B, KeepFirst, /*KeepUnmatched=*/true,
/*SkipShared=*/true);
}
TreeTy* remove(TreeTy* T, key_type_ref V) {
T = remove_internal(V,T);
recoverNodes(T);
@ -561,6 +595,115 @@ protected:
return createNode(L,V,R);
}
/// Combines \p L and \p R with the value \p V (every key in \p L less than
/// \p V, every key in \p R greater) into one balanced tree. Unlike
/// balanceTree this tolerates an arbitrary height difference between \p L and
/// \p R: it descends the taller side's spine and rebalances on the way back
/// up, exactly as an insertion would.
TreeTy *joinTrees(TreeTy *L, value_type_ref V, TreeTy *R) {
if (getHeight(L) > getHeight(R) + 2)
return balanceTree(getLeft(L), getValue(L), joinTrees(getRight(L), V, R));
if (getHeight(R) > getHeight(L) + 2)
return balanceTree(joinTrees(L, V, getLeft(R)), getValue(R), getRight(R));
return createNode(L, V, R);
}
/// Splits \p T into \p L (all keys less than \p K) and \p R (all keys greater
/// than \p K). If \p K is present in \p T, \p Match is set to point at its
/// element (which is dropped from \p L and \p R); otherwise \p Match is null.
void splitLookup(TreeTy *T, key_type_ref K, TreeTy *&L,
const value_type *&Match, TreeTy *&R) {
if (isEmpty(T)) {
L = R = getEmptyTree();
Match = nullptr;
return;
}
key_type_ref KCurrent = ImutInfo::KeyOfValue(getValue(T));
if (ImutInfo::isEqual(K, KCurrent)) {
L = getLeft(T);
R = getRight(T);
// Use the tree accessor, which returns a reference to the stored element
// (the factory's getValue returns value_type_ref, which is by value for
// pointer-like element types).
Match = &T->getValue();
} else if (ImutInfo::isLess(K, KCurrent)) {
TreeTy *LR;
splitLookup(getLeft(T), K, L, Match, LR);
R = joinTrees(LR, getValue(T), getRight(T));
} else {
TreeTy *RL;
splitLookup(getRight(T), K, RL, Match, R);
L = joinTrees(getLeft(T), getValue(T), RL);
}
}
/// Rebuilds \p T with the same shape but each element replaced by
/// \p Combine applied to it. \p FromB selects which side of \p Combine the
/// element is passed on (it is the sole non-null argument).
template <typename CombineFn>
TreeTy *transformTree(TreeTy *T, CombineFn &Combine, bool FromB) {
if (isEmpty(T))
return T;
TreeTy *L = transformTree(getLeft(T), Combine, FromB);
TreeTy *R = transformTree(getRight(T), Combine, FromB);
const value_type &E = getValue(T);
return createNode(L, FromB ? Combine(nullptr, &E) : Combine(&E, nullptr),
R);
}
/// Merges \p A and \p B by recursing over \p A's structure and splitting \p B
/// at each of \p A's keys. For a key in both, the stored element is
/// Combine(AElem, BElem). \p KeepUnmatched controls keys unique to one side:
/// when true, such elements (and whole non-overlapping subtrees) are taken
/// unchanged and shared, and \p Combine is invoked only on keys present in
/// both (valid when \p Combine is an identity for a missing side, e.g. a set
/// union or a lattice join with an identity element); when false every key is
/// passed through \p Combine with the absent side null (needed for a join
/// that transforms unmatched keys, e.g. liveness downgrading Must to Maybe).
template <typename CombineFn>
TreeTy *merge_internal(TreeTy *A, TreeTy *B, CombineFn &Combine,
bool KeepUnmatched, bool SkipShared) {
// When A and B are the same tree (which happens all the time once B is
// derived from A by a small edit, since the untouched side is shared by
// pointer), an idempotent merge returns it unchanged in O(1). Only valid
// when merge(x, x) == x, so the caller opts in via SkipShared.
if (SkipShared && A == B)
return A;
if (isEmpty(A))
return KeepUnmatched ? B : transformTree(B, Combine, /*FromB=*/true);
if (isEmpty(B))
return KeepUnmatched ? A : transformTree(A, Combine, /*FromB=*/false);
const value_type &AElem = getValue(A);
TreeTy *BL, *BR;
const value_type *BMatch;
splitLookup(B, ImutInfo::KeyOfValue(AElem), BL, BMatch, BR);
TreeTy *NewL =
merge_internal(getLeft(A), BL, Combine, KeepUnmatched, SkipShared);
TreeTy *NewR =
merge_internal(getRight(A), BR, Combine, KeepUnmatched, SkipShared);
if (!BMatch) {
// Key present only in A.
if (KeepUnmatched) {
if (NewL == getLeft(A) && NewR == getRight(A))
return A;
return joinTrees(NewL, AElem, NewR);
}
return joinTrees(NewL, Combine(&AElem, nullptr), NewR);
}
// Key present in both: combine the two elements. Preserve sharing when the
// combined value is unchanged and neither subtree moved, so that a join
// that only touches a few keys does not rebuild the whole spine.
auto NewElem = Combine(&AElem, BMatch);
if (NewL == getLeft(A) && NewR == getRight(A) &&
ImutInfo::isDataEqual(ImutInfo::DataOfValue(NewElem),
ImutInfo::DataOfValue(AElem)))
return A;
return joinTrees(NewL, NewElem, NewR);
}
/// add_internal - Creates a new tree that includes the specified
/// data and the data from the original tree. If the original tree
/// already contained the data item, the original tree is returned.
@ -994,6 +1137,30 @@ public:
return ImmutableSet(NewT);
}
/// Returns the union of \p A and \p B, computed in a single traversal that
/// shares subtrees of both operands wherever possible (see
/// ImutAVLFactory::unionTrees). This is more efficient than repeatedly
/// adding \p B's elements to \p A when \p B is large.
[[nodiscard]] ImmutableSet unionSets(ImmutableSet A, ImmutableSet B) {
if (A.Root.get() == B.Root.get() || B.isEmpty())
return A;
if (A.isEmpty())
return B;
// Drive the recursion with the taller tree so the shorter one is the one
// being split.
if (A.getHeight() < B.getHeight())
std::swap(A, B);
if constexpr (Canonicalize) {
// The bulk path does not canonicalize the nodes it creates, so fall
// back to per-element insertion for canonicalizing factories.
for (value_type_ref V : B)
A = add(A, V);
return A;
} else {
return ImmutableSet(F.unionTrees(A.Root.get(), B.Root.get()));
}
}
/// Creates a new immutable set that contains all of the values
/// of the original set with the exception of the specified value. If
/// the original set did not contain the value, the original set is

82
llvm/unittests/ADT/ImmutableMapTest.cpp
View file

@ -8,9 +8,18 @@
#include "llvm/ADT/ImmutableMap.h"
#include "gtest/gtest.h"
#include <map>
#include <random>
#include <utility>
using namespace llvm;
// Non-canonicalizing map, matching how the clang lifetime analyses use it and
// the only configuration for which the bulk mergeWith path is valid.
template <typename K, typename V>
using NCMap =
ImmutableMap<K, V, ImutKeyValueInfo<K, V>, /*Canonicalize=*/false>;
namespace {
TEST(ImmutableMapTest, EmptyIntMapTest) {
@ -87,4 +96,77 @@ TEST(ImmutableMapTest, MultiElemIntMapRefTest) {
EXPECT_TRUE(f.getEmptyMap() == f.getEmptyMap());
}
// Exercises Factory::mergeWith against a std::map oracle across the
// KeepA/KeepB configurations used by the set union (true/true), an asymmetric
// map join (true/false), and a symmetric map join (false/false), with a
// combiner that is order- and side-sensitive so the flag semantics are
// actually tested.
TEST(ImmutableMapTest, MergeWithStressTest) {
using Map = NCMap<int, int>;
using VT = Map::value_type; // const std::pair<int, int>
auto Combine = [](const VT *A, const VT *B) -> std::pair<int, int> {
int Key = A ? A->first : B->first;
int AV = A ? A->second : 0;
int BV = B ? B->second : 0;
return {Key, 1 + 3 * AV + 5 * BV};
};
std::mt19937 Rng(7);
std::uniform_int_distribution<int> SizeDist(0, 120);
std::uniform_int_distribution<int> KeyDist(0, 200);
std::uniform_int_distribution<int> ValDist(0, 1000);
for (bool KeepUnmatched : {true, false}) {
for (int Trial = 0; Trial < 300; ++Trial) {
Map::Factory F;
Map A = F.getEmptyMap(), B = F.getEmptyMap();
std::map<int, int> RefA, RefB;
int NA = SizeDist(Rng), NB = SizeDist(Rng);
for (int I = 0; I < NA; ++I) {
int K = KeyDist(Rng), V = ValDist(Rng);
A = F.add(A, K, V);
RefA[K] = V;
}
for (int I = 0; I < NB; ++I) {
int K = KeyDist(Rng), V = ValDist(Rng);
B = F.add(B, K, V);
RefB[K] = V;
}
Map U = F.mergeWith(A, B, Combine, KeepUnmatched);
if (const auto *Root = U.getRootWithoutRetain())
Root->validateTree();
// Reference result: matched keys always combined; unmatched keys are
// kept as-is when KeepUnmatched, else passed through Combine.
std::map<int, int> Exp;
for (auto &[K, AV] : RefA) {
auto It = RefB.find(K);
if (It != RefB.end())
Exp[K] = 1 + 3 * AV + 5 * It->second;
else
Exp[K] = KeepUnmatched ? AV : 1 + 3 * AV;
}
for (auto &[K, BV] : RefB)
if (!RefA.count(K))
Exp[K] = KeepUnmatched ? BV : 1 + 5 * BV;
std::map<int, int> Got;
for (Map::iterator It = U.begin(), E = U.end(); It != E; ++It)
Got[It.getKey()] = It.getData();
EXPECT_EQ(Exp, Got) << "KeepUnmatched=" << KeepUnmatched;
// Inputs unchanged.
std::map<int, int> GotA, GotB;
for (Map::iterator It = A.begin(), E = A.end(); It != E; ++It)
GotA[It.getKey()] = It.getData();
for (Map::iterator It = B.begin(), E = B.end(); It != E; ++It)
GotB[It.getKey()] = It.getData();
EXPECT_EQ(RefA, GotA);
EXPECT_EQ(RefB, GotB);
}
}
}
}

54
llvm/unittests/ADT/ImmutableSetTest.cpp
View file

@ -387,4 +387,58 @@ TEST_F(ImmutableSetTest, IteratorEquality) {
EXPECT_TRUE(A == TreeIter());
}
}
// Exercises Factory::unionSets against a std::set oracle over many random
// input pairs, for both the canonicalizing (per-element fallback) and
// non-canonicalizing (bulk tree-tree) code paths.
template <bool Canonicalize> static void unionSetsStress() {
using Set = ImmutableSet<int, ImutContainerInfo<int>, Canonicalize>;
std::mt19937 Rng(42);
std::uniform_int_distribution<int> SizeDist(0, 120);
std::uniform_int_distribution<int> ValDist(0, 200);
for (int Trial = 0; Trial < 300; ++Trial) {
typename Set::Factory F;
Set A = F.getEmptySet();
Set B = F.getEmptySet();
std::set<int> RefA, RefB;
int NA = SizeDist(Rng), NB = SizeDist(Rng);
for (int I = 0; I < NA; ++I) {
int V = ValDist(Rng);
A = F.add(A, V);
RefA.insert(V);
}
for (int I = 0; I < NB; ++I) {
int V = ValDist(Rng);
B = F.add(B, V);
RefB.insert(V);
}
Set U = F.unionSets(A, B);
// The result is a structurally valid AVL tree.
if (const auto *Root = U.getRootWithoutRetain())
Root->validateTree();
// Its contents are exactly the union of the inputs.
std::set<int> RefU = RefA;
RefU.insert(RefB.begin(), RefB.end());
std::set<int> Got(U.begin(), U.end());
EXPECT_EQ(RefU, Got);
// The inputs are unchanged (persistence).
EXPECT_EQ(RefA, std::set<int>(A.begin(), A.end()));
EXPECT_EQ(RefB, std::set<int>(B.begin(), B.end()));
// Unioning with a subset / self is the identity.
EXPECT_TRUE(F.unionSets(U, A) == U);
EXPECT_TRUE(F.unionSets(A, A) == A);
}
}
TEST_F(ImmutableSetTest, UnionSetsStressTest) {
unionSetsStress</*Canonicalize=*/false>();
unionSetsStress</*Canonicalize=*/true>();
}
} // namespace