refactor: Unify upperBound and lowerBound into boundHelper (#7943)

Co-authored-by: Bart <11445373+bthomee@users.noreply.github.com>
Co-authored-by: vlntb <13349202+vlntb@users.noreply.github.com>
This commit is contained in:
Bart
2026-09-24 12:48:25 +00:00
committed by GitHub
parent 30640a626f
commit 1d7669f528
3 changed files with 88 additions and 44 deletions

View File

@@ -592,6 +592,20 @@ private:
SHAMapLeafNode*
belowHelper(NodePathStack& stack, BelowDirection direction) const;
/**
* Returns the nearest item strictly past `id`, in the given direction.
*
* Walks back up the path to `id`. At each inner node the branches beyond the one `id` takes
* hold the candidates, so the first non-empty one is the closest and the extreme leaf below
* it is the answer.
*
* @param id The key to search from, which need not be in the map.
* @param direction First to search upwards from `id`, Last to search downwards.
* @return An iterator to the item found, or end() if no item lies on that side of `id`.
*/
ConstIterator
boundHelper(uint256 const& id, BelowDirection direction) const;
// Simple descent
// Get a child of the specified node
SHAMapTreeNode*

View File

@@ -575,8 +575,10 @@ SHAMap::peekItem(uint256 const& id, SHAMapHash& hash) const
}
SHAMap::ConstIterator
SHAMap::upperBound(uint256 const& id) const
SHAMap::boundHelper(uint256 const& id, BelowDirection direction) const
{
auto const searchingForward = direction == BelowDirection::First;
NodePathStack stack;
walkTowardsKey(id, &stack);
while (!stack.empty())
@@ -584,63 +586,45 @@ SHAMap::upperBound(uint256 const& id) const
auto const [node, nodeID] = stack.top();
if (node->isLeaf())
{
auto leaf = safeDowncast<SHAMapLeafNode*>(node.get());
if (leaf->peekItem()->key() > id)
return ConstIterator(this, leaf->peekItem().get(), std::move(stack));
auto const& item = safeDowncast<SHAMapLeafNode const&>(*node).peekItem();
if (searchingForward ? (item->key() > id) : (item->key() < id))
return ConstIterator(this, item.get(), std::move(stack));
}
else
{
auto& inner = safeDowncast<SHAMapInnerNode&>(*node);
for (auto branch = selectBranch(nodeID, id) + 1; branch < kBranchFactor; ++branch)
auto const taken = selectBranch(nodeID, id);
auto const remaining = searchingForward ? (kBranchFactor - 1u - taken) : taken;
for (auto scanned = 0u; scanned < remaining; ++scanned)
{
if (!inner.isEmptyBranch(branch))
{
stack.pushChild(descendThrow(inner, branch), branch);
auto leaf = belowHelper(stack, BelowDirection::First);
if (leaf == nullptr)
Throw<SHAMapMissingNode>(type_, id);
return ConstIterator(this, leaf->peekItem().get(), std::move(stack));
}
auto const branch =
searchingForward ? (taken + 1u + scanned) : (taken - 1u - scanned);
if (inner.isEmptyBranch(branch))
continue;
stack.pushChild(descendThrow(inner, branch), branch);
auto const leaf = belowHelper(stack, direction);
if (leaf == nullptr)
Throw<SHAMapMissingNode>(type_, id);
return ConstIterator(this, leaf->peekItem().get(), std::move(stack));
}
}
stack.pop();
}
return end();
}
SHAMap::ConstIterator
SHAMap::upperBound(uint256 const& id) const
{
return boundHelper(id, BelowDirection::First);
}
SHAMap::ConstIterator
SHAMap::lowerBound(uint256 const& id) const
{
NodePathStack stack;
walkTowardsKey(id, &stack);
while (!stack.empty())
{
auto const [node, nodeID] = stack.top();
if (node->isLeaf())
{
auto leaf = safeDowncast<SHAMapLeafNode*>(node.get());
if (leaf->peekItem()->key() < id)
return ConstIterator(this, leaf->peekItem().get(), std::move(stack));
}
else
{
auto& inner = safeDowncast<SHAMapInnerNode&>(*node);
for (auto branch = selectBranch(nodeID, id); branch > 0u;)
{
--branch;
if (!inner.isEmptyBranch(branch))
{
stack.pushChild(descendThrow(inner, branch), branch);
auto leaf = belowHelper(stack, BelowDirection::Last);
if (leaf == nullptr)
Throw<SHAMapMissingNode>(type_, id);
return ConstIterator(this, leaf->peekItem().get(), std::move(stack));
}
}
}
stack.pop();
}
// TODO: what to return here?
return end();
return boundHelper(id, BelowDirection::Last);
}
bool

View File

@@ -455,6 +455,52 @@ TEST_F(SHAMapTraversal, bounds_agree_with_iteration_for_absent_keys)
}
}
TEST_F(SHAMapTraversal, bounds_on_empty_map_return_end)
{
tests::TestNodeFamily f{j_};
SHAMap map{SHAMapType::FREE, f};
map.setUnbacked();
// The root is a childless inner node, so boundHelper's inner-node branch scans every branch on
// the requested side of the one id selects, finds them all empty, and falls through to end()
// rather than dereference a child.
EXPECT_EQ(map.upperBound(uint256{}), map.end());
EXPECT_EQ(map.lowerBound(uint256{}), map.end());
uint256 probe;
std::fill_n(probe.begin(), probe.size(), std::uint8_t{0xff});
EXPECT_EQ(map.upperBound(probe), map.end());
EXPECT_EQ(map.lowerBound(probe), map.end());
}
TEST_F(SHAMapTraversal, bounds_on_single_item_map_use_the_leaf_below_the_root)
{
tests::TestNodeFamily f{j_};
SHAMap map{SHAMapType::FREE, f};
auto const key = deepFanOutKeys().front();
fillMap(map, {key});
// fillMap adds items in-process, so root_ stays an inner node with the single leaf below it.
// The stack holds both, so boundHelper examines the leaf first.
uint256 below = key;
--below;
uint256 above = key;
++above;
auto const upper = map.upperBound(below);
ASSERT_NE(upper, map.end());
EXPECT_EQ(upper->key(), key);
EXPECT_EQ(map.upperBound(key), map.end());
EXPECT_EQ(map.upperBound(above), map.end());
auto const lower = map.lowerBound(above);
ASSERT_NE(lower, map.end());
EXPECT_EQ(lower->key(), key);
EXPECT_EQ(map.lowerBound(key), map.end());
EXPECT_EQ(map.lowerBound(below), map.end());
}
TEST_F(SHAMapTraversal, iteration_survives_deletions)
{
tests::TestNodeFamily f{j_};