sum_tree: Store context on cursor (#34904)

This gets rid of the need to pass context to all cursor functions. In
practice context is always immutable when interacting with cursors.

A nicety of this is in the follow-up PR we will be able to implement
Iterator for all Cursors/filter cursors (hell, we may be able to get rid
of filter cursor altogether, as it is just a custom `filter` impl on
iterator trait).
Release Notes:

- N/A
This commit is contained in:
Piotr Osiewicz 2025-07-22 18:20:48 +02:00 committed by GitHub
parent fa3e1ccc37
commit 64d0fec699
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 749 additions and 876 deletions

View file

@ -54,7 +54,7 @@ impl<K: Clone + Ord, V: Clone> TreeMap<K, V> {
pub fn get(&self, key: &K) -> Option<&V> {
let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(&());
cursor.seek(&MapKeyRef(Some(key)), Bias::Left, &());
cursor.seek(&MapKeyRef(Some(key)), Bias::Left);
if let Some(item) = cursor.item() {
if Some(key) == item.key().0.as_ref() {
Some(&item.value)
@ -86,12 +86,12 @@ impl<K: Clone + Ord, V: Clone> TreeMap<K, V> {
let mut removed = None;
let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(&());
let key = MapKeyRef(Some(key));
let mut new_tree = cursor.slice(&key, Bias::Left, &());
if key.cmp(&cursor.end(&()), &()) == Ordering::Equal {
let mut new_tree = cursor.slice(&key, Bias::Left);
if key.cmp(&cursor.end(), &()) == Ordering::Equal {
removed = Some(cursor.item().unwrap().value.clone());
cursor.next(&());
cursor.next();
}
new_tree.append(cursor.suffix(&()), &());
new_tree.append(cursor.suffix(), &());
drop(cursor);
self.0 = new_tree;
removed
@ -101,9 +101,9 @@ impl<K: Clone + Ord, V: Clone> TreeMap<K, V> {
let start = MapSeekTargetAdaptor(start);
let end = MapSeekTargetAdaptor(end);
let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(&());
let mut new_tree = cursor.slice(&start, Bias::Left, &());
cursor.seek(&end, Bias::Left, &());
new_tree.append(cursor.suffix(&()), &());
let mut new_tree = cursor.slice(&start, Bias::Left);
cursor.seek(&end, Bias::Left);
new_tree.append(cursor.suffix(), &());
drop(cursor);
self.0 = new_tree;
}
@ -112,15 +112,15 @@ impl<K: Clone + Ord, V: Clone> TreeMap<K, V> {
pub fn closest(&self, key: &K) -> Option<(&K, &V)> {
let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(&());
let key = MapKeyRef(Some(key));
cursor.seek(&key, Bias::Right, &());
cursor.prev(&());
cursor.seek(&key, Bias::Right);
cursor.prev();
cursor.item().map(|item| (&item.key, &item.value))
}
pub fn iter_from<'a>(&'a self, from: &K) -> impl Iterator<Item = (&'a K, &'a V)> + 'a {
let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(&());
let from_key = MapKeyRef(Some(from));
cursor.seek(&from_key, Bias::Left, &());
cursor.seek(&from_key, Bias::Left);
cursor.map(|map_entry| (&map_entry.key, &map_entry.value))
}
@ -131,15 +131,15 @@ impl<K: Clone + Ord, V: Clone> TreeMap<K, V> {
{
let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(&());
let key = MapKeyRef(Some(key));
let mut new_tree = cursor.slice(&key, Bias::Left, &());
let mut new_tree = cursor.slice(&key, Bias::Left);
let mut result = None;
if key.cmp(&cursor.end(&()), &()) == Ordering::Equal {
if key.cmp(&cursor.end(), &()) == Ordering::Equal {
let mut updated = cursor.item().unwrap().clone();
result = Some(f(&mut updated.value));
new_tree.push(updated, &());
cursor.next(&());
cursor.next();
}
new_tree.append(cursor.suffix(&()), &());
new_tree.append(cursor.suffix(), &());
drop(cursor);
self.0 = new_tree;
result
@ -149,12 +149,12 @@ impl<K: Clone + Ord, V: Clone> TreeMap<K, V> {
let mut new_map = SumTree::<MapEntry<K, V>>::default();
let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(&());
cursor.next(&());
cursor.next();
while let Some(item) = cursor.item() {
if predicate(&item.key, &item.value) {
new_map.push(item.clone(), &());
}
cursor.next(&());
cursor.next();
}
drop(cursor);