Add `Cursor::next_item`

Antonio Scandurra created 2 years ago

Change summary

crates/sum_tree/src/cursor.rs   | 36 +++++++++++++++++++++++++++++++++++
crates/sum_tree/src/sum_tree.rs | 31 ++++++++++++++++++++++++++++++
2 files changed, 67 insertions(+)

Detailed changes

crates/sum_tree/src/cursor.rs

@@ -97,6 +97,42 @@ where
         }
     }
 
+    pub fn next_item(&self) -> Option<&'a T> {
+        self.assert_did_seek();
+        if let Some(entry) = self.stack.last() {
+            if entry.index == entry.tree.0.items().len() - 1 {
+                if let Some(next_leaf) = self.next_leaf() {
+                    Some(next_leaf.0.items().first().unwrap())
+                } else {
+                    None
+                }
+            } else {
+                match *entry.tree.0 {
+                    Node::Leaf { ref items, .. } => Some(&items[entry.index + 1]),
+                    _ => unreachable!(),
+                }
+            }
+        } else if self.at_end {
+            None
+        } else {
+            self.tree.first()
+        }
+    }
+
+    fn next_leaf(&self) -> Option<&'a SumTree<T>> {
+        for entry in self.stack.iter().rev().skip(1) {
+            if entry.index < entry.tree.0.child_trees().len() - 1 {
+                match *entry.tree.0 {
+                    Node::Internal {
+                        ref child_trees, ..
+                    } => return Some(child_trees[entry.index + 1].leftmost_leaf()),
+                    Node::Leaf { .. } => unreachable!(),
+                };
+            }
+        }
+        None
+    }
+
     pub fn prev_item(&self) -> Option<&'a T> {
         self.assert_did_seek();
         if let Some(entry) = self.stack.last() {

crates/sum_tree/src/sum_tree.rs

@@ -816,6 +816,14 @@ mod tests {
                         assert_eq!(cursor.item(), None);
                     }
 
+                    if before_start {
+                        assert_eq!(cursor.next_item(), reference_items.get(0));
+                    } else if pos + 1 < reference_items.len() {
+                        assert_eq!(cursor.next_item().unwrap(), &reference_items[pos + 1]);
+                    } else {
+                        assert_eq!(cursor.next_item(), None);
+                    }
+
                     if i < 5 {
                         cursor.next(&());
                         if pos < reference_items.len() {
@@ -861,14 +869,17 @@ mod tests {
         );
         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(&());
         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(&());
         assert_eq!(cursor.item(), None);
         assert_eq!(cursor.prev_item(), None);
+        assert_eq!(cursor.next_item(), None);
         assert_eq!(cursor.start().sum, 0);
 
         // Single-element tree
@@ -881,22 +892,26 @@ mod tests {
         );
         assert_eq!(cursor.item(), Some(&1));
         assert_eq!(cursor.prev_item(), None);
+        assert_eq!(cursor.next_item(), None);
         assert_eq!(cursor.start().sum, 0);
 
         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(&());
         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.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, &());
@@ -908,6 +923,7 @@ mod tests {
         );
         assert_eq!(cursor.item(), None);
         assert_eq!(cursor.prev_item(), Some(&1));
+        assert_eq!(cursor.next_item(), None);
         assert_eq!(cursor.start().sum, 1);
 
         // Multiple-element tree
@@ -918,67 +934,80 @@ mod tests {
         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(&());
         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(&());
         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(&());
         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(&());
         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(&());
         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(&());
         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(&());
         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(&());
         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(&());
         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(&());
         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(&());
         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(&());
         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);
 
         let mut cursor = tree.cursor::<IntegersSummary>();
@@ -990,6 +1019,7 @@ mod tests {
         );
         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.seek(&Count(3), Bias::Right, &());
@@ -1001,6 +1031,7 @@ mod tests {
         );
         assert_eq!(cursor.item(), None);
         assert_eq!(cursor.prev_item(), Some(&6));
+        assert_eq!(cursor.next_item(), None);
         assert_eq!(cursor.start().sum, 21);
 
         // Seeking can bias left or right