diff --git a/benches/native.rs b/benches/native.rs
index 36db8f1..dcd3f5d 100644
--- a/benches/native.rs
+++ b/benches/native.rs
@@ -8,7 +8,7 @@ extern crate im;
extern crate rand;
extern crate test;
-use rand::{rngs::SmallRng, SeedableRng, Rng};
+use rand::{rngs::SmallRng, Rng, SeedableRng};
use std::collections::{BTreeMap, HashMap, VecDeque};
use std::iter::FromIterator;
use test::Bencher;
diff --git a/src/hash/set.rs b/src/hash/set.rs
index 4f1dfee..53d3d6b 100644
--- a/src/hash/set.rs
+++ b/src/hash/set.rs
@@ -33,8 +33,8 @@ use std::ops::{Add, Deref, Mul};
use crate::nodes::hamt::{hash_key, Drain as NodeDrain, HashValue, Iter as NodeIter, Node};
use crate::ordset::OrdSet;
-use crate::Vector;
use crate::util::{Pool, PoolRef, Ref};
+use crate::Vector;
/// Construct a set from a sequence of values.
///
diff --git a/src/nodes/btree.rs b/src/nodes/btree.rs
index 79a5bc1..3eeb7fb 100644
--- 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 {
@@ -92,9 +92,14 @@ pub(crate) enum Remove {
Update(A, Node),
}
+enum Boundary {
+ Lowest,
+ Highest,
+}
+
enum RemoveAction {
DeleteAt(usize),
- PullUp(usize, usize, usize),
+ PullUp(Boundary, usize, usize),
Merge(usize),
StealFromLeft(usize),
StealFromRight(usize),
@@ -395,7 +400,17 @@ impl Node {
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 +439,22 @@ impl Node {
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)
@@ -618,14 +641,32 @@ impl Node {
A::Key: Borrow,
{
let index = A::search_key(&self.keys, key);
- self.remove_index(pool, index, key)
+ self.remove_index(pool, index, Ok(key))
+ }
+
+ fn remove_target(
+ &mut self,
+ pool: &Pool>,
+ target: Result<&BK, Boundary>,
+ ) -> Remove
+ where
+ A: Clone,
+ BK: Ord + ?Sized,
+ A::Key: Borrow,
+ {
+ let index = match target {
+ Ok(key) => A::search_key(&self.keys, key),
+ Err(Boundary::Lowest) => Err(0),
+ Err(Boundary::Highest) => Err(self.keys.len()),
+ };
+ self.remove_index(pool, index, target)
}
fn remove_index(
&mut self,
pool: &Pool>,
index: Result,
- key: &BK,
+ target: Result<&BK, Boundary>,
) -> Remove
where
A: Clone,
@@ -638,41 +679,30 @@ impl Node {
match (&self.children[index], &self.children[index + 1]) {
// If we're a leaf, just delete the entry.
(&None, &None) => RemoveAction::DeleteAt(index),
- // Right is empty. Attempt to steal from left if enough capacity,
- // otherwise pull the predecessor up.
- (&Some(ref left), &None) => {
- if !left.too_small() {
- RemoveAction::StealFromLeft(index + 1)
- } else {
- RemoveAction::PullUp(left.keys.len() - 1, index, index)
- }
- }
- // Left is empty. Attempt to steal from right if enough capacity,
- // otherwise pull the successor up.
- (&None, &Some(ref right)) => {
- if !right.too_small() {
- RemoveAction::StealFromRight(index)
- } else {
- RemoveAction::PullUp(0, index, index + 1)
- }
- }
- // Both left and right are non empty. Attempt to steal from left or
- // right if enough capacity, otherwise just merge the children.
+ // First consider pulling either predecessor (from left) or successor (from right).
+ // otherwise just merge the two small children.
(&Some(ref left), &Some(ref right)) => {
- if left.has_room() && !right.too_small() {
- RemoveAction::StealFromRight(index)
- } else if right.has_room() && !left.too_small() {
- RemoveAction::StealFromLeft(index + 1)
+ if !left.too_small() {
+ RemoveAction::PullUp(Boundary::Highest, index, index)
+ } else if !right.too_small() {
+ RemoveAction::PullUp(Boundary::Lowest, index, index + 1)
} else {
RemoveAction::Merge(index)
}
}
+ _ => unreachable!("Branch missing children"),
}
}
- // Key is adjacent to some key in node
+ // Target is adjacent to some key in node
Err(index) => match self.children[index] {
- // No child at location means key isn't in map.
- None => return Remove::NoChange,
+ // We're deading with a leaf node
+ None => match target {
+ // No child at location means key isn't in map.
+ Ok(_key) => return Remove::NoChange,
+ // Looking for the lowest or highest key
+ Err(Boundary::Lowest) => RemoveAction::DeleteAt(0),
+ Err(Boundary::Highest) => RemoveAction::DeleteAt(self.keys.len() - 1),
+ },
// Child at location, but it's at minimum capacity.
Some(ref child) if child.too_small() => {
let left = if index > 0 {
@@ -708,13 +738,13 @@ impl Node {
self.children.remove(index);
Remove::Removed(pair)
}
- RemoveAction::PullUp(target_index, pull_to, child_index) => {
+ RemoveAction::PullUp(boundary, pull_to, child_index) => {
let children = &mut self.children;
let mut update = None;
let value;
if let Some(&mut Some(ref mut child_ref)) = children.get_mut(child_index) {
let child = PoolRef::make_mut(pool, child_ref);
- match child.remove_index(pool, Ok(target_index), key) {
+ match child.remove_target(pool, Err(boundary)) {
Remove::NoChange => unreachable!(),
Remove::Removed(pulled_value) => {
value = self.keys.set(pull_to, pulled_value);
@@ -741,7 +771,7 @@ impl Node {
PoolRef::unwrap_or_clone(left),
PoolRef::unwrap_or_clone(right),
);
- let (removed, new_child) = match merged_child.remove(pool, key) {
+ let (removed, new_child) = match merged_child.remove_target(pool, target) {
Remove::NoChange => unreachable!(),
Remove::Removed(removed) => (removed, merged_child),
Remove::Update(removed, updated_child) => (removed, updated_child),
@@ -760,13 +790,7 @@ impl Node {
{
let mut children = self.children.as_mut_slice()[index - 1..=index]
.iter_mut()
- .map(|n| {
- if let Some(ref mut o) = *n {
- o
- } else {
- unreachable!()
- }
- });
+ .map(|n| n.as_mut().unwrap());
let left = PoolRef::make_mut(pool, children.next().unwrap());
let child = PoolRef::make_mut(pool, children.next().unwrap());
// Prepare the rebalanced node.
@@ -774,7 +798,7 @@ impl Node {
left.children.last().unwrap().clone(),
self.keys[index - 1].clone(),
);
- match child.remove(pool, key) {
+ match child.remove_target(pool, target) {
Remove::NoChange => {
// Key wasn't there, we need to revert the steal.
child.pop_min();
@@ -806,18 +830,12 @@ impl Node {
{
let mut children = self.children.as_mut_slice()[index..index + 2]
.iter_mut()
- .map(|n| {
- if let Some(ref mut o) = *n {
- o
- } else {
- unreachable!()
- }
- });
+ .map(|n| n.as_mut().unwrap());
let child = PoolRef::make_mut(pool, children.next().unwrap());
let right = PoolRef::make_mut(pool, children.next().unwrap());
// Prepare the rebalanced node.
child.push_max(right.children[0].clone(), self.keys[index].clone());
- match child.remove(pool, key) {
+ match child.remove_target(pool, target) {
Remove::NoChange => {
// Key wasn't there, we need to revert the steal.
child.pop_max();
@@ -844,11 +862,17 @@ impl Node {
Remove::Removed(out_value)
}
RemoveAction::MergeFirst(index) => {
- if self.keys[index].cmp_keys(key) != Ordering::Equal
- && !self.child_contains(index, key)
- && !self.child_contains(index + 1, key)
- {
- return Remove::NoChange;
+ if let Ok(key) = target {
+ // Bail early if we're looking for a not existing key
+ match self.keys[index].cmp_keys(key) {
+ Ordering::Less if !self.child_contains(index + 1, key) => {
+ return Remove::NoChange
+ }
+ Ordering::Greater if !self.child_contains(index, key) => {
+ return Remove::NoChange
+ }
+ _ => (),
+ }
}
let left = self.children.remove(index).unwrap();
let right = mem::replace(&mut self.children[index], None).unwrap();
@@ -860,7 +884,7 @@ impl Node {
);
let update;
let out_value;
- match merged.remove(pool, key) {
+ match merged.remove_target(pool, target) {
Remove::NoChange => {
panic!("nodes::btree::Node::remove: caught an absent key too late while merging");
}
@@ -887,7 +911,7 @@ impl Node {
let out_value;
if let Some(&mut Some(ref mut child_ref)) = self.children.get_mut(index) {
let child = PoolRef::make_mut(pool, child_ref);
- match child.remove(pool, key) {
+ match child.remove_target(pool, target) {
Remove::NoChange => return Remove::NoChange,
Remove::Removed(value) => {
out_value = value;
diff --git a/src/ord/map.rs b/src/ord/map.rs
index 14931f7..55294d8 100644
--- 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::>();
+ let omap = data.collect::>();
+
+ 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();