Skip to content

Commit

Permalink
fix prev/next search in btree
Browse files Browse the repository at this point in the history
Previously, we weren't accounting for separator keys correctly in the recursive search. This commit fixes that and adds property tests.
  • Loading branch information
rphmeier committed Nov 1, 2024
1 parent 71331ea commit 38409ed
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 2 deletions.
31 changes: 29 additions & 2 deletions src/nodes/btree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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),
},
},
}
}
Expand All @@ -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),
},
},
}
}
Expand Down
14 changes: 14 additions & 0 deletions src/ord/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)));
}
}
}
}
20 changes: 20 additions & 0 deletions src/ord/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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));
}
}
}
}

0 comments on commit 38409ed

Please sign in to comment.