Repository: tnballo/scapegoat (language: rust)

Issue:
Add `range` APIs
To mirror `BTreeSet`'s [`range`](https://doc.rust-lang.org/stable/std/collections/struct.BTreeSet.html#method.range), and `BTreeMap`'s [`range`](https://doc.rust-lang.org/std/collections/btree_map/struct.BTreeMap.html#method.range)/[`range_mut`](https://doc.rust-lang.org/std/collections/btree_map/struct.BTreeMap.html#method.range_mut).

To help navigate the codebase:

* `src/tree/tree.rs` is where the majority of logic is implemented, both sets and maps are public wrappers around this tree. Since v2.0, this module is internal-only (not part of crate's public API). So you're free to make any changes necessary to support the `range` APIs. Current iterator implementations live in `src/tree/iter.rs`.

* Return types like [`Range`](https://doc.rust-lang.org/std/collections/btree_map/struct.Range.html) should be added to `src/set_types.rs` and `src/map_types.rs`.


Discussion / hints:
I'd like to work on this.
Sweet, welcome back! Assigned to you.

Fix patch:
diff --git a/src/map.rs b/src/map.rs
index 93328da..1999af5 100644
--- a/src/map.rs
+++ b/src/map.rs
@@ -1,11 +1,11 @@
 use core::borrow::Borrow;
 use core::fmt::{self, Debug};
 use core::iter::FromIterator;
-use core::ops::Index;
+use core::ops::{Index, RangeBounds};
 
 use crate::map_types::{
-    Entry, IntoIter, IntoKeys, IntoValues, Iter, IterMut, Keys, OccupiedEntry, VacantEntry, Values,
-    ValuesMut,
+    Entry, IntoIter, IntoKeys, IntoValues, Iter, IterMut, Keys, OccupiedEntry, Range, RangeMut,
+    VacantEntry, Values, ValuesMut,
 };
 use crate::tree::{SgError, SgTree};
 
@@ -977,6 +977,86 @@ impl<K: Ord + Default, V: Default, const N: usize> SgMap<K, V, N> {
             table: self,
         })
     }
+
+    /// Constructs a 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(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((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive
+    /// range from 4 to 10.
+    ///
+    /// # Panics
+    ///
+    /// Panics if range `start > end`.
+    /// Panics if range `start == end` and both bounds are `Excluded`.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    /// use core::ops::Bound::Included;
+    ///
+    /// let mut map = SgMap::<_, _, 10>::new();
+    /// map.insert(3, "a");
+    /// map.insert(5, "b");
+    /// map.insert(8, "c");
+    /// for (&key, &value) in map.range((Included(&4), Included(&8))) {
+    ///     println!("{}: {}", key, value);
+    /// }
+    /// assert_eq!(Some((&5, &"b")), map.range(4..).next());
+    /// ```
+    pub fn range<T, R>(&self, range: R) -> Range<'_, K, V, N>
+    where
+        T: Ord + ?Sized,
+        K: Borrow<T> + Ord,
+        R: RangeBounds<T>,
+    {
+        Range {
+            table: self,
+            node_idx_iter: self.bst.range_search(range).into_iter(),
+        }
+    }
+
+    /// 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(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((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive
+    /// range from 4 to 10.
+    ///
+    /// # Panics
+    ///
+    /// Panics if range `start > end`.
+    /// Panics if range `start == end` and both bounds are `Excluded`.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map: SgMap<_, _, 10> = ["Alice", "Bob", "Carol", "Cheryl"]
+    ///     .iter()
+    ///     .map(|&s| (s, 0))
+    ///     .collect();
+    /// for (_, balance) in map.range_mut("B".."Cheryl") {
+    ///     *balance += 100;
+    /// }
+    /// for (name, balance) in &map {
+    ///     println!("{} => {}", name, balance);
+    /// }
+    /// ```
+    pub fn range_mut<T, R>(&mut self, range: R) -> RangeMut<'_, T, R, K, V, N>
+    where
+        T: Ord + ?Sized,
+        K: Borrow<T> + Ord,
+        R: RangeBounds<T>,
+    {
+        RangeMut::new(self, range)
+    }
 }
 
 // Convenience Traits --------------------------------------------------------------------------------------------------
