immutable_chunkmap/
set.rs

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