From 976edfedf735bd8928e84acc75d62d09fce3440f Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Fri, 23 Jun 2023 18:25:27 +0200 Subject: [PATCH] Add `Cursor::next_item` --- crates/sum_tree/src/cursor.rs | 36 +++++++++++++++++++++++++++++++++ crates/sum_tree/src/sum_tree.rs | 31 ++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/crates/sum_tree/src/cursor.rs b/crates/sum_tree/src/cursor.rs index b78484e6de..59165283f6 100644 --- a/crates/sum_tree/src/cursor.rs +++ b/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> { + 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() { diff --git a/crates/sum_tree/src/sum_tree.rs b/crates/sum_tree/src/sum_tree.rs index 577a942889..6ac0ef6350 100644 --- a/crates/sum_tree/src/sum_tree.rs +++ b/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::(); 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::(); @@ -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