diff --git a/src/map_types.rs b/src/map_types.rs
index 3869af4..80a94fa 100644
--- a/src/map_types.rs
+++ b/src/map_types.rs
@@ -1,3 +1,9 @@
+use core::borrow::Borrow;
+use core::iter::FusedIterator;
+use core::iter::Peekable;
+use core::marker::PhantomData;
+use core::ops::RangeBounds;
+
 use crate::map::SgMap;
 use crate::tree::{
     Idx, IntoIter as TreeIntoIter, Iter as TreeIter, IterMut as TreeIterMut, SmallNode,
@@ -231,6 +237,7 @@ pub enum Entry<'a, K: Ord + Default, V: Default, const N: usize> {
     Occupied(OccupiedEntry<'a, K, V, N>),
 }
 
+use tinyvec::ArrayVec;
 use Entry::*;
 
 impl<'a, K: Ord + Default, V: Default, const N: usize> Entry<'a, K, V, N> {
@@ -600,3 +607,101 @@ impl<'a, K: Ord + Default, V: Default, const N: usize> OccupiedEntry<'a, K, V, N
         self.remove_entry().1
     }
 }
+
+/// An iterator over a sub-range of entries in a `SgMap`.
+///
+/// This `struct` is created by the [`range`] method on [`SgMap`][crate::map::SgMap]. See its
+/// documentation for more.
+pub struct Range<'a, K: Ord + Default, V: Default, const N: usize> {
+    pub(crate) table: &'a SgMap<K, V, N>,
+    pub(crate) node_idx_iter: <ArrayVec<[usize; N]> as IntoIterator>::IntoIter,
+}
+
+impl<'a, K: Ord + Default, V: Default, const N: usize> Range<'a, K, V, N> {
+    fn to_node_ref(&self, idx: usize) -> (&'a K, &'a V) {
+        let node = &self.table.bst.arena[idx];
+        (node.key(), node.val())
+    }
+}
+
+impl<'a, K: Ord + Default, V: Default, const N: usize> Iterator for Range<'a, K, V, N> {
+    type Item = (&'a K, &'a V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let node_idx = self.node_idx_iter.next()?;
+        Some(self.to_node_ref(node_idx))
+    }
+}
+
+impl<'a, K: Ord + Default, V: Default, const N: usize> DoubleEndedIterator for Range<'a, K, V, N> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        let node_idx = self.node_idx_iter.next_back()?;
+        Some(self.to_node_ref(node_idx))
+    }
+}
+
+impl<'a, K: Ord + Default, V: Default, const N: usize> FusedIterator for Range<'a, K, V, N> {}
+
+/// A mutable iterator over a sub-range of entries in a `SgMap`.
+///
+/// This `struct` is created by the [`range_mut`] method on [`SgMap`]. See its
+/// documentation for more.
+///
+/// [`range_mut`]: SgMap::range_mut
+pub struct RangeMut<'a, T, R, K, V, const N: usize>
+where
+    T: Ord + ?Sized,
+    K: Borrow<T> + Ord + Default,
+    V: Default,
+    R: RangeBounds<T>,
+{
+    inner: Peekable<TreeIterMut<'a, K, V, N>>,
+    range: R,
+    _marker: PhantomData<T>,
+}
+
+impl<'a, T, R, K, V, const N: usize> RangeMut<'a, T, R, K, V, N>
+where
+    T: Ord + ?Sized,
+    K: Borrow<T> + Ord + Default,
+    V: Default,
+    R: RangeBounds<T>,
+{
+    pub(crate) fn new(map: &'a mut SgMap<K, V, N>, range: R) -> Self {
+        let mut peekable = map.bst.iter_mut().peekable();
+
+        while let Some(node) = peekable.peek() {
+            if range.contains(node.0.borrow()) {
+                break;
+            }
+
+            peekable.next();
+        }
+
+        Self {
+            inner: peekable,
+            range,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<'a, T, R, K, V, const N: usize> Iterator for RangeMut<'a, T, R, K, V, N>
+where
+    T: Ord + ?Sized,
+    K: Borrow<T> + Ord + Default,
+    V: Default,
+    R: RangeBounds<T>,
+{
+    type Item = (&'a K, &'a mut V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let node = self.inner.next()?;
+
+        if !self.range.contains(node.0.borrow()) {
+            None
+        } else {
+            Some(node)
+        }
+    }
+}
diff --git a/src/set.rs b/src/set.rs
index f662624..3eaf06c 100644
--- a/src/set.rs
+++ b/src/set.rs
@@ -1,9 +1,12 @@
 use core::borrow::Borrow;
 use core::fmt::{self, Debug};
 use core::iter::FromIterator;
+use core::ops::RangeBounds;
 use core::ops::{BitAnd, BitOr, BitXor, Sub};
 
-use crate::set_types::{Difference, Intersection, IntoIter, Iter, SymmetricDifference, Union};
+use crate::set_types::{
+    Difference, Intersection, IntoIter, Iter, Range, SymmetricDifference, Union,
+};
 use crate::tree::{SgError, SgTree};
 
 /// Safe, fallible, embedded-friendly ordered set.
@@ -673,6 +676,40 @@ impl<T: Ord + Default, const N: usize> SgSet<T, N> {
         self.bst.len()
     }
 
+    /// Constructs a double-ended iterator over a sub-range of elements in the set.
+    /// The simplest way is to use the range syntax `min..max`, thus `range(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((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive
+    /// range from 4 to 10.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgSet;
+    /// use core::ops::Bound::Included;
+    ///
+    /// let mut set = SgSet::<_, 5>::new();
+    /// set.insert(3);
+    /// set.insert(5);
+    /// set.insert(8);
+    /// for &elem in set.range((Included(&4), Included(&8))) {
+    ///     println!("{}", elem);
+    /// }
+    /// assert_eq!(Some(&5), set.range(4..).next());
+    /// ```
+    pub fn range<K, R>(&self, range: R) -> Range<'_, T, N>
+    where
+        K: Ord + ?Sized,
+        T: Borrow<K> + Ord,
+        R: RangeBounds<K>,
+    {
+        Range {
+            table: self,
+            node_idx_iter: self.bst.range_search(range).into_iter(),
+        }
+    }
+
     /// Returns an iterator over values representing set difference, e.g., values in `self` but not in `other`, in ascending order.
     ///
     /// # Examples
diff --git a/src/set_types.rs b/src/set_types.rs
index 76704a8..a859b72 100644
--- a/src/set_types.rs
+++ b/src/set_types.rs
@@ -1,7 +1,8 @@
 use core::cmp::Ordering;
+use core::iter::FusedIterator;
 
 use crate::set::SgSet;
-use crate::tree::{Idx, IntoIter as TreeIntoIter, Iter as TreeIter};
+use crate::tree::{Idx, IntoIter as TreeIntoIter, Iter as TreeIter, SmallNode};
 
 use smallnum::SmallUnsigned;
 use tinyvec::{ArrayVec, ArrayVecIterator};
@@ -398,3 +399,34 @@ impl<'a, T: Ord + Default, const N: usize> ExactSizeIterator for Union<'a, T, N>
         self.total_cnt - self.spent_cnt
     }
 }
+
+/// An iterator over a sub-range of items in a `SgSet`.
+///
+/// This `struct` is created by the [`range`] method on [`SgSet`].
+/// See its documentation for more.
+///
+/// [`range`]: SgSet::range
+pub struct Range<'a, T: Ord + Default, const N: usize> {
+    pub(crate) table: &'a SgSet<T, N>,
+    pub(crate) node_idx_iter: <ArrayVec<[usize; N]> as IntoIterator>::IntoIter,
+}
+
+impl<'a, T: Ord + Default, const N: usize> Iterator for Range<'a, T, N> {
+    type Item = &'a T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let node_idx = self.node_idx_iter.next()?;
+        let node = &self.table.bst.arena[node_idx];
+        Some(node.key())
+    }
+}
+
+impl<'a, T: Ord + Default, const N: usize> DoubleEndedIterator for Range<'a, T, N> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        let node_idx = self.node_idx_iter.next_back()?;
+        let node = &self.table.bst.arena[node_idx];
+        Some(node.key())
+    }
+}
+
+impl<'a, T: Ord + Default, const N: usize> FusedIterator for Range<'a, T, N> {}
diff --git a/src/tree/tree.rs b/src/tree/tree.rs
index c0a3ee0..206e6e1 100644
--- a/src/tree/tree.rs
+++ b/src/tree/tree.rs
@@ -4,6 +4,7 @@ use core::fmt::{self, Debug};
 use core::hash::{Hash, Hasher};
 use core::iter::FromIterator;
 use core::mem;
