diff --git a/crates/sum_tree/src/cursor.rs b/crates/sum_tree/src/cursor.rs index fff675597c..7048389d68 100644 --- a/crates/sum_tree/src/cursor.rs +++ b/crates/sum_tree/src/cursor.rs @@ -205,7 +205,7 @@ where #[track_caller] pub fn prev(&mut self) { - self.search_backward(|_| Ordering::Greater) + self.search_backward(|_| Ordering::Equal) } #[track_caller] @@ -283,7 +283,7 @@ where #[track_caller] pub fn next(&mut self) { - self.search_forward(|_| Ordering::Less) + self.search_forward(|_| Ordering::Equal) } #[track_caller] @@ -322,25 +322,18 @@ where if entry.index < child_summaries.len() { let index = child_summaries[entry.index..] .partition_point(|item| filter_node(item).is_lt()); - if index < child_summaries.len() - entry.index { - entry.index += index; - } - let position = Some(entry.index) - .filter(|index| *index < child_summaries.len()) - .unwrap_or(child_summaries.len()); + entry.index += index; - if let Some(summary) = child_summaries.get(position) { + if let Some(summary) = child_summaries.get(entry.index) { entry.position.add_summary(summary, self.cx); self.position.add_summary(summary, self.cx); } } - dbg!((entry.index, child_trees.len())); child_trees.get(entry.index) } Node::Leaf { item_summaries, .. } => { - dbg!("Ayo"); if !descend { let item_summary = &item_summaries[entry.index]; entry.index += 1; @@ -351,15 +344,10 @@ where if entry.index < item_summaries.len() { let index = item_summaries[entry.index..] .partition_point(|item| filter_node(item).is_lt()); - if index < item_summaries.len() - entry.index { - entry.index += index; - } - entry.index += index; - let position = Some(entry.index) - .filter(|index| *index < item_summaries.len()) - .unwrap_or(item_summaries.len()); - if let Some(summary) = item_summaries.get(position) { + entry.index += index; + + if let Some(summary) = item_summaries.get(entry.index) { entry.position.add_summary(summary, self.cx); self.position.add_summary(summary, self.cx); } diff --git a/crates/sum_tree/src/sum_tree.rs b/crates/sum_tree/src/sum_tree.rs index 74fd65ab37..a0074aa2bc 100644 --- a/crates/sum_tree/src/sum_tree.rs +++ b/crates/sum_tree/src/sum_tree.rs @@ -5,6 +5,7 @@ use arrayvec::ArrayVec; pub use cursor::{Cursor, FilterCursor, Iter}; use rayon::prelude::*; use std::borrow::Cow; +use std::fmt::Debug; use std::marker::PhantomData; use std::mem; use std::{cmp::Ordering, fmt, iter::FromIterator, sync::Arc}; @@ -54,7 +55,6 @@ impl Summary for &'static () { } fn add_summary(&mut self, _: &Self, _: &()) {} - // fn sub_summary(&mut self, _: &Self, _: &()) {} } /// Each [`Summary`] type can have more than one [`Dimension`] type that it measures. @@ -214,14 +214,20 @@ impl SumTree { let mut iter = iter.into_iter().fuse().peekable(); while iter.peek().is_some() { let items: ArrayVec = iter.by_ref().take(2 * TREE_BASE).collect(); - let item_summaries: ArrayVec = - items.iter().map(|item| item.summary(cx)).collect(); - - let mut summary = item_summaries[0].clone(); - for item_summary in &item_summaries[1..] { - ::add_summary(&mut summary, item_summary, cx); - } - + let item_summaries: ArrayVec = items + .iter() + .scan(None, |previous_item, item| { + let summary = item.summary(cx); + let current_item = if let Some(mut base) = previous_item.take() { + ::add_summary(&mut base, &summary, cx); + base + } else { + summary + }; + _ = previous_item.insert(current_item.clone()); + Some(current_item) + }) + .collect(); nodes.push(Node::Leaf { items, item_summaries, @@ -1189,9 +1195,11 @@ mod tests { // Multiple-element tree let mut tree = SumTree::default(); tree.extend(vec![1, 2, 3, 4, 5, 6], &()); + let mut cursor = tree.cursor::(&()); - assert_eq!(cursor.slice(&Count(2), Bias::Right).items(&()), [1, 2]); + let slice = cursor.slice(&Count(2), Bias::Right); + assert_eq!(slice.items(&()), [1, 2]); assert_eq!(cursor.item(), Some(&3)); assert_eq!(cursor.prev_item(), Some(&2)); assert_eq!(cursor.next_item(), Some(&4)); diff --git a/crates/sum_tree/src/tree_map.rs b/crates/sum_tree/src/tree_map.rs index fcf6da6cd6..c966c1050e 100644 --- a/crates/sum_tree/src/tree_map.rs +++ b/crates/sum_tree/src/tree_map.rs @@ -406,9 +406,9 @@ mod tests { map.insert("baa", 3); map.insert("baaab", 4); map.insert("c", 5); - dbg!(&map); - let result = map - .iter_from(&"ba") + + let items = map.iter_from(&"ba"); + let result = items .take_while(|(key, _)| key.starts_with("ba")) .collect::>();