Skip to main content

immutable_chunkmap/
set.rs

1use crate::avl::{Iter, NodeHandle, NodeRef, Tree, WeakTree};
2pub use crate::chunk::DEFAULT_SIZE;
3use core::{
4    borrow::Borrow,
5    cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd},
6    default::Default,
7    fmt::{self, Debug, Formatter},
8    hash::{Hash, Hasher},
9    iter::FromIterator,
10    ops::{RangeBounds, RangeFull},
11};
12
13#[cfg(feature = "serde")]
14use serde::{
15    de::{SeqAccess, Visitor},
16    ser::SerializeSeq,
17    Deserialize, Deserializer, Serialize, Serializer,
18};
19
20#[cfg(feature = "serde")]
21use core::marker::PhantomData;
22
23#[cfg(feature = "rayon")]
24use rayon::{
25    iter::{FromParallelIterator, IntoParallelIterator},
26    prelude::*,
27};
28
29/// This set uses a similar strategy to BTreeSet to ensure cache
30/// efficient performance on modern hardware while still providing
31/// log(N) get, insert, and remove operations.
32/// # Examples
33/// ```
34/// # extern crate alloc;
35/// use alloc::string::String;
36/// use self::immutable_chunkmap::set::SetM;
37///
38/// let m =
39///    SetM::new()
40///    .insert(String::from("1")).0
41///    .insert(String::from("2")).0
42///    .insert(String::from("3")).0;
43///
44/// assert_eq!(m.contains("1"), true);
45/// assert_eq!(m.contains("2"), true);
46/// assert_eq!(m.contains("3"), true);
47/// assert_eq!(m.contains("4"), false);
48///
49/// for k in &m { println!("{}", k) }
50/// ```
51#[derive(Clone)]
52#[repr(transparent)]
53pub struct Set<K: Ord + Clone, const SIZE: usize>(Tree<K, (), SIZE>);
54
55/// set with a smaller chunk size, faster to update, slower to search
56pub type SetS<K> = Set<K, { DEFAULT_SIZE / 2 }>;
57
58/// set with the default chunk size, a good balance of search and update performance
59pub type SetM<K> = Set<K, DEFAULT_SIZE>;
60
61/// set with a larger chunk size, faster to search, slower to update
62pub type SetL<K> = Set<K, { DEFAULT_SIZE * 2 }>;
63
64#[derive(Clone)]
65pub struct WeakSetRef<K: Ord + Clone, const SIZE: usize>(WeakTree<K, (), SIZE>);
66
67pub type WeakSetRefS<K> = WeakSetRef<K, 32>;
68pub type WeakSetRefM<K> = WeakSetRef<K, 128>;
69pub type WeakSetRefL<K> = WeakSetRef<K, 512>;
70
71impl<K, const SIZE: usize> WeakSetRef<K, SIZE>
72where
73    K: Ord + Clone,
74{
75    pub fn upgrade(&self) -> Option<Set<K, SIZE>> {
76        self.0.upgrade().map(Set)
77    }
78}
79
80impl<K, const SIZE: usize> Hash for Set<K, SIZE>
81where
82    K: Hash + Ord + Clone,
83{
84    fn hash<H: Hasher>(&self, state: &mut H) {
85        self.0.hash(state)
86    }
87}
88
89impl<K, const SIZE: usize> Default for Set<K, SIZE>
90where
91    K: Ord + Clone,
92{
93    fn default() -> Set<K, SIZE> {
94        Set::new()
95    }
96}
97
98impl<K, const SIZE: usize> PartialEq for Set<K, SIZE>
99where
100    K: Ord + Clone,
101{
102    fn eq(&self, other: &Set<K, SIZE>) -> bool {
103        self.0 == other.0
104    }
105}
106
107impl<K, const SIZE: usize> Eq for Set<K, SIZE> where K: Eq + Ord + Clone {}
108
109impl<K, const SIZE: usize> PartialOrd for Set<K, SIZE>
110where
111    K: Ord + Clone,
112{
113    fn partial_cmp(&self, other: &Set<K, SIZE>) -> Option<Ordering> {
114        self.0.partial_cmp(&other.0)
115    }
116}
117
118impl<K, const SIZE: usize> Ord for Set<K, SIZE>
119where
120    K: Ord + Clone,
121{
122    fn cmp(&self, other: &Set<K, SIZE>) -> Ordering {
123        self.0.cmp(&other.0)
124    }
125}
126
127impl<K, const SIZE: usize> Debug for Set<K, SIZE>
128where
129    K: Debug + Ord + Clone,
130{
131    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
132        f.debug_set().entries(self.into_iter()).finish()
133    }
134}
135
136impl<K, const SIZE: usize> FromIterator<K> for Set<K, SIZE>
137where
138    K: Ord + Clone,
139{
140    fn from_iter<T: IntoIterator<Item = K>>(iter: T) -> Self {
141        Set::new().insert_many(iter)
142    }
143}
144
145pub struct SetIter<
146    'a,
147    R: RangeBounds<Q> + 'a,
148    Q: Ord + ?Sized,
149    K: 'a + Clone + Ord + Borrow<Q>,
150    const SIZE: usize,
151>(Iter<'a, R, Q, K, (), SIZE>);
152
153impl<'a, R, Q, K, const SIZE: usize> Iterator for SetIter<'a, R, Q, K, SIZE>
154where
155    Q: Ord + ?Sized,
156    R: RangeBounds<Q> + 'a,
157    K: 'a + Clone + Ord + Borrow<Q>,
158{
159    type Item = &'a K;
160    fn next(&mut self) -> Option<Self::Item> {
161        self.0.next().map(|(k, ())| k)
162    }
163}
164
165impl<'a, R, Q, K, const SIZE: usize> DoubleEndedIterator for SetIter<'a, R, Q, K, SIZE>
166where
167    Q: Ord + ?Sized,
168    R: RangeBounds<Q> + 'a,
169    K: 'a + Clone + Ord + Borrow<Q>,
170{
171    fn next_back(&mut self) -> Option<Self::Item> {
172        self.0.next_back().map(|(k, ())| k)
173    }
174}
175
176impl<'a, K, const SIZE: usize> IntoIterator for &'a Set<K, SIZE>
177where
178    K: 'a + Ord + Clone,
179{
180    type Item = &'a K;
181    type IntoIter = SetIter<'a, RangeFull, K, K, SIZE>;
182    fn into_iter(self) -> Self::IntoIter {
183        SetIter(self.0.into_iter())
184    }
185}
186
187#[cfg(feature = "serde")]
188impl<V, const SIZE: usize> Serialize for Set<V, SIZE>
189where
190    V: Serialize + Clone + Ord,
191{
192    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
193    where
194        S: Serializer,
195    {
196        let mut seq = serializer.serialize_seq(Some(self.len()))?;
197        for v in self {
198            seq.serialize_element(v)?
199        }
200        seq.end()
201    }
202}
203
204#[cfg(feature = "serde")]
205struct SetVisitor<V: Clone + Ord, const SIZE: usize> {
206    marker: PhantomData<fn() -> Set<V, SIZE>>,
207}
208
209#[cfg(feature = "serde")]
210impl<'a, V, const SIZE: usize> Visitor<'a> for SetVisitor<V, SIZE>
211where
212    V: Deserialize<'a> + Clone + Ord,
213{
214    type Value = Set<V, SIZE>;
215
216    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
217        formatter.write_str("expecting an immutable_chunkmap::Set")
218    }
219
220    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
221    where
222        A: SeqAccess<'a>,
223    {
224        let mut t = Set::<V, SIZE>::new();
225        while let Some(v) = seq.next_element()? {
226            t.insert_cow(v);
227        }
228        Ok(t)
229    }
230}
231
232#[cfg(feature = "serde")]
233impl<'a, V, const SIZE: usize> Deserialize<'a> for Set<V, SIZE>
234where
235    V: Deserialize<'a> + Clone + Ord,
236{
237    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
238    where
239        D: Deserializer<'a>,
240    {
241        deserializer.deserialize_seq(SetVisitor {
242            marker: PhantomData,
243        })
244    }
245}
246
247#[cfg(feature = "rayon")]
248impl<'a, V, const SIZE: usize> IntoParallelIterator for &'a Set<V, SIZE>
249where
250    V: 'a + Ord + Clone + Send + Sync,
251{
252    type Item = &'a V;
253    type Iter = rayon::vec::IntoIter<&'a V>;
254
255    fn into_par_iter(self) -> Self::Iter {
256        self.into_iter().collect::<Vec<_>>().into_par_iter()
257    }
258}
259
260#[cfg(feature = "rayon")]
261impl<V, const SIZE: usize> FromParallelIterator<V> for Set<V, SIZE>
262where
263    V: Ord + Clone + Send + Sync,
264{
265    fn from_par_iter<I>(i: I) -> Self
266    where
267        I: IntoParallelIterator<Item = V>,
268    {
269        i.into_par_iter()
270            .fold_with(Set::new(), |mut m, v| {
271                m.insert_cow(v);
272                m
273            })
274            .reduce_with(|m0, m1| m0.union(&m1))
275            .unwrap_or_else(Set::new)
276    }
277}
278
279impl<K, const SIZE: usize> Set<K, SIZE>
280where
281    K: Ord + Clone,
282{
283    /// Create a new empty set
284    pub fn new() -> Self {
285        Set(Tree::new())
286    }
287
288    /// The root of the set's tree, `None` when empty; see
289    /// [`Map::root`](crate::map::Map::root).
290    pub fn root(&self) -> Option<NodeRef<'_, K, (), SIZE>> {
291        self.0.root()
292    }
293
294    /// A set over the tree rooted at `root`.
295    pub fn from_root(root: Option<NodeHandle<K, (), SIZE>>) -> Self {
296        Set(Tree::from_root(root))
297    }
298
299    /// Create a weak reference to this set
300    pub fn downgrade(&self) -> WeakSetRef<K, SIZE> {
301        WeakSetRef(self.0.downgrade())
302    }
303
304    /// Return the number of strong references to this set (see Arc)
305    pub fn strong_count(&self) -> usize {
306        self.0.strong_count()
307    }
308
309    /// Return the number of weak references to this set (see Arc)
310    pub fn weak_count(&self) -> usize {
311        self.0.weak_count()
312    }
313
314    /// This will insert many elements at once, and is
315    /// potentially a lot faster than inserting one by one,
316    /// especially if the data is sorted.
317    ///
318    /// #Examples
319    ///```
320    /// use self::immutable_chunkmap::set::SetM;
321    ///
322    /// let mut v = vec![1, 10, -12, 44, 50];
323    /// v.sort_unstable();
324    ///
325    /// let m = SetM::new().insert_many(v.iter().map(|k| *k));
326    ///
327    /// for k in &v {
328    ///   assert_eq!(m.contains(k), true)
329    /// }
330    /// ```
331    pub fn insert_many<E: IntoIterator<Item = K>>(&self, elts: E) -> Self {
332        let root = self.0.insert_many(elts.into_iter().map(|k| (k, ())));
333        Set(root)
334    }
335
336    /// Remove multiple elements in a single pass. Similar performance
337    /// to insert_many.
338    pub fn remove_many<Q, E>(&self, elts: E) -> Self
339    where
340        Q: Ord,
341        K: Borrow<Q>,
342        E: IntoIterator<Item = Q>,
343    {
344        let root = self
345            .0
346            .update_many(elts.into_iter().map(|k| (k, ())), &mut |_, _, _| None);
347        Set(root)
348    }
349
350    /// This is just slightly wierd, however if you have a bunch of
351    /// borrowed forms of members of the set, and you want to look at
352    /// the real entries and possibly add/update/remove them, then
353    /// this method is for you.
354    pub fn update_many<Q, E, F>(&self, elts: E, mut f: F) -> Self
355    where
356        Q: Ord,
357        K: Borrow<Q>,
358        E: IntoIterator<Item = Q>,
359        F: FnMut(Q, Option<&K>) -> Option<K>,
360    {
361        let root =
362            self.0
363                .update_many(elts.into_iter().map(|k| (k, ())), &mut |q, (), cur| {
364                    let cur = cur.map(|(k, ())| k);
365                    f(q, cur).map(|k| (k, ()))
366                });
367        Set(root)
368    }
369
370    /// return a new set with k inserted into it. If k already
371    /// exists in the old set return true, else false. If the
372    /// element already exists in the set memory will not be
373    /// allocated.
374    pub fn insert(&self, k: K) -> (Self, bool) {
375        if self.contains(&k) {
376            (self.clone(), true)
377        } else {
378            (Set(self.0.insert(k, ()).0), false)
379        }
380    }
381
382    /// insert `k` with copy on write semantics. if `self` is a unique
383    /// reference to the set, then k will be inserted in
384    /// place. Otherwise, only the parts of the set necessary to
385    /// insert `k` will be copied, and then the copies will be
386    /// mutated. self will share all the parts that weren't modfied
387    /// with any previous clones.
388    pub fn insert_cow(&mut self, k: K) -> bool {
389        self.0.insert_cow(k, ()).is_some()
390    }
391
392    /// return true if the set contains k, else false. Runs in
393    /// log(N) time and constant space. where N is the size of
394    /// the set.
395    pub fn contains<'a, Q>(&'a self, k: &Q) -> bool
396    where
397        Q: ?Sized + Ord,
398        K: Borrow<Q>,
399    {
400        self.0.get(k).is_some()
401    }
402
403    /// return a reference to the item in the set that is equal to the
404    /// given value, or None if no such value exists.
405    pub fn get<'a, Q>(&'a self, k: &Q) -> Option<&'a K>
406    where
407        Q: ?Sized + Ord,
408        K: Borrow<Q>,
409    {
410        self.0.get_key(k)
411    }
412
413    /// return a new set with k removed. Runs in log(N) time
414    /// and log(N) space, where N is the size of the set
415    pub fn remove<Q: Sized + Ord>(&self, k: &Q) -> (Self, bool)
416    where
417        K: Borrow<Q>,
418    {
419        let (t, prev) = self.0.remove(k);
420        (Set(t), prev.is_some())
421    }
422
423    /// remove `k` from the set in place with copy on write semantics
424    /// (see `insert_cow`). return true if `k` was in the set.
425    pub fn remove_cow<Q: Sized + Ord>(&mut self, k: &Q) -> bool
426    where
427        K: Borrow<Q>,
428    {
429        self.0.remove_cow(k).is_some()
430    }
431
432    /// return the union of 2 sets. Runs in O(log(N) + M) time and
433    /// space, where N is the largest of the two sets, and M is the
434    /// number of chunks that intersect, which is roughly proportional
435    /// to the size of the intersection.
436    ///
437    /// # Examples
438    /// ```
439    /// use core::iter::FromIterator;
440    /// use self::immutable_chunkmap::set::SetM;
441    ///
442    /// let s0 = SetM::from_iter(0..10);
443    /// let s1 = SetM::from_iter(5..15);
444    /// let s2 = s0.union(&s1);
445    /// for i in 0..15 {
446    ///     assert!(s2.contains(&i));
447    /// }
448    /// ```
449    pub fn union(&self, other: &Set<K, SIZE>) -> Self {
450        Set(Tree::union(&self.0, &other.0, &mut |_, (), ()| Some(())))
451    }
452
453    /// return the intersection of 2 sets. Runs in O(log(N) + M) time
454    /// and space, where N is the smallest of the two sets, and M is
455    /// the number of intersecting chunks.
456    ///
457    /// # Examples
458    /// use core::iter::FromIterator;
459    /// use self::immutable_chunkmap::set::SetM;
460    ///
461    /// let s0 = SetM::from_iter(0..100);
462    /// let s1 = SetM::from_iter(20..50);
463    /// let s2 = s0.intersect(&s1);
464    ///
465    /// assert!(s2.len() == 30);
466    /// for i in 0..100 {
467    ///     if i < 20 || i >= 50 {
468    ///         assert!(!s2.contains(&i));
469    ///     } else {
470    ///         assert!(s2.contains(&i));
471    ///     }
472    /// }
473    pub fn intersect(&self, other: &Set<K, SIZE>) -> Self {
474        Set(Tree::intersect(
475            &self.0,
476            &other.0,
477            &mut |_, (), ()| Some(()),
478        ))
479    }
480
481    /// Return the difference of two sets. Runs in O(log(N) + M) time
482    /// and space, where N is the smallest of the two sets, and M is
483    /// the number of intersecting chunks.
484    ///
485    /// # Examples
486    /// ```
487    /// use core::iter::FromIterator;
488    /// use self::immutable_chunkmap::set::SetM;
489    ///
490    /// let s0 = SetM::from_iter(0..100);
491    /// let s1 = SetM::from_iter(0..50);
492    /// let s2 = s0.diff(&s1);
493    ///
494    /// assert!(s2.len() == 50);
495    /// for i in 0..50 {
496    ///     assert!(!s2.contains(&i));
497    /// }
498    /// for i in 50..100 {
499    ///     assert!(s2.contains(&i));
500    /// }
501    /// ```
502    pub fn diff(&self, other: &Set<K, SIZE>) -> Self
503    where
504        K: Debug,
505    {
506        Set(Tree::diff(&self.0, &other.0, &mut |_, (), ()| None))
507    }
508
509    /// get the number of elements in the map O(1) time and space
510    pub fn len(&self) -> usize {
511        self.0.len()
512    }
513
514    /// return an iterator over the subset of elements in the
515    /// set that are within the specified range.
516    ///
517    /// The returned iterator runs in O(log(N) + M) time, and
518    /// constant space. N is the number of elements in the
519    /// tree, and M is the number of elements you examine.
520    ///
521    /// if lbound >= ubound the returned iterator will be empty
522    pub fn range<'a, Q, R>(&'a self, r: R) -> SetIter<'a, R, Q, K, SIZE>
523    where
524        Q: Ord + ?Sized + 'a,
525        K: 'a + Clone + Ord + Borrow<Q>,
526        R: RangeBounds<Q> + 'a,
527    {
528        SetIter(self.0.range(r))
529    }
530}
531
532impl<K, const SIZE: usize> Set<K, SIZE>
533where
534    K: Ord + Clone + Debug,
535{
536    #[allow(dead_code)]
537    pub(crate) fn invariant(&self) -> () {
538        self.0.invariant()
539    }
540}