+use core::ops::RangeBounds;
 use core::ops::{Index, Sub};
 
 use super::arena::Arena;
@@ -529,6 +530,30 @@ impl<K: Ord + Default, V: Default, const N: usize> SgTree<K, V, N> {
         Idx::MAX as usize
     }
 
+    pub(crate) fn range_search<T, R>(&self, range: R) -> ArrayVec<[usize; N]>
+    where
+        T: Ord + ?Sized,
+        R: RangeBounds<T>,
+        K: Borrow<T> + Ord,
+    {
+        let mut node_idxs = ArrayVec::<[usize; N]>::new();
+
+        for (idx, node) in self
+            .arena
+            .iter()
+            .enumerate()
+            .filter_map(|(i, node)| Some((i, node.as_ref()?)))
+        {
+            if range.contains(node.key().borrow()) {
+                node_idxs.push(idx);
+            }
+        }
+
+        node_idxs.sort_unstable_by_key(|idx| self.arena[*idx].key());
+
+        node_idxs
+    }
+
     // Private API -----------------------------------------------------------------------------------------------------
 
     // Iterative search. If key found, returns node idx, parent idx, and a bool indicating if node is right child


Test patch:
diff --git a/tests/test_map_api.rs b/tests/test_map_api.rs
index 80db4c0..ec097d9 100644
--- a/tests/test_map_api.rs
+++ b/tests/test_map_api.rs
@@ -283,7 +283,7 @@ fn test_map_append_fallible() {
         a.len(),
         a.capacity(),
         b.len(),
-        a.iter().filter(|(k, _)| b.contains_key(&k)).count()
+        a.iter().filter(|(k, _)| b.contains_key(k)).count()
     );
 
     assert!(a.try_append(&mut b).is_ok());
