[llvm] [ADT] Rewrite ImmutableSet/Map in-order iterator without per-node state (PR #205552)
via llvm-commits
llvm-commits at lists.llvm.org
Wed Jun 24 07:50:19 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-adt
Author: Gábor Horváth (Xazax-hun)
<details>
<summary>Changes</summary>
ImutAVLTreeInOrderIterator was layered on a "generic" iterator whose stack stored a 2-bit visit-state (None/Left/Right) packed into the low bits of each node pointer, and which was spun through several intermediate states per in-order step.
Replace it with a single iterator that keeps a stack of plain node pointers (the root-to-current ancestor chain) and recovers the traversal direction by inspecting whether it is ascending from a node's left or right child. A node's parent cannot be cached in the node itself because these trees are persistent and structurally shared (one node may be a child of different parents across tree versions), so the ancestor stack is the per-traversal parent chain.
The observable contract is unchanged: same in-order sequence, same skipSubTree() semantics (so tree canonicalization via isEqual is unaffected), and the iterator stays bidirectional. operator== now compares the current node rather than the whole path, which is O(1) and correct because within a single tree a node has a unique path.
Microbenchmarks (llvm/benchmarks/ImmutableSetIteratorBM.cpp, -O2 with assertions) show ~2x faster plain traversal and skipSubTree, and the node-identity operator== is ~1.5x faster on comparison-heavy loops (no change on the common compare-against-end() path).
Assisted by: Claude Opus 4.8
---
Patch is 22.26 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/205552.diff
4 Files Affected:
- (modified) llvm/benchmarks/CMakeLists.txt (+1)
- (added) llvm/benchmarks/ImmutableSetIteratorBM.cpp (+121)
- (modified) llvm/include/llvm/ADT/ImmutableSet.h (+79-134)
- (modified) llvm/unittests/ADT/ImmutableSetTest.cpp (+187)
``````````diff
diff --git a/llvm/benchmarks/CMakeLists.txt b/llvm/benchmarks/CMakeLists.txt
index 69ebeaa78344b..fbfc9a63e1e8c 100644
--- a/llvm/benchmarks/CMakeLists.txt
+++ b/llvm/benchmarks/CMakeLists.txt
@@ -15,6 +15,7 @@ add_benchmark(MustacheBench Mustache.cpp PARTIAL_SOURCES_INTENDED)
add_benchmark(SpecialCaseListBM SpecialCaseListBM.cpp PARTIAL_SOURCES_INTENDED)
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(RuntimeLibcallsBench RuntimeLibcalls.cpp PARTIAL_SOURCES_INTENDED)
diff --git a/llvm/benchmarks/ImmutableSetIteratorBM.cpp b/llvm/benchmarks/ImmutableSetIteratorBM.cpp
new file mode 100644
index 0000000000000..3ee704ede4374
--- /dev/null
+++ b/llvm/benchmarks/ImmutableSetIteratorBM.cpp
@@ -0,0 +1,121 @@
+//===- ImmutableSetIteratorBM.cpp - Benchmark ImmutableSet iterators ------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Benchmarks in-order traversal of the ImutAVLTree backing ImmutableSet using
+// only its public iterator API. It does not compare iterator implementations
+// directly; instead, run the binary before and after a change to the iterator
+// and compare the two reports (e.g. with llvm/utils/compare.py) to see the
+// effect of that change.
+//
+// Two access patterns are measured:
+// * Iterate - a plain forward walk over ImmutableSet::iterator, the common
+// client usage.
+// * Skip - a walk that calls skipSubTree on the tree iterator at every
+// other node, the pattern used by ImutAVLTree::isEqual and the
+// tree canonicalization that the clang static analyzer relies on.
+//
+//===----------------------------------------------------------------------===//
+
+#include "benchmark/benchmark.h"
+#include "llvm/ADT/ImmutableSet.h"
+#include <algorithm>
+#include <cstdint>
+#include <numeric>
+#include <random>
+#include <vector>
+
+using namespace llvm;
+
+namespace {
+
+using Tree = ImmutableSet<int>::TreeTy;
+
+// Holds a factory plus a built set so the (non-trivial) tree construction is
+// kept out of the timed region. The factory must outlive the tree.
+struct Fixture {
+ ImmutableSet<int>::Factory F{/*canonicalize=*/false};
+ ImmutableSet<int> Set = F.getEmptySet();
+
+ explicit Fixture(size_t N) {
+ std::vector<int> Vals(N);
+ std::iota(Vals.begin(), Vals.end(), 0);
+ std::mt19937 Rng(0xC0FFEE);
+ std::shuffle(Vals.begin(), Vals.end(), Rng);
+ for (int V : Vals)
+ Set = F.add(Set, V);
+ }
+};
+
+// Plain forward iteration over the public ImmutableSet::iterator.
+static void BM_Iterate(benchmark::State &State) {
+ const size_t N = State.range(0);
+ Fixture Fix(N);
+ const ImmutableSet<int> &S = Fix.Set;
+ benchmark::DoNotOptimize(S.getRootWithoutRetain());
+
+ for (auto _ : State) {
+ int64_t Sum = 0;
+ for (int V : S)
+ Sum += V;
+ benchmark::DoNotOptimize(Sum);
+ }
+ State.SetItemsProcessed(State.iterations() * N);
+}
+
+// Iterate while skipping every other subtree, exercising skipSubTree (the
+// pattern behind ImutAVLTree::isEqual / canonicalization) rather than plain ++.
+static void BM_IterateWithSkips(benchmark::State &State) {
+ const size_t N = State.range(0);
+ Fixture Fix(N);
+ const Tree *Root = Fix.Set.getRootWithoutRetain();
+ benchmark::DoNotOptimize(Root);
+
+ for (auto _ : State) {
+ int64_t Sum = 0;
+ unsigned K = 0;
+ for (Tree::iterator I(Root), E; I != E; ++K) {
+ Sum += I->getValue();
+ if (K & 1)
+ I.skipSubTree();
+ else
+ ++I;
+ }
+ benchmark::DoNotOptimize(Sum);
+ }
+}
+
+// Compare two iterators positioned at the same (non-end) node on every step.
+// This isolates operator==: it is the case that differs between comparing the
+// current node only (O(1)) and comparing the whole ancestor path (O(depth)).
+// It is a microbenchmark of the comparison itself, not a typical client loop
+// (those compare against end(), which is O(1) either way).
+static void BM_CompareSamePosition(benchmark::State &State) {
+ const size_t N = State.range(0);
+ Fixture Fix(N);
+ const Tree *Root = Fix.Set.getRootWithoutRetain();
+ benchmark::DoNotOptimize(Root);
+
+ for (auto _ : State) {
+ bool R = false;
+ Tree::iterator E;
+ for (Tree::iterator A(Root), B(Root); A != E; ++A, ++B)
+ R ^= (A == B);
+ benchmark::DoNotOptimize(R);
+ }
+ State.SetItemsProcessed(State.iterations() * N);
+}
+
+} // namespace
+
+#define ITER_SIZES Arg(16)->Arg(256)->Arg(4096)->Arg(65536)
+
+BENCHMARK(BM_Iterate)->Name("Iterate")->ITER_SIZES;
+BENCHMARK(BM_IterateWithSkips)->Name("Skip")->ITER_SIZES;
+BENCHMARK(BM_CompareSamePosition)->Name("CompareSamePos")->ITER_SIZES;
+
+BENCHMARK_MAIN();
diff --git a/llvm/include/llvm/ADT/ImmutableSet.h b/llvm/include/llvm/ADT/ImmutableSet.h
index b9c3e68e51365..b6ddcf0a23d2e 100644
--- a/llvm/include/llvm/ADT/ImmutableSet.h
+++ b/llvm/include/llvm/ADT/ImmutableSet.h
@@ -40,7 +40,6 @@ namespace llvm {
template <typename ImutInfo> class ImutAVLFactory;
template <typename ImutInfo> class ImutIntervalAVLFactory;
template <typename ImutInfo> class ImutAVLTreeInOrderIterator;
-template <typename ImutInfo> class ImutAVLTreeGenericIterator;
template <typename ImutInfo >
class ImutAVLTree {
@@ -53,7 +52,6 @@ class ImutAVLTree {
friend class ImutAVLFactory<ImutInfo>;
friend class ImutIntervalAVLFactory<ImutInfo>;
- friend class ImutAVLTreeGenericIterator<ImutInfo>;
//===----------------------------------------------------===//
// Public Interface.
@@ -652,12 +650,24 @@ class ImutAVLFactory {
};
//===----------------------------------------------------------------------===//
-// Immutable AVL-Tree Iterators.
+// Immutable AVL-Tree Iterator.
//===----------------------------------------------------------------------===//
-template <typename ImutInfo> class ImutAVLTreeGenericIterator {
- SmallVector<uintptr_t,20> stack;
-
+/// Bidirectional in-order iterator over the nodes of an ImutAVLTree.
+///
+/// The iterator keeps the chain of ancestors from the root down to the current
+/// node on an explicit stack of plain node pointers, and decides which way to
+/// move next by inspecting whether it is ascending from a node's left or right
+/// child. This avoids storing any per-node visit-state: there is no need to
+/// remember "have I already visited this node's left/right subtree", because
+/// that is recovered by comparing the child we just left against the parent's
+/// left and right pointers.
+///
+/// A node's parent cannot be cached in the node itself, because these trees are
+/// persistent and structurally shared: a single node may appear as the child of
+/// different parents across different tree versions. The ancestor stack is
+/// therefore the per-traversal parent chain.
+template <typename ImutInfo> class ImutAVLTreeInOrderIterator {
public:
using iterator_category = std::bidirectional_iterator_tag;
using value_type = ImutAVLTree<ImutInfo>;
@@ -665,164 +675,99 @@ template <typename ImutInfo> class ImutAVLTreeGenericIterator {
using pointer = value_type *;
using reference = value_type &;
- enum VisitFlag { VisitedNone=0x0, VisitedLeft=0x1, VisitedRight=0x3,
- Flags=0x3 };
-
using TreeTy = ImutAVLTree<ImutInfo>;
- ImutAVLTreeGenericIterator() = default;
- ImutAVLTreeGenericIterator(const TreeTy *Root) {
- if (Root) stack.push_back(reinterpret_cast<uintptr_t>(Root));
- }
-
- TreeTy &operator*() const {
- assert(!stack.empty());
- return *reinterpret_cast<TreeTy *>(stack.back() & ~Flags);
- }
- TreeTy *operator->() const { return &*this; }
-
- uintptr_t getVisitState() const {
- assert(!stack.empty());
- return stack.back() & Flags;
- }
-
- bool atEnd() const { return stack.empty(); }
-
- bool atBeginning() const {
- return stack.size() == 1 && getVisitState() == VisitedNone;
- }
-
- void skipToParent() {
- assert(!stack.empty());
- stack.pop_back();
- if (stack.empty())
- return;
- switch (getVisitState()) {
- case VisitedNone:
- stack.back() |= VisitedLeft;
- break;
- case VisitedLeft:
- stack.back() |= VisitedRight;
- break;
- default:
- llvm_unreachable("Unreachable.");
- }
- }
+private:
+ // Path[0] is the root and Path.back() is the current node. An empty path is
+ // the end iterator. The invariant is that Path always holds the exact chain
+ // of ancestors of the current node, root-most first.
+ SmallVector<TreeTy *, 20> Path;
- bool operator==(const ImutAVLTreeGenericIterator &x) const {
- return stack == x.stack;
+ // Descend along left children, pushing each node; lands on the minimum of the
+ // subtree rooted at T (i.e. the first node in an in-order traversal of T).
+ void descendToMin(TreeTy *T) {
+ for (; T; T = T->getLeft())
+ Path.push_back(T);
}
- bool operator!=(const ImutAVLTreeGenericIterator &x) const {
- return !(*this == x);
+ // Descend along right children, pushing each node; lands on the maximum of
+ // the subtree rooted at T (i.e. the last node in an in-order traversal of T).
+ void descendToMax(TreeTy *T) {
+ for (; T; T = T->getRight())
+ Path.push_back(T);
}
- ImutAVLTreeGenericIterator &operator++() {
- assert(!stack.empty());
- TreeTy* Current = reinterpret_cast<TreeTy*>(stack.back() & ~Flags);
- assert(Current);
- switch (getVisitState()) {
- case VisitedNone:
- if (TreeTy* L = Current->getLeft())
- stack.push_back(reinterpret_cast<uintptr_t>(L));
- else
- stack.back() |= VisitedLeft;
- break;
- case VisitedLeft:
- if (TreeTy* R = Current->getRight())
- stack.push_back(reinterpret_cast<uintptr_t>(R));
- else
- stack.back() |= VisitedRight;
- break;
- case VisitedRight:
- skipToParent();
- break;
- default:
- llvm_unreachable("Unreachable.");
- }
- return *this;
+ // Pop the current node and ascend until we reach an ancestor from its *left*
+ // child, i.e. the first ancestor whose subtree is not yet fully visited. That
+ // ancestor is the in-order successor of the subtree we just left; if there is
+ // none, Path is emptied (the end iterator). Shared by operator++ and
+ // skipSubTree, whose only difference is whether the current node's right
+ // subtree is descended into first.
+ void ascendFromRightChild() {
+ TreeTy *Child = Path.pop_back_val();
+ while (!Path.empty() && Path.back()->getRight() == Child)
+ Child = Path.pop_back_val();
}
- ImutAVLTreeGenericIterator &operator--() {
- assert(!stack.empty());
- TreeTy* Current = reinterpret_cast<TreeTy*>(stack.back() & ~Flags);
- assert(Current);
- switch (getVisitState()) {
- case VisitedNone:
- stack.pop_back();
- break;
- case VisitedLeft:
- stack.back() &= ~Flags; // Set state to "VisitedNone."
- if (TreeTy* L = Current->getLeft())
- stack.push_back(reinterpret_cast<uintptr_t>(L) | VisitedRight);
- break;
- case VisitedRight:
- stack.back() &= ~Flags;
- stack.back() |= VisitedLeft;
- if (TreeTy* R = Current->getRight())
- stack.push_back(reinterpret_cast<uintptr_t>(R) | VisitedRight);
- break;
- default:
- llvm_unreachable("Unreachable.");
- }
- return *this;
+ // Mirror of ascendFromRightChild for reverse traversal (operator--).
+ void ascendFromLeftChild() {
+ TreeTy *Child = Path.pop_back_val();
+ while (!Path.empty() && Path.back()->getLeft() == Child)
+ Child = Path.pop_back_val();
}
-};
-
-template <typename ImutInfo> class ImutAVLTreeInOrderIterator {
- using InternalIteratorTy = ImutAVLTreeGenericIterator<ImutInfo>;
-
- InternalIteratorTy InternalItr;
public:
- using iterator_category = std::bidirectional_iterator_tag;
- using value_type = ImutAVLTree<ImutInfo>;
- using difference_type = std::ptrdiff_t;
- using pointer = value_type *;
- using reference = value_type &;
-
- using TreeTy = ImutAVLTree<ImutInfo>;
-
- ImutAVLTreeInOrderIterator(const TreeTy* Root) : InternalItr(Root) {
- if (Root)
- ++*this; // Advance to first element.
+ ImutAVLTreeInOrderIterator() = default; // end() iterator.
+ ImutAVLTreeInOrderIterator(const TreeTy *Root) {
+ descendToMin(const_cast<TreeTy *>(Root));
}
- ImutAVLTreeInOrderIterator() : InternalItr() {}
-
+ // Two iterators are equal iff they sit on the same node (or are both end()).
+ // Within a single tree a node has a unique root-to-node path, so the current
+ // node alone identifies the position; comparing the whole path is therefore
+ // unnecessary. Comparing iterators from different trees is not meaningful, as
+ // for any standard container.
bool operator==(const ImutAVLTreeInOrderIterator &x) const {
- return InternalItr == x.InternalItr;
+ if (Path.empty() || x.Path.empty())
+ return Path.empty() == x.Path.empty();
+ return Path.back() == x.Path.back();
}
-
bool operator!=(const ImutAVLTreeInOrderIterator &x) const {
return !(*this == x);
}
- TreeTy &operator*() const { return *InternalItr; }
- TreeTy *operator->() const { return &*InternalItr; }
+ TreeTy &operator*() const { return *Path.back(); }
+ TreeTy *operator->() const { return Path.back(); }
ImutAVLTreeInOrderIterator &operator++() {
- do ++InternalItr;
- while (!InternalItr.atEnd() &&
- InternalItr.getVisitState() != InternalIteratorTy::VisitedLeft);
-
+ assert(!Path.empty() && "Incrementing the end iterator");
+ if (TreeTy *R = Path.back()->getRight())
+ // The in-order successor is the minimum of the right subtree.
+ descendToMin(R);
+ else
+ // No right subtree: the successor is the nearest ancestor reached from a
+ // left child.
+ ascendFromRightChild();
return *this;
}
ImutAVLTreeInOrderIterator &operator--() {
- do --InternalItr;
- while (!InternalItr.atBeginning() &&
- InternalItr.getVisitState() != InternalIteratorTy::VisitedLeft);
-
+ assert(!Path.empty() && "Decrementing the end iterator");
+ if (TreeTy *L = Path.back()->getLeft())
+ // The in-order predecessor is the maximum of the left subtree.
+ descendToMax(L);
+ else
+ // Mirror of operator++.
+ ascendFromLeftChild();
return *this;
}
+ /// Move to the in-order successor of the entire subtree rooted at the current
+ /// node, i.e. skip the current node together with its right subtree. This is
+ /// exactly the ascent half of operator++.
void skipSubTree() {
- InternalItr.skipToParent();
-
- while (!InternalItr.atEnd() &&
- InternalItr.getVisitState() != InternalIteratorTy::VisitedLeft)
- ++InternalItr;
+ assert(!Path.empty() && "Skipping past the end iterator");
+ ascendFromRightChild();
}
};
diff --git a/llvm/unittests/ADT/ImmutableSetTest.cpp b/llvm/unittests/ADT/ImmutableSetTest.cpp
index 87bc2a8da4bad..b3d7e03089675 100644
--- a/llvm/unittests/ADT/ImmutableSetTest.cpp
+++ b/llvm/unittests/ADT/ImmutableSetTest.cpp
@@ -8,6 +8,11 @@
#include "llvm/ADT/ImmutableSet.h"
#include "gtest/gtest.h"
+#include <algorithm>
+#include <numeric>
+#include <random>
+#include <set>
+#include <vector>
using namespace llvm;
@@ -195,4 +200,186 @@ TEST_F(ImmutableSetTest, RemoveIfNotFoundTest) {
ImmutableSet<long> U = f.remove(S, 3);
EXPECT_NE(S.getRoot(), U.getRoot());
}
+
+//===----------------------------------------------------------------------===//
+// In-order iterator correctness, validated against independent oracles.
+//
+// These checks do not assume any particular iterator implementation; they pin
+// the externally observable contract (in-order ordering, reverse traversal,
+// and skipSubTree). This is what the clang static analyzer and the tree
+// canonicalization machinery rely on.
+//===----------------------------------------------------------------------===//
+
+namespace {
+using Info = ImutContainerInfo<int>;
+using Tree = ImutAVLTree<Info>;
+using TreeIter = Tree::iterator; // ImutAVLTreeInOrderIterator
+
+// Build an ImmutableSet from the given values (in the given insertion order),
+// optionally removing some afterwards, so trees of varied shape are produced.
+ImmutableSet<int> buildSet(ImmutableSet<int>::Factory &F, ArrayRef<int> ToAdd,
+ ArrayRef<int> ToRemove = {}) {
+ ImmutableSet<int> S = F.getEmptySet();
+ for (int V : ToAdd)
+ S = F.add(S, V);
+ for (int V : ToRemove)
+ S = F.remove(S, V);
+ return S;
+}
+
+// A representative collection of trees: degenerate, hand-picked small shapes,
+// a large balanced one, and many pseudo-random insert/remove mixes.
+std::vector<ImmutableSet<int>> makeTestSets(ImmutableSet<int>::Factory &F) {
+ std::vector<ImmutableSet<int>> Sets;
+ Sets.push_back(F.getEmptySet());
+ Sets.push_back(buildSet(F, {42}));
+ Sets.push_back(buildSet(F, {1, 2, 3}));
+ Sets.push_back(buildSet(F, {3, 2, 1}));
+ Sets.push_back(buildSet(F, {2, 1, 3}));
+
+ std::vector<int> Sorted(200);
+ std::iota(Sorted.begin(), Sorted.end(), 0);
+ Sets.push_back(buildSet(F, Sorted));
+
+ std::mt19937 Rng(12345);
+ for (int Trial = 0; Trial < 25; ++Trial) {
+ std::vector<int> Vals(150);
+ std::iota(Vals.begin(), Vals.end(), 0);
+ std::shuffle(Vals.begin(), Vals.end(), Rng);
+ std::vector<int> Removals(Vals.begin(), Vals.begin() + (Trial % 40));
+ Sets.push_back(buildSet(F, Vals, Removals));
+ }
+ return Sets;
+}
+} // namespace
+
+// Forward iteration must visit keys in ascending order, exactly matching an
+// independent std::set holding the same elements.
+TEST_F(ImmutableSetTest, IteratorInOrderMatchesStdSet) {
+ ImmutableSet<int>::Factory F(/*canonicalize=*/false);
+ for (const ImmutableSet<int> &S : makeTestSets(F)) {
+ std::set<int> Oracle;
+ for (ImmutableSet<int>::iterator I = S.begin(), E = S.end(); I != E; ++I)
+ Oracle.insert(*I);
+
+ std::vector<int> Forward;
+ for (ImmutableSet<int>::iterator I = S.begin(), E = S.end(); I != E; ++I)
+ Forward.push_back(*I);
+
+ EXPECT_TRUE(std::is_sorted(Forward.begin(), Forward.end()));
+ EXPECT_TRUE(std::adjacent_find(Forward.begin(), Forward.end()) ==
+ Forward.end()); // no duplicates
+ EXPECT_TRUE(std::equal(Forward.begin(), Forward.end(), Oracle.begin(),
+ Oracle.end()));
+ }
+}
+
+// Walking backwards with operator-- from the last element must reproduce the
+// reverse of the forward traversal.
+TEST_F(ImmutableSetTest, IteratorReverseMatchesForward) {
+ ImmutableSet<int>::Factory F(/*canonicalize=*/false);
+ for (const ImmutableSet<int> &S : makeTestSets(F)) {
+ const Tree *Root = S.getRootWithoutRetain();
+
+ std::vector<const Tree *> Forward;
+ for (TreeIter I(Root), E; I != E; ++I)
+ Forward.push_back(&*I);
+ if (Forward.empty())
+ continue;
+
+ // Advance to the last element, then walk backwards.
+ TreeIter Last(Root);
+ for (TreeIter I(Root), E; I != E; ++I)
+ Last = I;
+
+ std::vector<const Tree *> Backward;
+ Backward.push_back(&*Last);
+ for (TreeIter B = Last; Backward.size() < Forward.size();) {
+ --B;
+ Backward.push_back(&*B);
+ }
+ std::reverse(Backward.begin(), Backward.end());
+ EXPECT_EQ(Forward, Backward);
+ }
+}
+
+// skipSubTree must land on the in-order successor of the *entire* subtree
+// rooted at the current node. Since an in-order traversal visits a subtree as a
+// contiguous run, the destination index is computable independently from the
+// node's right-subtree size.
+TEST_F(ImmutableSetTest, IteratorSkipSubTree) {
+ ImmutableSet<int>::Factory F(/*canonicalize=*/false);
+ for (const ImmutableSet<int> &S : makeTestSets(F)) {
+ const Tree *Root = S.getRootWithoutRetain();
+
+ std::vector<const Tree *> Order;
+ for (TreeIter I(Root), E; I != E; ++I)
+ Order.push_back(&*I);
+
+ // Fro...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/205552
More information about the llvm-commits
mailing list