Fix btree::Node::path_{next/prev}

bodil · Apr 29, 2022 · 41d9972 · 41d9972
1 parent cf08eec
commit 41d9972
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 9 deletions.
diff --git a/src/nodes/btree.rs b/src/nodes/btree.rs
@@ -16,7 +16,7 @@ use crate::util::{Pool, PoolClone, PoolDefault, PoolRef};
 use self::Insert::*;
 use self::InsertAction::*;
 
-const NODE_SIZE: usize = NodeSize::USIZE;
+pub(crate) const NODE_SIZE: usize = NodeSize::USIZE;
 const MEDIAN: usize = (NODE_SIZE + 1) >> 1;
 
 pub trait BTreeValue {
@@ -395,7 +395,17 @@ impl<A: BTreeValue> Node<A> {
                         path.push((self, index));
                         path
                     }
-                    None => Vec::new(),
+                    None => {
+                        // go back up to find next
+                        while let Some((node, idx)) = path.last() {
+                            if node.keys.len() == *idx {
+                                path.pop();
+                            } else {
+                                break;
+                            }
+                        }
+                        path
+                    }
                 },
                 Some(ref node) => {
                     path.push((self, index));
@@ -424,14 +434,22 @@ impl<A: BTreeValue> Node<A> {
                 path
             }
             Err(index) => match self.children[index] {
-                None if index == 0 => Vec::new(),
-                None => match self.keys.get(index - 1) {
-                    Some(_) => {
-                        path.push((self, index - 1));
-                        path
+                None if index == 0 => {
+                    // go back up to find prev
+                    while let Some((_, idx)) = path.last_mut() {
+                        if *idx == 0 {
+                            path.pop();
+                        } else {
+                            *idx -= 1;
+                            break;
+                        }
                     }
-                    None => Vec::new(),
-                },
+                    path
+                }
+                None => {
+                    path.push((self, index - 1));
+                    path
+                }
                 Some(ref node) => {
                     path.push((self, index));
                     node.path_prev(key, path)

diff --git a/src/ord/map.rs b/src/ord/map.rs
@@ -2381,6 +2381,30 @@ mod test {
         assert_eq!(vec![(1, 2), (2, 3), (3, 4), (4, 5), (5, 6)], range);
     }
 
+    #[test]
+    fn range_iter_big() {
+        use crate::nodes::btree::NODE_SIZE;
+        use std::ops::Bound::Included;
+        const N: usize = NODE_SIZE * NODE_SIZE * 5; // enough for a sizeable 3 level tree
+
+        let data = (1usize..N).filter(|i| i % 2 == 0).map(|i| (i, ()));
+        let bmap = data
+            .clone()
+            .collect::<std::collections::BTreeMap<usize, ()>>();
+        let omap = data.collect::<OrdMap<usize, ()>>();
+
+        for i in (0..NODE_SIZE * 5).chain(N - NODE_SIZE * 5..=N + 1) {
+            assert_eq!(omap.range(i..).count(), bmap.range(i..).count());
+            assert_eq!(omap.range(..i).count(), bmap.range(..i).count());
+            assert_eq!(omap.range(i..i + 7).count(), bmap.range(i..i + 7).count());
+            assert_eq!(omap.range(i..=i + 7).count(), bmap.range(i..=i + 7).count());
+            assert_eq!(
+                omap.range((Included(i), Included(i + 7))).count(),
+                bmap.range((Included(i), Included(i + 7))).count(),
+            );
+        }
+    }
+
     #[test]
     fn issue_124() {
         let mut map = OrdMap::new();