@@ -306,3 +306,31 @@ fn test_map_insert_panic() {
 
     a.insert(4, "4"); // panic
 }
+
+#[test]
+fn test_map_range() {
+    let array = [(1, "a"), (5, "e"), (3, "c"), (7, "g"), (9, "i")];
+    let map = SgMap::from(array);
+
+    let range = 3..8;
+
+    let keys: Vec<_> = map.range(range.clone()).collect();
+
+    assert!(keys.windows(2).all(|w| w[0] < w[1]));
+    assert!(keys.iter().all(|(x, _)| range.contains(*x)));
+}
+
+#[test]
+fn test_map_range_mut() {
+    let mut map: SgMap<_, _, DEFAULT_CAPACITY> =
+        ["a", "b", "c", "d", "e"].iter().map(|s| (*s, 0)).collect();
+
+    for (_, val) in map.range_mut("d"..) {
+        *val += 10;
+    }
+
+    assert_eq!(map["a"], 0);
+    assert_eq!(map["c"], 0);
+    assert_eq!(map["d"], 10);
+    assert_eq!(map["e"], 10);
+}
diff --git a/tests/test_set_api.rs b/tests/test_set_api.rs
index 62d971e..cb7b328 100644
--- a/tests/test_set_api.rs
+++ b/tests/test_set_api.rs
@@ -226,6 +226,19 @@ fn test_set_is_disjoint() {
     assert!(!a.is_disjoint(&c));
 }
 
+#[test]
+fn test_set_range() {
+    let array = [1, 5, 3, 7, 9];
+    let map = SgSet::from(array);
+
+    let range = 3..8;
+
+    let keys: Vec<_> = map.range(range.clone()).collect();
+
+    assert!(keys.windows(2).all(|w| w[0] < w[1]));
+    assert!(keys.iter().all(|x| range.contains(*x)));
+}
+
 // Fallible APIs -------------------------------------------------------------------------------------------------------
 
 #[test]
@@ -274,7 +287,7 @@ fn test_map_append_fallible() {
         a.len(),
         a.capacity(),
         b.len(),
-        a.iter().filter(|k| b.contains(&k)).count()
+        a.iter().filter(|k| b.contains(k)).count()
     );
 
     assert!(a.try_append(&mut b).is_ok());