Skip to content
Merged
31 changes: 24 additions & 7 deletions src/btreemap.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1188,16 +1188,33 @@ where
self.range_internal(key_range).into()
}

/// Returns an iterator pointing to the first element below the given bound.
/// Returns an empty iterator if there are no keys below the given bound.
pub fn iter_upper_bound(&self, bound: &K) -> Iter<K, V, M> {
/// Returns an iterator starting just before the given key.
///
/// Finds the largest key strictly less than `bound` and starts from it.
/// Useful when `range(bound..)` skips the previous element.
///
/// Returns an empty iterator if no smaller key exists.
pub fn iter_from_prev_key(&self, bound: &K) -> Iter<K, V, M> {
if let Some((start_key, _)) = self.range(..bound).next_back() {
IterInternal::new_in_range(self, (Bound::Included(start_key), Bound::Unbounded)).into()
} else {
IterInternal::null(self).into()
}
}

/// **Deprecated**: use [`iter_from_prev_key`] instead.
///
/// The name `iter_upper_bound` was misleading — it suggested an inclusive
/// upper bound. In reality, it starts from the largest key strictly less
/// than the given bound.
///
/// The new name, [`iter_from_prev_key`], better reflects this behavior and
/// improves code clarity.
#[deprecated(note = "use `iter_from_prev_key` instead")]
pub fn iter_upper_bound(&self, bound: &K) -> Iter<K, V, M> {
self.iter_from_prev_key(bound)
}

/// Returns an iterator over the keys of the map.
pub fn keys(&self) -> KeysIter<K, V, M> {
self.iter_internal().into()
Expand Down Expand Up @@ -2955,28 +2972,28 @@ mod test {
}
btree_test!(test_bruteforce_range_search, bruteforce_range_search);

fn test_iter_upper_bound<K: TestKey, V: TestValue>() {
fn test_iter_from_prev_key<K: TestKey, V: TestValue>() {
let (key, value) = (K::build, V::build);
run_btree_test(|mut btree| {
for j in 0..100 {
btree.insert(key(j), value(j));
for i in 0..=j {
assert_eq!(
btree.iter_upper_bound(&key(i + 1)).next(),
btree.iter_from_prev_key(&key(i + 1)).next(),
Some((key(i), value(i))),
"failed to get an upper bound for key({})",
i + 1
);
}
assert_eq!(
btree.iter_upper_bound(&key(0)).next(),
btree.iter_from_prev_key(&key(0)).next(),
None,
"key(0) must not have an upper bound"
);
}
});
}
btree_test!(test_test_iter_upper_bound, test_iter_upper_bound);
btree_test!(test_test_iter_from_prev_key, test_iter_from_prev_key);

// A buggy implementation of storable where the max_size is smaller than the serialized size.
#[derive(Clone, Ord, PartialOrd, Eq, PartialEq)]
Expand Down
4 changes: 2 additions & 2 deletions src/btreemap/proptests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -141,12 +141,12 @@ fn map_min_max(#[strategy(pvec(any::<u64>(), 10..100))] keys: Vec<u64>) {
}

#[proptest]
fn map_upper_bound_iter(#[strategy(pvec(0u64..u64::MAX -1 , 10..100))] keys: Vec<u64>) {
fn map_iter_from_prev_key(#[strategy(pvec(0u64..u64::MAX -1 , 10..100))] keys: Vec<u64>) {
run_btree_test(|mut map| {
for k in keys.iter() {
map.insert(*k, ());

prop_assert_eq!(Some((*k, ())), map.iter_upper_bound(&(k + 1)).next());
prop_assert_eq!(Some((*k, ())), map.iter_from_prev_key(&(k + 1)).next());
}

Ok(())
Expand Down
86 changes: 0 additions & 86 deletions src/btreeset.rs
Original file line number Diff line number Diff line change
Expand Up @@ -510,31 +510,6 @@ where
Iter::new(self.map.range(key_range))
}

/// Returns an iterator pointing to the first element strictly below the given bound.
/// Returns an empty iterator if there are no keys strictly below the given bound.
///
/// # Complexity
/// O(log n) for creating the iterator, where n is the number of elements in the set.
///
/// # Example
///
/// ```rust
/// use ic_stable_structures::{BTreeSet, DefaultMemoryImpl};
/// use ic_stable_structures::memory_manager::{MemoryId, MemoryManager};
///
/// let mem_mgr = MemoryManager::init(DefaultMemoryImpl::default());
/// let mut set: BTreeSet<u64, _> = BTreeSet::new(mem_mgr.get(MemoryId::new(0)));
/// set.insert(1);
/// set.insert(2);
/// set.insert(3);
///
/// let upper_bound: Option<u64> = set.iter_upper_bound(&3).next();
/// assert_eq!(upper_bound, Some(2));
/// ```
pub fn iter_upper_bound(&self, bound: &K) -> Iter<K, M> {
Comment thread
maksymar marked this conversation as resolved.
Iter::new(self.map.iter_upper_bound(bound))
}

/// Returns an iterator over the union of this set and another.
///
/// The union of two sets is a set containing all elements that are in either set.
Expand Down Expand Up @@ -1110,31 +1085,6 @@ mod test {
assert!(!btreeset.contains(&1u32));
}

#[test]
fn test_iter_upper_bound() {
let mem = make_memory();
let mut btreeset = BTreeSet::new(mem);

for i in 0u32..100 {
btreeset.insert(i);

// Test that `iter_upper_bound` returns the largest element strictly below the bound.
for j in 1u32..=i {
assert_eq!(
btreeset.iter_upper_bound(&(j + 1)).next(),
Some(j),
"failed to get an upper bound for {}",
j + 1
);
}
assert_eq!(
btreeset.iter_upper_bound(&0).next(),
None,
"0 must not have an upper bound"
);
}
}

#[test]
fn test_iter() {
let mem = make_memory();
Expand Down Expand Up @@ -1234,20 +1184,6 @@ mod test {
assert_eq!(elements[999], 999);
}

#[test]
fn test_iter_upper_bound_large_set() {
let mem = make_memory();
let mut btreeset: BTreeSet<u32, _> = BTreeSet::new(mem);

for i in 0u32..1000 {
btreeset.insert(i);
}

assert_eq!(btreeset.iter_upper_bound(&500).next(), Some(499));
assert_eq!(btreeset.iter_upper_bound(&0).next(), None);
assert_eq!(btreeset.iter_upper_bound(&1000).next(), Some(999));
}

#[test]
fn test_range_large_set() {
let mem = make_memory();
Expand Down Expand Up @@ -1304,14 +1240,6 @@ mod test {
assert_eq!(btreeset.pop_last(), None);
}

#[test]
fn test_iter_upper_bound_empty() {
let mem = make_memory();
let btreeset: BTreeSet<u32, _> = BTreeSet::new(mem);

assert_eq!(btreeset.iter_upper_bound(&42u32).next(), None);
}

#[test]
fn test_range_empty() {
let mem = make_memory();
Expand Down Expand Up @@ -1414,20 +1342,6 @@ mod test {
assert!(btreeset.is_empty());
}

#[test]
fn test_iter_upper_bound_edge_cases() {
let mem = make_memory();
let mut btreeset: BTreeSet<u32, _> = BTreeSet::new(mem);

for i in 1..=10 {
btreeset.insert(i);
}

assert_eq!(btreeset.iter_upper_bound(&1).next(), None); // No element strictly below 1
assert_eq!(btreeset.iter_upper_bound(&5).next(), Some(4)); // Largest element below 5
assert_eq!(btreeset.iter_upper_bound(&11).next(), Some(10)); // Largest element below 11
}

#[test]
fn test_is_disjoint_with_disjoint_sets() {
let mem1 = make_memory();
Expand Down
13 changes: 0 additions & 13 deletions src/btreeset/proptests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,19 +67,6 @@ fn set_min_max(#[strategy(pvec(any::<u64>(), 10..100))] keys: Vec<u64>) {
});
}

#[proptest]
fn set_upper_bound_iter(#[strategy(pvec(0u64..u64::MAX - 1, 10..100))] keys: Vec<u64>) {
crate::btreeset::test::run_btree_test(|mut set| {
for k in keys.iter() {
set.insert(*k);

prop_assert_eq!(Some(*k), set.iter_upper_bound(&(k + 1)).next());
}

Ok(())
});
}

// Given an operation, executes it on the given stable btreeset and standard btreeset, verifying
// that the result of the operation is equal in both btrees.
fn execute_operation<M: Memory>(
Expand Down
9 changes: 5 additions & 4 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,11 @@ mod base_vec;
pub mod btreemap;
pub mod cell;
pub use cell::{Cell as StableCell, Cell};
pub mod btreeset;
pub mod file_mem;
#[cfg(target_arch = "wasm32")]
mod ic0_memory; // Memory API for canisters.
pub mod log;
pub use log::{Log as StableLog, Log};
pub mod btreeset;
pub mod memory_manager;
pub mod min_heap;
pub mod reader;
Expand All @@ -17,20 +16,22 @@ pub mod storable;
mod tests;
mod types;
pub mod vec;
pub use min_heap::{MinHeap, MinHeap as StableMinHeap};
pub use vec::{Vec as StableVec, Vec};
pub mod vec_mem;
pub mod writer;

pub use btreemap::{BTreeMap, BTreeMap as StableBTreeMap};
pub use btreeset::{BTreeSet, BTreeSet as StableBTreeSet};
pub use file_mem::FileMemory;
#[cfg(target_arch = "wasm32")]
pub use ic0_memory::Ic0StableMemory;
pub use log::{Log as StableLog, Log};
pub use min_heap::{MinHeap, MinHeap as StableMinHeap};
use std::error;
use std::fmt::{Display, Formatter};
use std::mem::MaybeUninit;
pub use storable::Storable;
use types::Address;
pub use vec::{Vec as StableVec, Vec};
pub use vec_mem::VectorMemory;

#[cfg(target_arch = "wasm32")]
Expand Down
Loading