Skip to content

Commit

Permalink
Add RangeMut iterator
Browse files Browse the repository at this point in the history 
**Description**
 - Add the `RangeMut` iterator by generalizing the existing range
   iterator using a `macro_rules` macro.
 - Also add the `TreeMap::range_mut` method and docs + example

**Motivation**
This change is needed to reach feature parity with the `BTreeMap`
which has a `range_mut` method.

**Testing Done**
Added doc test and new unit test. The `RangeMut` test coverage is
pretty low, but I'm banking on the fact that the implementation is
largely the same as `Range`.
  • Loading branch information
@declanvk
declanvk committed Sep 11, 2024
1 parent 9236921 commit d5ddd4a
Show file tree
Hide file tree
Showing 2 changed files with 203 additions and 137 deletions.
71 changes: 37 additions & 34 deletions src/collections/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -943,54 +943,57 @@ impl<K, V, const PREFIX_LEN: usize> TreeMap<K, V, PREFIX_LEN> {
)
}

/*
/// Constructs a mutable double-ended iterator over a sub-range of elements
/// in the map.
///
/// The simplest way is to use the range syntax `min..max`, thus
/// `range+_mut(min..max)` will yield elements from min (inclusive) to max
/// `range_mut(min..max)` will yield elements from min (inclusive) to max
/// (exclusive). The range may also be entered as `(Bound<T>, Bound<T>)`, so
/// for example `range_mut((Excluded(4), Included(10)))` will yield a
/// left-exclusive, right-inclusive range from 4 to 10.
//
// # Examples
//
// ```rust,should_panic
// use blart::TreeMap;
//
// let mut map: TreeMap<_, i32> = TreeMap::new();
//
// for (key, value) in [("Alice", 0), ("Bob", 0), ("Carol", 0), ("Cheryl", 0)] {
// let _ = map.try_insert(key, value).unwrap();
// }
//
// for (name, balance) in map.range_mut("B".."Cheryl") {
// *balance += 100;
//
// if name.starts_with('C') {
// *balance *= 2;
// }
// }
//
// for (name, balance) in &map {
// println!("{name} => {balance}");
// }
//
// assert_eq!(map["Alice"], 0);
// assert_eq!(map["Bob"], 100);
// assert_eq!(map["Carol"], 200);
// assert_eq!(map["Cheryl"], 200);
// ```
#[allow(dead_code)]
pub(crate) fn range_mut<Q, R>(&mut self, _range: R) -> iterators::RangeMut<K, V>
///
/// # Examples
///
/// ```rust
/// use blart::TreeMap;
///
/// let mut map: TreeMap<_, i32> = TreeMap::new();
///
/// for (key, value) in [("Alice", 0), ("Bob", 0), ("Carol", 0), ("Cheryl", 0)] {
/// let _ = map.try_insert(key, value).unwrap();
/// }
///
/// for (name, balance) in map.range_mut("B"..="Cheryl") {
/// *balance += 100;
///
/// if name.starts_with('C') {
/// *balance *= 2;
/// }
/// }
///
/// for (name, balance) in &map {
/// println!("{name} => {balance}");
/// }
///
/// assert_eq!(map["Alice"], 0);
/// assert_eq!(map["Bob"], 100);
/// assert_eq!(map["Carol"], 200);
/// assert_eq!(map["Cheryl"], 200);
/// ```
pub fn range_mut<Q, R>(&mut self, range: R) -> iterators::RangeMut<K, V, PREFIX_LEN>
where
Q: AsBytes + ?Sized,
K: Borrow<Q> + AsBytes,
R: RangeBounds<Q>,
{
todo!()
iterators::RangeMut::new(
self,
range.start_bound().map(AsBytes::as_bytes),
range.end_bound().map(AsBytes::as_bytes),
)
}

/*
/// Splits the collection into two at the given key. Returns everything
/// after the given key, including the key.
//
Expand Down
269 changes: 166 additions & 103 deletions src/collections/map/iterators/range.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -381,125 +381,161 @@ unsafe fn find_leaf_pointer_for_bound<K: AsBytes, V, const PREFIX_LEN: usize>(
})
}

/// An iterator over a sub-range of entries in a [`TreeMap`].
///
/// This struct is created by the [`range`][TreeMap::range] method on `TreeMap`.
/// See its documentation for more details.
pub struct Range<'a, K, V, const PREFIX_LEN: usize> {
inner: RawIterator<K, V, PREFIX_LEN>,
_tree: &'a TreeMap<K, V, PREFIX_LEN>,
}
macro_rules! implement_range_iter {
(
$(#[$outer:meta])*
struct $name:ident {
tree: $tree_ty:ty,
item: $item_ty:ty,
$leaf_accessor_func:ident
}
) => {
$(#[$outer])*
pub struct $name<'a, K, V, const PREFIX_LEN: usize> {
inner: RawIterator<K, V, PREFIX_LEN>,
_tree: $tree_ty,
}

impl<'a, K: AsBytes, V, const PREFIX_LEN: usize> Range<'a, K, V, PREFIX_LEN> {
/// Create a new range iterator over the given tree, starting and ending
/// according to the given bounds.
///
/// # Panics
///
/// This function will panic if `start_bound` is greater than `end_bound`.
pub(crate) fn new(
tree: &'a TreeMap<K, V, PREFIX_LEN>,
start_bound: Bound<&[u8]>,
end_bound: Bound<&[u8]>,
) -> Self {
let Some(tree_state) = &tree.state else {
return Self {
_tree: tree,
inner: RawIterator::empty(),
};
};
impl<'a, K: AsBytes, V, const PREFIX_LEN: usize> $name<'a, K, V, PREFIX_LEN> {
/// Create a new range iterator over the given tree, starting and ending
/// according to the given bounds.
///
/// # Panics
///
/// This function will panic if `start_bound` is greater than `end_bound`.
pub(crate) fn new(
tree: $tree_ty,
start_bound: Bound<&[u8]>,
end_bound: Bound<&[u8]>,
) -> Self {
let Some(tree_state) = &tree.state else {
return Self {
_tree: tree,
inner: RawIterator::empty(),
};
};

{
match (start_bound, end_bound) {
(Bound::Excluded(s), Bound::Excluded(e)) if s == e => {
panic!("range start and end are equal and excluded: ({s:?})")
},
(
Bound::Included(s) | Bound::Excluded(s),
Bound::Included(e) | Bound::Excluded(e),
) if s > e => {
panic!("range start ({s:?}) is greater than range end ({e:?})")
},
_ => {},
}
}
{
match (start_bound, end_bound) {
(Bound::Excluded(s), Bound::Excluded(e)) if s == e => {
panic!("range start and end are equal and excluded: ({s:?})")
},
(
Bound::Included(s) | Bound::Excluded(s),
Bound::Included(e) | Bound::Excluded(e),
) if s > e => {
panic!("range start ({s:?}) is greater than range end ({e:?})")
},
_ => {},
}
}

// SAFETY: Since we have a shared reference to the original TreeMap, we know
// there will be no concurrent mutation
let possible_start = unsafe { find_leaf_pointer_for_bound(tree_state, start_bound, true) };
let Some(start) = possible_start else {
return Self {
_tree: tree,
inner: RawIterator::empty(),
};
};
// SAFETY: Since we have a (shared or mutable) reference to the original TreeMap, we know
// there will be no concurrent mutation
let possible_start =
unsafe { find_leaf_pointer_for_bound(tree_state, start_bound, true) };
let Some(start) = possible_start else {
return Self {
_tree: tree,
inner: RawIterator::empty(),
};
};

// SAFETY: Since we have a shared reference to the original TreeMap, we know
// there will be no concurrent mutation
let possible_end = unsafe { find_leaf_pointer_for_bound(tree_state, end_bound, false) };
let Some(end) = possible_end else {
return Self {
_tree: tree,
inner: RawIterator::empty(),
};
};
// SAFETY: Since we have a (shared or mutable) reference to the original TreeMap, we know
// there will be no concurrent mutation
let possible_end =
unsafe { find_leaf_pointer_for_bound(tree_state, end_bound, false) };
let Some(end) = possible_end else {
return Self {
_tree: tree,
inner: RawIterator::empty(),
};
};

// SAFETY: Since we have a (shared or mutable) reference to the original TreeMap, we know
// there will be no concurrent mutation. Also, the reference lifetimes created
// are bounded to this `unsafe` block, and don't overlap with any mutation.
unsafe {
let start_leaf_bytes = start.as_key_ref().as_bytes();
let end_leaf_bytes = end.as_key_ref().as_bytes();

if start_leaf_bytes > end_leaf_bytes {
// Resolved start leaf is not less than or equal to resolved end leaf for
// iteration order
return Self {
_tree: tree,
inner: RawIterator::empty(),
};
}
}

// SAFETY: Since we have a shared reference to the original TreeMap, we know
// there will be no concurrent mutation. Also, the reference lifetimes created
// are bounded to this `unsafe` block, and don't overlap with any mutation.
unsafe {
let start_leaf_bytes = start.as_key_ref().as_bytes();
let end_leaf_bytes = end.as_key_ref().as_bytes();

if start_leaf_bytes > end_leaf_bytes {
// Resolved start leaf is not less than or equal to resolved end leaf for
// iteration order
return Self {
Self {
_tree: tree,
inner: RawIterator::empty(),
};
// SAFETY: `start` is guaranteed to be less than or equal to `end` in the iteration
// order because of the check we do on the bytes of the resolved leaf pointers, just
// above this line
inner: unsafe { RawIterator::new(start, end) },
}
}
}

Self {
_tree: tree,
// SAFETY: `start` is guaranteed to be less than or equal to `end` in the iteration
// order because of the check we do on the bytes of the resolved leaf pointers, just
// above this line
inner: unsafe { RawIterator::new(start, end) },
}
}
}
impl<'a, K, V, const PREFIX_LEN: usize> Iterator for $name<'a, K, V, PREFIX_LEN> {
type Item = $item_ty;

impl<'a, K, V, const PREFIX_LEN: usize> Iterator for Range<'a, K, V, PREFIX_LEN> {
type Item = (&'a K, &'a V);
fn next(&mut self) -> Option<Self::Item> {
// SAFETY: This iterator has a reference (either shared or mutable) to the
// original `TreeMap` it is iterating over, preventing any other modification.
let leaf_ptr = unsafe { self.inner.next()? };

fn next(&mut self) -> Option<Self::Item> {
// SAFETY: This iterator has a reference (either shared or mutable) to the
// original `TreeMap` it is iterating over, preventing any other modification.
let leaf_ptr = unsafe { self.inner.next()? };
// SAFETY: THe lifetimes returned from this function are returned as bounded by
// lifetime 'a, meaning that they cannot outlive this iterator's reference
// (shared or mutable) to the original TreeMap.
Some(unsafe { leaf_ptr.$leaf_accessor_func() })
}
}

// SAFETY: THe lifetimes returned from this function are returned as bounded by
// lifetime 'a, meaning that they cannot outlive this iterator's reference
// (shared or mutable) to the original TreeMap.
Some(unsafe { leaf_ptr.as_key_value_ref() })
}
}
impl<'a, K, V, const PREFIX_LEN: usize> DoubleEndedIterator
for $name<'a, K, V, PREFIX_LEN>
{
fn next_back(&mut self) -> Option<Self::Item> {
// SAFETY: This iterator has a reference (either shared or mutable) to the
// original `TreeMap` it is iterating over, preventing any other modification.
let leaf_ptr = unsafe { self.inner.next_back()? };

// SAFETY: THe lifetimes returned from this function are returned as bounded by
// lifetime 'a, meaning that they cannot outlive this iterator's reference
// (shared or mutable) to the original TreeMap.
Some(unsafe { leaf_ptr.$leaf_accessor_func() })
}
}

impl<'a, K, V, const PREFIX_LEN: usize> DoubleEndedIterator for Range<'a, K, V, PREFIX_LEN> {
fn next_back(&mut self) -> Option<Self::Item> {
// SAFETY: This iterator has a reference (either shared or mutable) to the
// original `TreeMap` it is iterating over, preventing any other modification.
let leaf_ptr = unsafe { self.inner.next_back()? };
impl<'a, K, V, const PREFIX_LEN: usize> FusedIterator for $name<'a, K, V, PREFIX_LEN> {}
};
}

// SAFETY: THe lifetimes returned from this function are returned as bounded by
// lifetime 'a, meaning that they cannot outlive this iterator's reference
// (shared or mutable) to the original TreeMap.
Some(unsafe { leaf_ptr.as_key_value_ref() })
implement_range_iter!(
/// An iterator over a sub-range of entries in a [`TreeMap`].
///
/// This struct is created by the [`range`][TreeMap::range] method on `TreeMap`.
/// See its documentation for more details.
struct Range {
tree: &'a TreeMap<K, V, PREFIX_LEN>,
item: (&'a K, &'a V),
as_key_value_ref
}
}
);

impl<'a, K, V, const PREFIX_LEN: usize> FusedIterator for Range<'a, K, V, PREFIX_LEN> {}
implement_range_iter!(
/// A mutable iterator over a sub-range of entries in a [`TreeMap`].
///
/// This struct is created by the [`range_mut`][TreeMap::range_mut] method on `TreeMap`.
/// See its documentation for more details.
struct RangeMut {
tree: &'a mut TreeMap<K, V, PREFIX_LEN>,
item: (&'a K, &'a mut V),
as_key_ref_value_mut
}
);

#[cfg(test)]
mod tests {
Expand Down Expand Up @@ -774,4 +810,31 @@ mod tests {
assert_eq!(it.next().unwrap(), &[127, 127, u8::MAX]);
assert_eq!(it.next(), None);
}

#[test]
fn range_mut_mutate_some_keys() {
let mut tree = fixture_tree();

tree.values_mut().for_each(|value| *value = 0);

tree.range_mut([126, 0, 0]..=[128u8, 0, 0])
.for_each(|(_, value)| *value = 1024);

assert_eq!(
tree.into_iter().collect::<Vec<_>>(),
[
([0, 0, 0], 0),
([0, 0, 255], 0),
([0, 255, 0], 0),
([0, 255, 255], 0),
([127, 127, 0], 1024),
([127, 127, 127], 1024),
([127, 127, 255], 1024),
([255, 0, 0], 0),
([255, 0, 255], 0),
([255, 255, 0], 0),
([255, 255, 255], 0)
]
);
}
}

0 comments on commit d5ddd4a

Please sign in to comment.