Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

sum_tree: Store context on cursor (#34904)

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

173
crates/sum_tree/src/cursor.rs
View file

@ -25,6 +25,7 @@ pub struct Cursor<'a, T: Item, D> {
position: D,
did_seek: bool,
at_end: bool,
cx: &'a <T::Summary as Summary>::Context,
}
impl<T: Item + fmt::Debug, D: fmt::Debug> fmt::Debug for Cursor<'_, T, D>
@ -52,21 +53,22 @@ where
T: Item,
D: Dimension<'a, T::Summary>,
{
pub fn new(tree: &'a SumTree<T>, cx: &<T::Summary as Summary>::Context) -> Self {
pub fn new(tree: &'a SumTree<T>, cx: &'a <T::Summary as Summary>::Context) -> Self {
Self {
tree,
stack: ArrayVec::new(),
position: D::zero(cx),
did_seek: false,
at_end: tree.is_empty(),
cx,
}
}
fn reset(&mut self, cx: &<T::Summary as Summary>::Context) {
fn reset(&mut self) {
self.did_seek = false;
self.at_end = self.tree.is_empty();
self.stack.truncate(0);
self.position = D::zero(cx);
self.position = D::zero(self.cx);
}
pub fn start(&self) -> &D {
@ -74,10 +76,10 @@ where
}
#[track_caller]
pub fn end(&self, cx: &<T::Summary as Summary>::Context) -> D {
pub fn end(&self) -> D {
if let Some(item_summary) = self.item_summary() {
let mut end = self.start().clone();
end.add_summary(item_summary, cx);
end.add_summary(item_summary, self.cx);
end
} else {
self.start().clone()
@ -202,12 +204,12 @@ where
}
#[track_caller]
pub fn prev(&mut self, cx: &<T::Summary as Summary>::Context) {
self.search_backward(|_| true, cx)
pub fn prev(&mut self) {
self.search_backward(|_| true)
}
#[track_caller]
pub fn search_backward<F>(&mut self, mut filter_node: F, cx: &<T::Summary as Summary>::Context)
pub fn search_backward<F>(&mut self, mut filter_node: F)
where
F: FnMut(&T::Summary) -> bool,
{
@ -217,13 +219,13 @@ where
}
if self.at_end {
self.position = D::zero(cx);
self.position = D::zero(self.cx);
self.at_end = self.tree.is_empty();
if !self.tree.is_empty() {
self.stack.push(StackEntry {
tree: self.tree,
index: self.tree.0.child_summaries().len(),
position: D::from_summary(self.tree.summary(), cx),
position: D::from_summary(self.tree.summary(), self.cx),
});
}
}
@ -233,7 +235,7 @@ where
if let Some(StackEntry { position, .. }) = self.stack.iter().rev().nth(1) {
self.position = position.clone();
} else {
self.position = D::zero(cx);
self.position = D::zero(self.cx);
}
let entry = self.stack.last_mut().unwrap();
@ -247,7 +249,7 @@ where
}
for summary in &entry.tree.0.child_summaries()[..entry.index] {
self.position.add_summary(summary, cx);
self.position.add_summary(summary, self.cx);
}
entry.position = self.position.clone();
@ -257,7 +259,7 @@ where
if descending {
let tree = &child_trees[entry.index];
self.stack.push(StackEntry {
position: D::zero(cx),
position: D::zero(self.cx),
tree,
index: tree.0.child_summaries().len() - 1,
})
@ -273,12 +275,12 @@ where
}
#[track_caller]
pub fn next(&mut self, cx: &<T::Summary as Summary>::Context) {
self.search_forward(|_| true, cx)
pub fn next(&mut self) {
self.search_forward(|_| true)
}
#[track_caller]
pub fn search_forward<F>(&mut self, mut filter_node: F, cx: &<T::Summary as Summary>::Context)
pub fn search_forward<F>(&mut self, mut filter_node: F)
where
F: FnMut(&T::Summary) -> bool,
{
@ -289,7 +291,7 @@ where
self.stack.push(StackEntry {
tree: self.tree,
index: 0,
position: D::zero(cx),
position: D::zero(self.cx),
});
descend = true;
}
@ -316,8 +318,8 @@ where
break;
} else {
entry.index += 1;
entry.position.add_summary(next_summary, cx);
self.position.add_summary(next_summary, cx);
entry.position.add_summary(next_summary, self.cx);
self.position.add_summary(next_summary, self.cx);
}
}
@ -327,8 +329,8 @@ where
if !descend {
let item_summary = &item_summaries[entry.index];
entry.index += 1;
entry.position.add_summary(item_summary, cx);
self.position.add_summary(item_summary, cx);
entry.position.add_summary(item_summary, self.cx);
self.position.add_summary(item_summary, self.cx);
}
loop {
@ -337,8 +339,8 @@ where
return;
} else {
entry.index += 1;
entry.position.add_summary(next_item_summary, cx);
self.position.add_summary(next_item_summary, cx);
entry.position.add_summary(next_item_summary, self.cx);
self.position.add_summary(next_item_summary, self.cx);
}
} else {
break None;
@ -380,71 +382,51 @@ where
D: Dimension<'a, T::Summary>,
{
#[track_caller]
pub fn seek<Target>(
&mut self,
pos: &Target,
bias: Bias,
cx: &<T::Summary as Summary>::Context,
) -> bool
pub fn seek<Target>(&mut self, pos: &Target, bias: Bias) -> bool
where
Target: SeekTarget<'a, T::Summary, D>,
{
self.reset(cx);
self.seek_internal(pos, bias, &mut (), cx)
self.reset();
self.seek_internal(pos, bias, &mut ())
}
#[track_caller]
pub fn seek_forward<Target>(
&mut self,
pos: &Target,
bias: Bias,
cx: &<T::Summary as Summary>::Context,
) -> bool
pub fn seek_forward<Target>(&mut self, pos: &Target, bias: Bias) -> bool
where
Target: SeekTarget<'a, T::Summary, D>,
{
self.seek_internal(pos, bias, &mut (), cx)
self.seek_internal(pos, bias, &mut ())
}
/// Advances the cursor and returns traversed items as a tree.
#[track_caller]
pub fn slice<Target>(
&mut self,
end: &Target,
bias: Bias,
cx: &<T::Summary as Summary>::Context,
) -> SumTree<T>
pub fn slice<Target>(&mut self, end: &Target, bias: Bias) -> SumTree<T>
where
Target: SeekTarget<'a, T::Summary, D>,
{
let mut slice = SliceSeekAggregate {
tree: SumTree::new(cx),
tree: SumTree::new(self.cx),
leaf_items: ArrayVec::new(),
leaf_item_summaries: ArrayVec::new(),
leaf_summary: <T::Summary as Summary>::zero(cx),
leaf_summary: <T::Summary as Summary>::zero(self.cx),
};
self.seek_internal(end, bias, &mut slice, cx);
self.seek_internal(end, bias, &mut slice);
slice.tree
}
#[track_caller]
pub fn suffix(&mut self, cx: &<T::Summary as Summary>::Context) -> SumTree<T> {
self.slice(&End::new(), Bias::Right, cx)
pub fn suffix(&mut self) -> SumTree<T> {
self.slice(&End::new(), Bias::Right)
}
#[track_caller]
pub fn summary<Target, Output>(
&mut self,
end: &Target,
bias: Bias,
cx: &<T::Summary as Summary>::Context,
) -> Output
pub fn summary<Target, Output>(&mut self, end: &Target, bias: Bias) -> Output
where
Target: SeekTarget<'a, T::Summary, D>,
Output: Dimension<'a, T::Summary>,
{
let mut summary = SummarySeekAggregate(Output::zero(cx));
self.seek_internal(end, bias, &mut summary, cx);
let mut summary = SummarySeekAggregate(Output::zero(self.cx));
self.seek_internal(end, bias, &mut summary);
summary.0
}
@ -455,10 +437,9 @@ where
target: &dyn SeekTarget<'a, T::Summary, D>,
bias: Bias,
aggregate: &mut dyn SeekAggregate<'a, T>,
cx: &<T::Summary as Summary>::Context,
) -> bool {
assert!(
target.cmp(&self.position, cx) >= Ordering::Equal,
target.cmp(&self.position, self.cx) >= Ordering::Equal,
"cannot seek backward",
);
@ -467,7 +448,7 @@ where
self.stack.push(StackEntry {
tree: self.tree,
index: 0,
position: D::zero(cx),
position: D::zero(self.cx),
});
}
@ -489,14 +470,14 @@ where
.zip(&child_summaries[entry.index..])
{
let mut child_end = self.position.clone();
child_end.add_summary(child_summary, cx);
child_end.add_summary(child_summary, self.cx);
let comparison = target.cmp(&child_end, cx);
let comparison = target.cmp(&child_end, self.cx);
if comparison == Ordering::Greater
|| (comparison == Ordering::Equal && bias == Bias::Right)
{
self.position = child_end;
aggregate.push_tree(child_tree, child_summary, cx);
aggregate.push_tree(child_tree, child_summary, self.cx);
entry.index += 1;
entry.position = self.position.clone();
} else {
@ -522,22 +503,22 @@ where
.zip(&item_summaries[entry.index..])
{
let mut child_end = self.position.clone();
child_end.add_summary(item_summary, cx);
child_end.add_summary(item_summary, self.cx);
let comparison = target.cmp(&child_end, cx);
let comparison = target.cmp(&child_end, self.cx);
if comparison == Ordering::Greater
|| (comparison == Ordering::Equal && bias == Bias::Right)
{
self.position = child_end;
aggregate.push_item(item, item_summary, cx);
aggregate.push_item(item, item_summary, self.cx);
entry.index += 1;
} else {
aggregate.end_leaf(cx);
aggregate.end_leaf(self.cx);
break 'outer;
}
}
aggregate.end_leaf(cx);
aggregate.end_leaf(self.cx);
}
}
@ -551,11 +532,11 @@ where
let mut end = self.position.clone();
if bias == Bias::Left {
if let Some(summary) = self.item_summary() {
end.add_summary(summary, cx);
end.add_summary(summary, self.cx);
}
}
target.cmp(&end, cx) == Ordering::Equal
target.cmp(&end, self.cx) == Ordering::Equal
}
}
@ -624,21 +605,19 @@ impl<'a, T: Item> Iterator for Iter<'a, T> {
}
}
impl<'a, T, S, D> Iterator for Cursor<'a, T, D>
impl<'a, T: Item, D> Iterator for Cursor<'a, T, D>
where
T: Item<Summary = S>,
S: Summary<Context = ()>,
D: Dimension<'a, T::Summary>,
{
type Item = &'a T;
fn next(&mut self) -> Option<Self::Item> {
if !self.did_seek {
self.next(&());
self.next();
}
if let Some(item) = self.item() {
self.next(&());
self.next();
Some(item)
} else {
None
@ -651,7 +630,7 @@ pub struct FilterCursor<'a, F, T: Item, D> {
filter_node: F,
}
impl<'a, F, T, D> FilterCursor<'a, F, T, D>
impl<'a, F, T: Item, D> FilterCursor<'a, F, T, D>
where
F: FnMut(&T::Summary) -> bool,
T: Item,
@ -659,7 +638,7 @@ where
{
pub fn new(
tree: &'a SumTree<T>,
cx: &<T::Summary as Summary>::Context,
cx: &'a <T::Summary as Summary>::Context,
filter_node: F,
) -> Self {
let cursor = tree.cursor::<D>(cx);
@ -673,8 +652,8 @@ where
self.cursor.start()
}
pub fn end(&self, cx: &<T::Summary as Summary>::Context) -> D {
self.cursor.end(cx)
pub fn end(&self) -> D {
self.cursor.end()
}
pub fn item(&self) -> Option<&'a T> {
@ -685,31 +664,29 @@ where
self.cursor.item_summary()
}
pub fn next(&mut self, cx: &<T::Summary as Summary>::Context) {
self.cursor.search_forward(&mut self.filter_node, cx);
pub fn next(&mut self) {
self.cursor.search_forward(&mut self.filter_node);
}
pub fn prev(&mut self, cx: &<T::Summary as Summary>::Context) {
self.cursor.search_backward(&mut self.filter_node, cx);
pub fn prev(&mut self) {
self.cursor.search_backward(&mut self.filter_node);
}
}
impl<'a, F, T, S, U> Iterator for FilterCursor<'a, F, T, U>
impl<'a, F, T: Item, U> Iterator for FilterCursor<'a, F, T, U>
where
F: FnMut(&T::Summary) -> bool,
T: Item<Summary = S>,
S: Summary<Context = ()>, //Context for the summary must be unit type, as .next() doesn't take arguments
U: Dimension<'a, T::Summary>,
{
type Item = &'a T;
fn next(&mut self) -> Option<Self::Item> {
if !self.cursor.did_seek {
self.next(&());
self.next();
}
if let Some(item) = self.item() {
self.cursor.search_forward(&mut self.filter_node, &());
self.cursor.search_forward(&mut self.filter_node);
Some(item)
} else {
None
@ -795,3 +772,23 @@ where
self.0.add_summary(summary, cx);
}
}
struct End<D>(PhantomData<D>);
impl<D> End<D> {
fn new() -> Self {
Self(PhantomData)
}
}
impl<'a, S: Summary, D: Dimension<'a, S>> SeekTarget<'a, S, D> for End<D> {
fn cmp(&self, _: &D, _: &S::Context) -> Ordering {
Ordering::Greater
}
}
impl<D> fmt::Debug for End<D> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("End").finish()
}
}

222
crates/sum_tree/src/sum_tree.rs
View file

@ -38,7 +38,6 @@ pub trait Summary: Clone {
type Context;
fn zero(cx: &Self::Context) -> Self;
fn add_summary(&mut self, summary: &Self, cx: &Self::Context);
}
@ -138,26 +137,6 @@ where
}
}
struct End<D>(PhantomData<D>);
impl<D> End<D> {
fn new() -> Self {
Self(PhantomData)
}
}
impl<'a, S: Summary, D: Dimension<'a, S>> SeekTarget<'a, S, D> for End<D> {
fn cmp(&self, _: &D, _: &S::Context) -> Ordering {
Ordering::Greater
}
}
impl<D> fmt::Debug for End<D> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("End").finish()
}
}
/// Bias is used to settle ambiguities when determining positions in an ordered sequence.
///
/// The primary use case is for text, where Bias influences
@ -372,10 +351,10 @@ impl<T: Item> SumTree<T> {
pub fn items(&self, cx: &<T::Summary as Summary>::Context) -> Vec<T> {
let mut items = Vec::new();
let mut cursor = self.cursor::<()>(cx);
cursor.next(cx);
cursor.next();
while let Some(item) = cursor.item() {
items.push(item.clone());
cursor.next(cx);
cursor.next();
}
items
}
@ -384,7 +363,7 @@ impl<T: Item> SumTree<T> {
Iter::new(self)
}
pub fn cursor<'a, S>(&'a self, cx: &<T::Summary as Summary>::Context) -> Cursor<'a, T, S>
pub fn cursor<'a, S>(&'a self, cx: &'a <T::Summary as Summary>::Context) -> Cursor<'a, T, S>
where
S: Dimension<'a, T::Summary>,
{
@ -395,7 +374,7 @@ impl<T: Item> SumTree<T> {
/// that is returned cannot be used with Rust's iterators.
pub fn filter<'a, F, U>(
&'a self,
cx: &<T::Summary as Summary>::Context,
cx: &'a <T::Summary as Summary>::Context,
filter_node: F,
) -> FilterCursor<'a, F, T, U>
where
@ -525,10 +504,6 @@ impl<T: Item> SumTree<T> {
}
}
pub fn ptr_eq(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.0, &other.0)
}
fn push_tree_recursive(
&mut self,
other: SumTree<T>,
@ -686,11 +661,6 @@ impl<T: Item> SumTree<T> {
} => child_trees.last().unwrap().rightmost_leaf(),
}
}
#[cfg(debug_assertions)]
pub fn _debug_entries(&self) -> Vec<&T> {
self.iter().collect::<Vec<_>>()
}
}
impl<T: Item + PartialEq> PartialEq for SumTree<T> {
@ -710,15 +680,15 @@ impl<T: KeyedItem> SumTree<T> {
let mut replaced = None;
*self = {
let mut cursor = self.cursor::<T::Key>(cx);
let mut new_tree = cursor.slice(&item.key(), Bias::Left, cx);
let mut new_tree = cursor.slice(&item.key(), Bias::Left);
if let Some(cursor_item) = cursor.item() {
if cursor_item.key() == item.key() {
replaced = Some(cursor_item.clone());
cursor.next(cx);
cursor.next();
}
}
new_tree.push(item, cx);
new_tree.append(cursor.suffix(cx), cx);
new_tree.append(cursor.suffix(), cx);
new_tree
};
replaced
@ -728,14 +698,14 @@ impl<T: KeyedItem> SumTree<T> {
let mut removed = None;
*self = {
let mut cursor = self.cursor::<T::Key>(cx);
let mut new_tree = cursor.slice(key, Bias::Left, cx);
let mut new_tree = cursor.slice(key, Bias::Left);
if let Some(item) = cursor.item() {
if item.key() == *key {
removed = Some(item.clone());
cursor.next(cx);
cursor.next();
}
}
new_tree.append(cursor.suffix(cx), cx);
new_tree.append(cursor.suffix(), cx);
new_tree
};
removed
@ -758,7 +728,7 @@ impl<T: KeyedItem> SumTree<T> {
let mut new_tree = SumTree::new(cx);
let mut buffered_items = Vec::new();
cursor.seek(&T::Key::zero(cx), Bias::Left, cx);
cursor.seek(&T::Key::zero(cx), Bias::Left);
for edit in edits {
let new_key = edit.key();
let mut old_item = cursor.item();
@ -768,7 +738,7 @@ impl<T: KeyedItem> SumTree<T> {
.map_or(false, |old_item| old_item.key() < new_key)
{
new_tree.extend(buffered_items.drain(..), cx);
let slice = cursor.slice(&new_key, Bias::Left, cx);
let slice = cursor.slice(&new_key, Bias::Left);
new_tree.append(slice, cx);
old_item = cursor.item();
}
@ -776,7 +746,7 @@ impl<T: KeyedItem> SumTree<T> {
if let Some(old_item) = old_item {
if old_item.key() == new_key {
removed.push(old_item.clone());
cursor.next(cx);
cursor.next();
}
}
@ -789,70 +759,25 @@ impl<T: KeyedItem> SumTree<T> {
}
new_tree.extend(buffered_items, cx);
new_tree.append(cursor.suffix(cx), cx);
new_tree.append(cursor.suffix(), cx);
new_tree
};
removed
}
pub fn get(&self, key: &T::Key, cx: &<T::Summary as Summary>::Context) -> Option<&T> {
pub fn get<'a>(
&'a self,
key: &T::Key,
cx: &'a <T::Summary as Summary>::Context,
) -> Option<&'a T> {
let mut cursor = self.cursor::<T::Key>(cx);
if cursor.seek(key, Bias::Left, cx) {
if cursor.seek(key, Bias::Left) {
cursor.item()
} else {
None
}
}
#[inline]
pub fn contains(&self, key: &T::Key, cx: &<T::Summary as Summary>::Context) -> bool {
self.get(key, cx).is_some()
}
pub fn update<F, R>(
&mut self,
key: &T::Key,
cx: &<T::Summary as Summary>::Context,
f: F,
) -> Option<R>
where
F: FnOnce(&mut T) -> R,
{
let mut cursor = self.cursor::<T::Key>(cx);
let mut new_tree = cursor.slice(key, Bias::Left, cx);
let mut result = None;
if Ord::cmp(key, &cursor.end(cx)) == Ordering::Equal {
let mut updated = cursor.item().unwrap().clone();
result = Some(f(&mut updated));
new_tree.push(updated, cx);
cursor.next(cx);
}
new_tree.append(cursor.suffix(cx), cx);
drop(cursor);
*self = new_tree;
result
}
pub fn retain<F: FnMut(&T) -> bool>(
&mut self,
cx: &<T::Summary as Summary>::Context,
mut predicate: F,
) {
let mut new_map = SumTree::new(cx);
let mut cursor = self.cursor::<T::Key>(cx);
cursor.next(cx);
while let Some(item) = cursor.item() {
if predicate(&item) {
new_map.push(item.clone(), cx);
}
cursor.next(cx);
}
drop(cursor);
*self = new_map;
}
}
impl<T, S> Default for SumTree<T>
@ -1061,14 +986,14 @@ mod tests {
tree = {
let mut cursor = tree.cursor::<Count>(&());
let mut new_tree = cursor.slice(&Count(splice_start), Bias::Right, &());
let mut new_tree = cursor.slice(&Count(splice_start), Bias::Right);
if rng.r#gen() {
new_tree.extend(new_items, &());
} else {
new_tree.par_extend(new_items, &());
}
cursor.seek(&Count(splice_end), Bias::Right, &());
new_tree.append(cursor.slice(&tree_end, Bias::Right, &()), &());
cursor.seek(&Count(splice_end), Bias::Right);
new_tree.append(cursor.slice(&tree_end, Bias::Right), &());
new_tree
};
@ -1090,10 +1015,10 @@ mod tests {
.collect::<Vec<_>>();
let mut item_ix = if rng.r#gen() {
filter_cursor.next(&());
filter_cursor.next();
0
} else {
filter_cursor.prev(&());
filter_cursor.prev();
expected_filtered_items.len().saturating_sub(1)
};
while item_ix < expected_filtered_items.len() {
@ -1103,19 +1028,19 @@ mod tests {
assert_eq!(actual_item, &reference_item);
assert_eq!(filter_cursor.start().0, reference_index);
log::info!("next");
filter_cursor.next(&());
filter_cursor.next();
item_ix += 1;
while item_ix > 0 && rng.gen_bool(0.2) {
log::info!("prev");
filter_cursor.prev(&());
filter_cursor.prev();
item_ix -= 1;
if item_ix == 0 && rng.gen_bool(0.2) {
filter_cursor.prev(&());
filter_cursor.prev();
assert_eq!(filter_cursor.item(), None);
assert_eq!(filter_cursor.start().0, 0);
filter_cursor.next(&());
filter_cursor.next();
}
}
}
@ -1124,9 +1049,9 @@ mod tests {
let mut before_start = false;
let mut cursor = tree.cursor::<Count>(&());
let start_pos = rng.gen_range(0..=reference_items.len());
cursor.seek(&Count(start_pos), Bias::Right, &());
cursor.seek(&Count(start_pos), Bias::Right);
let mut pos = rng.gen_range(start_pos..=reference_items.len());
cursor.seek_forward(&Count(pos), Bias::Right, &());
cursor.seek_forward(&Count(pos), Bias::Right);
for i in 0..10 {
assert_eq!(cursor.start().0, pos);
@ -1152,13 +1077,13 @@ mod tests {
}
if i < 5 {
cursor.next(&());
cursor.next();
if pos < reference_items.len() {
pos += 1;
before_start = false;
}
} else {
cursor.prev(&());
cursor.prev();
if pos == 0 {
before_start = true;
}
@ -1174,11 +1099,11 @@ mod tests {
let end_bias = if rng.r#gen() { Bias::Left } else { Bias::Right };
let mut cursor = tree.cursor::<Count>(&());
cursor.seek(&Count(start), start_bias, &());
let slice = cursor.slice(&Count(end), end_bias, &());
cursor.seek(&Count(start), start_bias);
let slice = cursor.slice(&Count(end), end_bias);
cursor.seek(&Count(start), start_bias, &());
let summary = cursor.summary::<_, Sum>(&Count(end), end_bias, &());
cursor.seek(&Count(start), start_bias);
let summary = cursor.summary::<_, Sum>(&Count(end), end_bias);
assert_eq!(summary.0, slice.summary().sum);
}
@ -1191,19 +1116,19 @@ mod tests {
let tree = SumTree::<u8>::default();
let mut cursor = tree.cursor::<IntegersSummary>(&());
assert_eq!(
cursor.slice(&Count(0), Bias::Right, &()).items(&()),
cursor.slice(&Count(0), Bias::Right).items(&()),
Vec::<u8>::new()
);
assert_eq!(cursor.item(), None);
assert_eq!(cursor.prev_item(), None);
assert_eq!(cursor.next_item(), None);
assert_eq!(cursor.start().sum, 0);
cursor.prev(&());
cursor.prev();
assert_eq!(cursor.item(), None);
assert_eq!(cursor.prev_item(), None);
assert_eq!(cursor.next_item(), None);
assert_eq!(cursor.start().sum, 0);
cursor.next(&());
cursor.next();
assert_eq!(cursor.item(), None);
assert_eq!(cursor.prev_item(), None);
assert_eq!(cursor.next_item(), None);
@ -1214,7 +1139,7 @@ mod tests {
tree.extend(vec![1], &());
let mut cursor = tree.cursor::<IntegersSummary>(&());
assert_eq!(
cursor.slice(&Count(0), Bias::Right, &()).items(&()),
cursor.slice(&Count(0), Bias::Right).items(&()),
Vec::<u8>::new()
);
assert_eq!(cursor.item(), Some(&1));
@ -1222,29 +1147,29 @@ mod tests {
assert_eq!(cursor.next_item(), None);
assert_eq!(cursor.start().sum, 0);
cursor.next(&());
cursor.next();
assert_eq!(cursor.item(), None);
assert_eq!(cursor.prev_item(), Some(&1));
assert_eq!(cursor.next_item(), None);
assert_eq!(cursor.start().sum, 1);
cursor.prev(&());
cursor.prev();
assert_eq!(cursor.item(), Some(&1));
assert_eq!(cursor.prev_item(), None);
assert_eq!(cursor.next_item(), None);
assert_eq!(cursor.start().sum, 0);
let mut cursor = tree.cursor::<IntegersSummary>(&());
assert_eq!(cursor.slice(&Count(1), Bias::Right, &()).items(&()), [1]);
assert_eq!(cursor.slice(&Count(1), Bias::Right).items(&()), [1]);
assert_eq!(cursor.item(), None);
assert_eq!(cursor.prev_item(), Some(&1));
assert_eq!(cursor.next_item(), None);
assert_eq!(cursor.start().sum, 1);
cursor.seek(&Count(0), Bias::Right, &());
cursor.seek(&Count(0), Bias::Right);
assert_eq!(
cursor
.slice(&tree.extent::<Count>(&()), Bias::Right, &())
.slice(&tree.extent::<Count>(&()), Bias::Right)
.items(&()),
[1]
);
@ -1258,80 +1183,80 @@ mod tests {
tree.extend(vec![1, 2, 3, 4, 5, 6], &());
let mut cursor = tree.cursor::<IntegersSummary>(&());
assert_eq!(cursor.slice(&Count(2), Bias::Right, &()).items(&()), [1, 2]);
assert_eq!(cursor.slice(&Count(2), Bias::Right).items(&()), [1, 2]);
assert_eq!(cursor.item(), Some(&3));
assert_eq!(cursor.prev_item(), Some(&2));
assert_eq!(cursor.next_item(), Some(&4));
assert_eq!(cursor.start().sum, 3);
cursor.next(&());
cursor.next();
assert_eq!(cursor.item(), Some(&4));
assert_eq!(cursor.prev_item(), Some(&3));
assert_eq!(cursor.next_item(), Some(&5));
assert_eq!(cursor.start().sum, 6);
cursor.next(&());
cursor.next();
assert_eq!(cursor.item(), Some(&5));
assert_eq!(cursor.prev_item(), Some(&4));
assert_eq!(cursor.next_item(), Some(&6));
assert_eq!(cursor.start().sum, 10);
cursor.next(&());
cursor.next();
assert_eq!(cursor.item(), Some(&6));
assert_eq!(cursor.prev_item(), Some(&5));
assert_eq!(cursor.next_item(), None);
assert_eq!(cursor.start().sum, 15);
cursor.next(&());
cursor.next(&());
cursor.next();
cursor.next();
assert_eq!(cursor.item(), None);
assert_eq!(cursor.prev_item(), Some(&6));
assert_eq!(cursor.next_item(), None);
assert_eq!(cursor.start().sum, 21);
cursor.prev(&());
cursor.prev();
assert_eq!(cursor.item(), Some(&6));
assert_eq!(cursor.prev_item(), Some(&5));
assert_eq!(cursor.next_item(), None);
assert_eq!(cursor.start().sum, 15);
cursor.prev(&());
cursor.prev();
assert_eq!(cursor.item(), Some(&5));
assert_eq!(cursor.prev_item(), Some(&4));
assert_eq!(cursor.next_item(), Some(&6));
assert_eq!(cursor.start().sum, 10);
cursor.prev(&());
cursor.prev();
assert_eq!(cursor.item(), Some(&4));
assert_eq!(cursor.prev_item(), Some(&3));
assert_eq!(cursor.next_item(), Some(&5));
assert_eq!(cursor.start().sum, 6);
cursor.prev(&());
cursor.prev();
assert_eq!(cursor.item(), Some(&3));
assert_eq!(cursor.prev_item(), Some(&2));
assert_eq!(cursor.next_item(), Some(&4));
assert_eq!(cursor.start().sum, 3);
cursor.prev(&());
cursor.prev();
assert_eq!(cursor.item(), Some(&2));
assert_eq!(cursor.prev_item(), Some(&1));
assert_eq!(cursor.next_item(), Some(&3));
assert_eq!(cursor.start().sum, 1);
cursor.prev(&());
cursor.prev();
assert_eq!(cursor.item(), Some(&1));
assert_eq!(cursor.prev_item(), None);
assert_eq!(cursor.next_item(), Some(&2));
assert_eq!(cursor.start().sum, 0);
cursor.prev(&());
cursor.prev();
assert_eq!(cursor.item(), None);
assert_eq!(cursor.prev_item(), None);
assert_eq!(cursor.next_item(), Some(&1));
assert_eq!(cursor.start().sum, 0);
cursor.next(&());
cursor.next();
assert_eq!(cursor.item(), Some(&1));
assert_eq!(cursor.prev_item(), None);
assert_eq!(cursor.next_item(), Some(&2));
@ -1340,7 +1265,7 @@ mod tests {
let mut cursor = tree.cursor::<IntegersSummary>(&());
assert_eq!(
cursor
.slice(&tree.extent::<Count>(&()), Bias::Right, &())
.slice(&tree.extent::<Count>(&()), Bias::Right)
.items(&()),
tree.items(&())
);
@ -1349,10 +1274,10 @@ mod tests {
assert_eq!(cursor.next_item(), None);
assert_eq!(cursor.start().sum, 21);
cursor.seek(&Count(3), Bias::Right, &());
cursor.seek(&Count(3), Bias::Right);
assert_eq!(
cursor
.slice(&tree.extent::<Count>(&()), Bias::Right, &())
.slice(&tree.extent::<Count>(&()), Bias::Right)
.items(&()),
[4, 5, 6]
);
@ -1362,25 +1287,16 @@ mod tests {
assert_eq!(cursor.start().sum, 21);
// Seeking can bias left or right
cursor.seek(&Count(1), Bias::Left, &());
cursor.seek(&Count(1), Bias::Left);
assert_eq!(cursor.item(), Some(&1));
cursor.seek(&Count(1), Bias::Right, &());
cursor.seek(&Count(1), Bias::Right);
assert_eq!(cursor.item(), Some(&2));
// Slicing without resetting starts from where the cursor is parked at.
cursor.seek(&Count(1), Bias::Right, &());
assert_eq!(
cursor.slice(&Count(3), Bias::Right, &()).items(&()),
vec![2, 3]
);
assert_eq!(
cursor.slice(&Count(6), Bias::Left, &()).items(&()),
vec![4, 5]
);
assert_eq!(
cursor.slice(&Count(6), Bias::Right, &()).items(&()),
vec![6]
);
cursor.seek(&Count(1), Bias::Right);
assert_eq!(cursor.slice(&Count(3), Bias::Right).items(&()), vec![2, 3]);
assert_eq!(cursor.slice(&Count(6), Bias::Left).items(&()), vec![4, 5]);
assert_eq!(cursor.slice(&Count(6), Bias::Right).items(&()), vec![6]);
}
#[test]

34
crates/sum_tree/src/tree_map.rs
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);