From 38409ed6dd0db3c8ef418e1342461d021ef489ee Mon Sep 17 00:00:00 2001
From: Robert Habermeier <rphmeier@gmail.com>
Date: Fri, 1 Nov 2024 14:56:23 -0500
Subject: [PATCH] fix prev/next search in btree

Previously, we weren't accounting for separator keys correctly in the recursive search. This commit fixes that and adds property tests.
---
 src/nodes/btree.rs | 31 +++++++++++++++++++++++++++++--
 src/ord/map.rs     | 14 ++++++++++++++
 src/ord/set.rs     | 20 ++++++++++++++++++++
 3 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/src/nodes/btree.rs b/src/nodes/btree.rs
index 84f63fa..3903d91 100644
--- a/src/nodes/btree.rs
+++ b/src/nodes/btree.rs
@@ -178,6 +178,25 @@ impl<A> Node<A> {
     }
 }
 
+impl<A: BTreeValue + std::fmt::Debug> Node<A> {
+    pub fn traverse_print(&self, indent: String) {
+        println!("{indent} keys n={} | {:?}", self.keys.len(), self.keys);
+        println!("{indent} children n={}", self.children.len());
+        let next_indent = indent.clone() + "    ";
+        for (i, child) in self.children.iter().enumerate() {
+            match child {
+                None => {
+                    println!("{next_indent} | child {i} empty");
+                }
+                Some(child) => {
+                    println!("{indent} | child {i}:");
+                    child.traverse_print(next_indent.clone());
+                }
+            }
+        }
+    }
+}
+
 impl<A: BTreeValue> Node<A> {
     fn child_contains<BK>(&self, index: usize, key: &BK) -> bool
     where
@@ -248,7 +267,11 @@ impl<A: BTreeValue> Node<A> {
             Err(index) => match self.children[index] {
                 None if index == 0 => None,
                 None => self.keys.get(index - 1).map(|_| &self.keys[index - 1]),
-                Some(ref node) => node.lookup_prev(key),
+                Some(ref node) => match node.lookup_prev(key) {
+                    None if index == 0 => None,
+                    None => self.keys.get(index - 1).map(|_| &self.keys[index - 1]),
+                    Some(key) => Some(key),
+                },
             },
         }
     }
@@ -265,7 +288,11 @@ impl<A: BTreeValue> Node<A> {
             Ok(index) => Some(&self.keys[index]),
             Err(index) => match self.children[index] {
                 None => self.keys.get(index).map(|_| &self.keys[index]),
-                Some(ref node) => node.lookup_next(key),
+                Some(ref node) => match node.lookup_next(key) {
+                    None if index == self.keys.len() => None,
+                    None => self.keys.get(index).map(|_| &self.keys[index]),
+                    Some(key) => Some(key),
+                },
             },
         }
     }
diff --git a/src/ord/map.rs b/src/ord/map.rs
index 1c32a38..3f674d4 100644
--- a/src/ord/map.rs
+++ b/src/ord/map.rs
@@ -2645,5 +2645,19 @@ mod test {
             }).collect();
             assert_eq!(expected, diff);
         }
+
+        #[test]
+        fn get_next_and_prev(count in 0..1000) {
+            let values = (0..count).map(|i| ((i + 1) * 2, ())).collect::<Vec<_>>();
+
+            let set = values.iter().cloned().collect::<OrdMap<i32, ()>>();
+            for value in &values {
+                let next = value.0 + 1;
+                assert_eq!(set.get_prev(&next), Some((&value.0, &value.1)));
+
+                let prev = value.0 - 1;
+                assert_eq!(set.get_next(&prev), Some((&value.0, &value.1)));
+            }
+        }
     }
 }
diff --git a/src/ord/set.rs b/src/ord/set.rs
index 2fd674a..1ff3776 100644
--- a/src/ord/set.rs
+++ b/src/ord/set.rs
@@ -178,6 +178,12 @@ pub struct OrdSet<A> {
     root: PoolRef<Node<Value<A>>>,
 }
 
+impl<A: Ord + Clone + Debug> OrdSet<A> {
+    pub fn traverse_print(&self) {
+        self.root.traverse_print(String::new());
+    }
+}
+
 impl<A> OrdSet<A> {
     /// Construct an empty set.
     #[must_use]
@@ -1239,5 +1245,19 @@ mod test {
             let result: Vec<i32> = set.range(..).rev().cloned().collect();
             assert_eq!(expected, result);
         }
+
+        #[test]
+        fn get_next_and_prev(count in 0..1000) {
+            let values = (0..count).map(|i| (i + 1) * 2).collect::<Vec<_>>();
+
+            let set = values.iter().cloned().collect::<OrdSet<i32>>();
+            for value in &values {
+                let next = *value + 1;
+                assert_eq!(set.get_prev(&next), Some(value));
+
+                let prev = *value - 1;
+                assert_eq!(set.get_next(&prev), Some(value));
+            }
+        }
     }
 }