Skip to main content

small_collections/sets/
btree_set.rs

1//! Sorted set that lives on the stack and spills to the heap.
2//!
3//! [`SmallBTreeSet`] is a thin wrapper around `SmallBTreeMap<T, (), N>`,
4//! inheriting all the sorted-order guarantees and the stack→heap spill protocol.
5
6use crate::AnySet;
7use crate::SmallBTreeMap;
8use std::borrow::Borrow;
9use std::fmt::{self, Debug};
10use std::iter::FromIterator;
11
12/// A sorted set that lives on the stack for up to `N` elements, then spills to the heap.
13///
14/// Implemented as a zero-overhead wrapper around `SmallBTreeMap<T, (), N>` so that `()`
15/// zero-sized values add no memory cost.  All iteration is in ascending key order.
16///
17/// # Generic parameters
18/// | Parameter | Meaning |
19/// |-----------|--------|
20/// | `T` | Element type; must implement `Ord` |
21/// | `N` | Stack capacity — max elements before spill |
22pub struct SmallBTreeSet<T, const N: usize> {
23    inner: SmallBTreeMap<T, (), N>,
24}
25
26impl<T, const N: usize> SmallBTreeSet<T, N>
27where
28    T: Ord,
29{
30    /// Creates a new empty sorted set.
31    pub fn new() -> Self {
32        Self {
33            inner: SmallBTreeMap::new(),
34        }
35    }
36
37    /// Creates a new empty set starting on the heap with the given initial capacity.
38    ///
39    /// If `cap <= N` this is equivalent to [`new`](SmallBTreeSet::new).
40    pub fn with_capacity(cap: usize) -> Self {
41        Self {
42            inner: SmallBTreeMap::with_capacity(cap),
43        }
44    }
45
46    /// Returns the number of elements in the set.
47    pub fn len(&self) -> usize {
48        self.inner.len()
49    }
50
51    /// Returns `true` if the set contains no elements.
52    pub fn is_empty(&self) -> bool {
53        self.inner.is_empty()
54    }
55
56    /// Clears the set, removing all elements.
57    pub fn clear(&mut self) {
58        self.inner.clear();
59    }
60
61    /// Adds a value to the set.
62    ///
63    /// Returns whether the value was newly inserted.
64    pub fn insert(&mut self, value: T) -> bool {
65        self.inner.insert(value, ()).is_none()
66    }
67
68    /// Returns `true` if the set contains a value.
69    pub fn contains<Q>(&self, value: &Q) -> bool
70    where
71        T: Borrow<Q>,
72        Q: Ord + ?Sized,
73    {
74        self.inner.get(value).is_some()
75    }
76
77    /// Removes a value from the set. Returns whether the value was present in the set.
78    pub fn remove<Q>(&mut self, value: &Q) -> bool
79    where
80        T: Borrow<Q>,
81        Q: Ord + ?Sized,
82    {
83        self.inner.remove(value).is_some()
84    }
85
86    /// Returns an iterator visiting all elements in ascending order.
87    pub fn iter(&self) -> Iter<'_, T> {
88        Iter {
89            inner: self.inner.iter(),
90        }
91    }
92
93    /// Returns `true` if the set is stored on the stack.
94    pub fn is_on_stack(&self) -> bool {
95        self.inner.is_on_stack()
96    }
97}
98
99impl<T, const N: usize> AnySet<T> for SmallBTreeSet<T, N>
100where
101    T: Ord,
102{
103    fn contains(&self, value: &T) -> bool {
104        self.contains(value)
105    }
106
107    fn len(&self) -> usize {
108        self.len()
109    }
110}
111
112/// A structure representing `Iter`.
113pub struct Iter<'a, T: Ord> {
114    inner: crate::maps::btree_map::Iter<'a, T, ()>,
115}
116
117impl<'a, T: Ord> Iterator for Iter<'a, T> {
118    type Item = &'a T;
119    fn next(&mut self) -> Option<Self::Item> {
120        self.inner.next().map(|(k, _)| k)
121    }
122}
123
124impl<T, const N: usize> IntoIterator for SmallBTreeSet<T, N>
125where
126    T: Ord,
127{
128    type Item = T;
129    type IntoIter = IntoIter<T, N>;
130
131    fn into_iter(self) -> Self::IntoIter {
132        IntoIter {
133            inner: self.inner.into_iter(),
134        }
135    }
136}
137
138/// A structure representing `IntoIter`.
139pub struct IntoIter<T: Ord, const N: usize> {
140    inner: crate::maps::btree_map::IntoIter<T, (), N>,
141}
142
143impl<T: Ord, const N: usize> Iterator for IntoIter<T, N> {
144    type Item = T;
145    fn next(&mut self) -> Option<Self::Item> {
146        self.inner.next().map(|(k, _)| k)
147    }
148}
149
150impl<T, const N: usize> Clone for SmallBTreeSet<T, N>
151where
152    T: Ord + Clone,
153{
154    fn clone(&self) -> Self {
155        Self {
156            inner: self.inner.clone(),
157        }
158    }
159}
160
161impl<T, const N: usize> Default for SmallBTreeSet<T, N>
162where
163    T: Ord,
164{
165    fn default() -> Self {
166        Self::new()
167    }
168}
169
170impl<T, const N: usize> Debug for SmallBTreeSet<T, N>
171where
172    T: Ord + Debug,
173{
174    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
175        f.debug_set().entries(self.iter()).finish()
176    }
177}
178
179impl<T, const N: usize> FromIterator<T> for SmallBTreeSet<T, N>
180where
181    T: Ord,
182{
183    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
184        let mut set = Self::new();
185        for value in iter {
186            set.insert(value);
187        }
188        set
189    }
190}
191
192impl<T, const N: usize> Extend<T> for SmallBTreeSet<T, N>
193where
194    T: Ord,
195{
196    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
197        for value in iter {
198            self.insert(value);
199        }
200    }
201}
202
203use std::cmp::Ordering;
204
205impl<T, const N: usize, S> PartialEq<S> for SmallBTreeSet<T, N>
206where
207    T: Ord,
208    S: AnySet<T>,
209{
210    fn eq(&self, other: &S) -> bool {
211        if self.len() != other.len() {
212            return false;
213        }
214        self.iter().all(|v| other.contains(v))
215    }
216}
217
218impl<T, const N: usize> Eq for SmallBTreeSet<T, N> where T: Ord + Eq {}
219
220impl<T, const N: usize> PartialOrd for SmallBTreeSet<T, N>
221where
222    T: Ord,
223{
224    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
225        self.iter().partial_cmp(other.iter())
226    }
227}
228
229impl<T, const N: usize> Ord for SmallBTreeSet<T, N>
230where
231    T: Ord,
232{
233    fn cmp(&self, other: &Self) -> Ordering {
234        self.iter().cmp(other.iter())
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    #[test]
243    fn test_btree_set_stack_ops_sorted_order() {
244        let mut set: SmallBTreeSet<i32, 4> = SmallBTreeSet::new();
245        set.insert(3);
246        set.insert(1);
247        set.insert(2);
248
249        assert_eq!(set.len(), 3);
250        assert!(set.contains(&1));
251        assert!(set.contains(&2));
252        assert!(set.contains(&3));
253        assert!(!set.contains(&4));
254
255        let values: Vec<_> = set.iter().cloned().collect();
256        assert_eq!(values, vec![1, 2, 3]);
257    }
258
259    #[test]
260    fn test_btree_set_spill_trigger_on_insert() {
261        let mut set: SmallBTreeSet<i32, 2> = SmallBTreeSet::new();
262        set.insert(1);
263        set.insert(2);
264        assert!(set.is_on_stack());
265
266        set.insert(0); // Spill
267        assert!(!set.is_on_stack());
268
269        let values: Vec<_> = set.iter().cloned().collect();
270        assert_eq!(values, vec![0, 1, 2]);
271    }
272
273    #[test]
274    fn test_btree_set_any_storage_remove() {
275        let mut set: SmallBTreeSet<i32, 4> = SmallBTreeSet::new();
276        set.insert(1);
277        set.insert(2);
278        assert!(set.remove(&1));
279        assert_eq!(set.len(), 1);
280        assert!(!set.contains(&1));
281
282        set.insert(3);
283        set.insert(4);
284        set.insert(5); // Spill
285        assert!(set.remove(&5));
286        assert_eq!(set.len(), 3);
287    }
288
289    #[test]
290    fn test_btree_set_any_storage_clear() {
291        let mut set: SmallBTreeSet<i32, 4> = SmallBTreeSet::new();
292        set.insert(1);
293        set.clear();
294        assert!(set.is_empty());
295
296        for i in 0..5 {
297            set.insert(i);
298        }
299        assert!(!set.is_on_stack());
300        set.clear();
301        assert!(set.is_empty());
302        assert!(!set.is_on_stack());
303    }
304
305    #[test]
306    fn test_btree_set_traits_exhaustive() {
307        let mut set: SmallBTreeSet<i32, 4> = SmallBTreeSet::new();
308        set.insert(1);
309        set.insert(2);
310
311        // Clone
312        let cloned = set.clone();
313        assert_eq!(cloned.len(), 2);
314        assert!(cloned.contains(&1));
315
316        // Debug
317        let debug = format!("{:?}", set);
318        assert!(debug.contains("1"));
319
320        // FromIterator
321        let collected: SmallBTreeSet<i32, 4> = vec![1, 2].into_iter().collect();
322        assert_eq!(collected.len(), 2);
323
324        // Extend
325        let mut set2 = SmallBTreeSet::<i32, 4>::new();
326        set2.extend(vec![1, 2]);
327        assert_eq!(set2.len(), 2);
328
329        // IntoIterator
330        let vec: Vec<_> = set2.into_iter().collect();
331        assert_eq!(vec.len(), 2);
332        assert!(vec.contains(&1));
333        assert!(vec.contains(&2));
334    }
335
336    #[test]
337    fn test_btree_set_traits_any_set_impl() {
338        let set: SmallBTreeSet<i32, 4> = vec![1, 2].into_iter().collect();
339        let any: &dyn AnySet<i32> = &set;
340        assert_eq!(any.len(), 2);
341        assert!(any.contains(&1));
342    }
343
344    #[test]
345    fn test_btree_set_traits_equality() {
346        let set: SmallBTreeSet<i32, 4> = vec![1, 2].into_iter().collect();
347        let s2 = set.clone();
348        assert_eq!(set, s2);
349    }
350
351    #[test]
352    fn test_btree_set_any_storage_clone_heap() {
353        let set: SmallBTreeSet<i32, 2> = vec![1, 2, 3].into_iter().collect();
354        assert!(!set.is_on_stack());
355
356        // clone heap
357        let cloned = set.clone();
358        assert_eq!(cloned.len(), 3);
359    }
360
361    #[test]
362    fn test_btree_set_any_storage_with_capacity() {
363        let s_cap = SmallBTreeSet::<i32, 2>::with_capacity(10);
364        assert!(!s_cap.is_on_stack());
365    }
366
367    #[test]
368    fn test_btree_set_traits_comparison() {
369        let set1: SmallBTreeSet<i32, 2> = vec![1, 2].into_iter().collect();
370        let set2: SmallBTreeSet<i32, 2> = vec![1, 2].into_iter().collect();
371        let set3: SmallBTreeSet<i32, 2> = vec![1, 3].into_iter().collect();
372
373        // PartialEq (generic)
374        assert_eq!(set1, set2);
375        assert_ne!(set1, set3);
376
377        // PartialOrd / Ord
378        assert!(set1 < set3);
379        assert!(set3 > set1);
380
381        // Interop with std::collections::BTreeSet
382        let std_set: std::collections::BTreeSet<i32> = vec![1, 2].into_iter().collect();
383        assert_eq!(set1, std_set);
384    }
385
386    #[test]
387    fn test_btree_set_coverage_gaps() {
388        // Default
389        let set: SmallBTreeSet<i32, 4> = Default::default();
390        assert!(set.is_empty());
391
392        // PartialEq differing lengths
393        let set1: SmallBTreeSet<i32, 4> = vec![1, 2].into_iter().collect();
394        let set2: SmallBTreeSet<i32, 4> = vec![1].into_iter().collect();
395        assert_ne!(set1, set2); // hits `if self.len() != other.len()`
396
397        // Ord cmp
398        use std::cmp::Ordering;
399        assert_eq!(set1.cmp(&set1), Ordering::Equal);
400        assert_eq!(set1.cmp(&set2), Ordering::Greater);
401    }
402}