roaring/treemap/
inherent.rs

1use alloc::collections::btree_map::{BTreeMap, Entry};
2use core::iter;
3use core::ops::RangeBounds;
4
5use crate::IntegerTooSmall;
6use crate::RoaringBitmap;
7use crate::RoaringTreemap;
8
9use super::util;
10
11#[cfg(not(feature = "std"))]
12use alloc::vec::Vec;
13
14impl RoaringTreemap {
15    /// Creates an empty `RoaringTreemap`.
16    ///
17    /// # Examples
18    ///
19    /// ```rust
20    /// use roaring::RoaringTreemap;
21    /// let rb = RoaringTreemap::new();
22    /// ```
23    pub fn new() -> RoaringTreemap {
24        RoaringTreemap { map: BTreeMap::new() }
25    }
26
27    /// Creates a full `RoaringTreemap`.
28    ///
29    /// # Examples
30    ///
31    /// ```rust,ignore
32    /// use roaring::RoaringTreemap;
33    /// let rb = RoaringTreemap::full();
34    /// ```
35    pub fn full() -> RoaringTreemap {
36        RoaringTreemap { map: (0..=u32::MAX).zip(iter::repeat(RoaringBitmap::full())).collect() }
37    }
38
39    /// Adds a value to the set. Returns `true` if the value was not already present in the set.
40    ///
41    /// # Examples
42    ///
43    /// ```rust
44    /// use roaring::RoaringTreemap;
45    ///
46    /// let mut rb = RoaringTreemap::new();
47    /// assert_eq!(rb.insert(3), true);
48    /// assert_eq!(rb.insert(3), false);
49    /// assert_eq!(rb.contains(3), true);
50    /// ```
51    pub fn insert(&mut self, value: u64) -> bool {
52        let (hi, lo) = util::split(value);
53        self.map.entry(hi).or_default().insert(lo)
54    }
55
56    /// Inserts a range of values.
57    ///
58    /// Returns the number of inserted values.
59    ///
60    /// # Examples
61    ///
62    /// ```rust
63    /// use roaring::RoaringTreemap;
64    ///
65    /// let mut rb = RoaringTreemap::new();
66    /// rb.insert_range(2..4);
67    /// assert!(rb.contains(2));
68    /// assert!(rb.contains(3));
69    /// assert!(!rb.contains(4));
70    /// ```
71    pub fn insert_range<R: RangeBounds<u64>>(&mut self, range: R) -> u64 {
72        let (start, end) = match util::convert_range_to_inclusive(range) {
73            Some(range) => (*range.start(), *range.end()),
74            None => return 0,
75        };
76
77        let (start_hi, start_lo) = util::split(start);
78        let (end_hi, end_lo) = util::split(end);
79
80        let mut counter = 0u64;
81
82        // Split the input range by the leading 32 bits
83        for hi in start_hi..=end_hi {
84            let entry = self.map.entry(hi);
85
86            // Calculate the sub-range from the lower 32 bits
87            counter += if hi == end_hi && hi == start_hi {
88                entry.or_default().insert_range(start_lo..=end_lo)
89            } else if hi == start_hi {
90                entry.or_default().insert_range(start_lo..=u32::MAX)
91            } else if hi == end_hi {
92                entry.or_default().insert_range(0..=end_lo)
93            } else {
94                // We insert a full bitmap if it doesn't already exist and return the size of it.
95                // But if the bitmap already exists at this spot we replace it with a full bitmap
96                // and specify that we didn't inserted the integers from the previous bitmap.
97                let full_bitmap = RoaringBitmap::full();
98                match entry {
99                    Entry::Vacant(entry) => entry.insert(full_bitmap).len(),
100                    Entry::Occupied(mut entry) => {
101                        full_bitmap.len() - entry.insert(full_bitmap).len()
102                    }
103                }
104            };
105        }
106
107        counter
108    }
109
110    /// Pushes `value` in the treemap only if it is greater than the current maximum value.
111    ///
112    /// Returns whether the value was inserted.
113    ///
114    /// # Examples
115    ///
116    /// ```rust
117    /// use roaring::RoaringTreemap;
118    ///
119    /// let mut rb = RoaringTreemap::new();
120    /// assert!(rb.push(1));
121    /// assert!(rb.push(3));
122    /// assert_eq!(rb.push(3), false);
123    /// assert!(rb.push(5));
124    ///
125    /// assert_eq!(rb.iter().collect::<Vec<u64>>(), vec![1, 3, 5]);
126    /// ```
127    #[deprecated(since = "0.11.0", note = "use `try_push` instead")]
128    pub fn push(&mut self, value: u64) -> bool {
129        let (hi, lo) = util::split(value);
130        self.map.entry(hi).or_default().try_push(lo).is_ok()
131    }
132
133    /// Pushes `value` in the treemap only if it is greater than the current maximum value.
134    ///
135    /// Returns an error if the value is not greater than the current maximum value.
136    ///
137    /// # Examples
138    ///
139    /// ```rust
140    /// use roaring::{RoaringTreemap, IntegerTooSmall};
141    ///
142    /// let mut rb = RoaringTreemap::new();
143    /// assert!(rb.try_push(1).is_ok());
144    /// assert!(rb.try_push(3).is_ok());
145    /// assert_eq!(rb.try_push(3), Err(IntegerTooSmall));
146    /// assert!(rb.try_push(5).is_ok());
147    ///
148    /// assert_eq!(rb.iter().collect::<Vec<u64>>(), vec![1, 3, 5]);
149    /// ```
150    pub fn try_push(&mut self, value: u64) -> Result<(), IntegerTooSmall> {
151        let (hi, lo) = util::split(value);
152        self.map.entry(hi).or_default().try_push(lo)
153    }
154
155    /// Pushes `value` in the treemap only if it is greater than the current maximum value.
156    /// It is up to the caller to have validated index > self.max()
157    ///
158    /// # Panics
159    ///
160    /// If debug_assertions enabled and index is > self.max()
161    pub(crate) fn push_unchecked(&mut self, value: u64) {
162        let (hi, lo) = util::split(value);
163        // BTreeMap last_mut not stabilized see https://github.com/rust-lang/rust/issues/62924
164        match self.map.iter_mut().next_back() {
165            Some((&key, bitmap)) if key == hi => bitmap.push_unchecked(lo),
166            Some((&key, _)) if cfg!(debug_assertions) && key > hi => {
167                panic!("last bitmap key > key of value")
168            }
169            _otherwise => {
170                // The tree is empty
171                let mut rb = RoaringBitmap::new();
172                rb.push_unchecked(lo);
173                self.map.insert(hi, rb);
174            }
175        }
176    }
177
178    /// Removes a value from the set. Returns `true` if the value was present in the set.
179    ///
180    /// # Examples
181    ///
182    /// ```rust
183    /// use roaring::RoaringTreemap;
184    ///
185    /// let mut rb = RoaringTreemap::new();
186    /// rb.insert(3);
187    /// assert_eq!(rb.remove(3), true);
188    /// assert_eq!(rb.remove(3), false);
189    /// assert_eq!(rb.contains(3), false);
190    /// ```
191    pub fn remove(&mut self, value: u64) -> bool {
192        let (hi, lo) = util::split(value);
193        match self.map.entry(hi) {
194            Entry::Vacant(_) => false,
195            Entry::Occupied(mut ent) => {
196                if ent.get_mut().remove(lo) {
197                    if ent.get().is_empty() {
198                        ent.remove();
199                    }
200                    true
201                } else {
202                    false
203                }
204            }
205        }
206    }
207
208    /// Removes a range of values.
209    /// Returns the number of removed values.
210    ///
211    /// # Examples
212    ///
213    /// ```rust
214    /// use roaring::RoaringTreemap;
215    ///
216    /// let mut rb = RoaringTreemap::new();
217    /// rb.insert(2);
218    /// rb.insert(3);
219    /// assert_eq!(rb.remove_range(2..4), 2);
220    /// ```
221    pub fn remove_range<R>(&mut self, range: R) -> u64
222    where
223        R: RangeBounds<u64>,
224    {
225        let (start, end) = match util::convert_range_to_inclusive(range) {
226            Some(range) => (*range.start(), *range.end()),
227            None => return 0,
228        };
229
230        let (start_container_key, start_index) = util::split(start);
231        let (end_container_key, end_index) = util::split(end);
232
233        let mut keys_to_remove = Vec::new();
234        let mut removed = 0;
235
236        for (&key, rb) in &mut self.map {
237            if key >= start_container_key && key <= end_container_key {
238                let a = if key == start_container_key { start_index } else { 0 };
239                let b = if key == end_container_key { end_index } else { u32::MAX };
240                removed += rb.remove_range(a..=b);
241                if rb.is_empty() {
242                    keys_to_remove.push(key);
243                }
244            }
245        }
246
247        for key in keys_to_remove {
248            self.map.remove(&key);
249        }
250
251        removed
252    }
253
254    /// Returns `true` if this set contains the specified integer.
255    ///
256    /// # Examples
257    ///
258    /// ```rust
259    /// use roaring::RoaringTreemap;
260    ///
261    /// let mut rb = RoaringTreemap::new();
262    /// rb.insert(1);
263    /// assert_eq!(rb.contains(0), false);
264    /// assert_eq!(rb.contains(1), true);
265    /// assert_eq!(rb.contains(100), false);
266    /// ```
267    pub fn contains(&self, value: u64) -> bool {
268        let (hi, lo) = util::split(value);
269        match self.map.get(&hi) {
270            None => false,
271            Some(r) => r.contains(lo),
272        }
273    }
274
275    /// Clears all integers in this set.
276    ///
277    /// # Examples
278    ///
279    /// ```rust
280    /// use roaring::RoaringTreemap;
281    ///
282    /// let mut rb = RoaringTreemap::new();
283    /// rb.insert(1);
284    /// assert_eq!(rb.contains(1), true);
285    /// rb.clear();
286    /// assert_eq!(rb.contains(1), false);
287    /// ```
288    pub fn clear(&mut self) {
289        self.map.clear();
290    }
291
292    /// Returns `true` if there are no integers in this set.
293    ///
294    /// # Examples
295    ///
296    /// ```rust
297    /// use roaring::RoaringTreemap;
298    ///
299    /// let mut rb = RoaringTreemap::new();
300    /// assert_eq!(rb.is_empty(), true);
301    ///
302    /// rb.insert(3);
303    /// assert_eq!(rb.is_empty(), false);
304    /// ```
305    pub fn is_empty(&self) -> bool {
306        self.map.values().all(RoaringBitmap::is_empty)
307    }
308
309    /// Returns `true` if there are every possible integers in this set.
310    ///
311    /// # Examples
312    ///
313    /// ```rust,ignore
314    /// use roaring::RoaringTreemap;
315    ///
316    /// let mut rb = RoaringTreemap::full();
317    /// assert!(!rb.is_empty());
318    /// assert!(rb.is_full());
319    /// ```
320    pub fn is_full(&self) -> bool {
321        self.map.len() == (u32::MAX as usize + 1) && self.map.values().all(RoaringBitmap::is_full)
322    }
323
324    /// Returns the number of distinct integers added to the set.
325    ///
326    /// # Examples
327    ///
328    /// ```rust
329    /// use roaring::RoaringTreemap;
330    ///
331    /// let mut rb = RoaringTreemap::new();
332    /// assert_eq!(rb.len(), 0);
333    ///
334    /// rb.insert(3);
335    /// assert_eq!(rb.len(), 1);
336    ///
337    /// rb.insert(3);
338    /// rb.insert(4);
339    /// assert_eq!(rb.len(), 2);
340    /// ```
341    pub fn len(&self) -> u64 {
342        self.map.values().map(RoaringBitmap::len).sum()
343    }
344
345    /// Returns the minimum value in the set (if the set is non-empty).
346    ///
347    /// # Examples
348    ///
349    /// ```rust
350    /// use roaring::RoaringTreemap;
351    ///
352    /// let mut rb = RoaringTreemap::new();
353    /// assert_eq!(rb.min(), None);
354    ///
355    /// rb.insert(3);
356    /// rb.insert(4);
357    /// assert_eq!(rb.min(), Some(3));
358    /// ```
359    pub fn min(&self) -> Option<u64> {
360        self.map
361            .iter()
362            .find(|&(_, rb)| rb.min().is_some())
363            .map(|(k, rb)| util::join(*k, rb.min().unwrap()))
364    }
365
366    /// Returns the maximum value in the set (if the set is non-empty).
367    ///
368    /// # Examples
369    ///
370    /// ```rust
371    /// use roaring::RoaringTreemap;
372    ///
373    /// let mut rb = RoaringTreemap::new();
374    /// assert_eq!(rb.max(), None);
375    ///
376    /// rb.insert(3);
377    /// rb.insert(4);
378    /// assert_eq!(rb.max(), Some(4));
379    /// ```
380    pub fn max(&self) -> Option<u64> {
381        self.map
382            .iter()
383            .rev()
384            .find(|&(_, rb)| rb.max().is_some())
385            .map(|(k, rb)| util::join(*k, rb.max().unwrap()))
386    }
387
388    /// Returns the number of integers that are <= value. rank(u64::MAX) == len()
389    ///
390    /// # Examples
391    ///
392    /// ```rust
393    /// use roaring::RoaringTreemap;
394    ///
395    /// let mut rb = RoaringTreemap::new();
396    /// assert_eq!(rb.rank(0), 0);
397    ///
398    /// rb.insert(3);
399    /// rb.insert(4);
400    /// assert_eq!(rb.rank(3), 1);
401    /// assert_eq!(rb.rank(10), 2)
402    /// ```
403    pub fn rank(&self, value: u64) -> u64 {
404        // if len becomes cached for RoaringTreemap: return len if len > value
405
406        let (hi, lo) = util::split(value);
407        let mut iter = self.map.range(..=hi).rev();
408
409        iter.next()
410            .map(|(&k, bitmap)| if k == hi { bitmap.rank(lo) } else { bitmap.len() })
411            .unwrap_or(0)
412            + iter.map(|(_, bitmap)| bitmap.len()).sum::<u64>()
413    }
414
415    /// Returns the `n`th integer in the set or `None` if `n >= len()`
416    ///
417    /// # Examples
418    ///
419    /// ```rust
420    /// use roaring::RoaringTreemap;
421    ///
422    /// let mut rb = RoaringTreemap::new();
423    /// assert_eq!(rb.select(0), None);
424    ///
425    /// rb.append(vec![0, 10, 100]);
426    ///
427    /// assert_eq!(rb.select(0), Some(0));
428    /// assert_eq!(rb.select(1), Some(10));
429    /// assert_eq!(rb.select(2), Some(100));
430    /// assert_eq!(rb.select(3), None);
431    /// ```
432    pub fn select(&self, mut n: u64) -> Option<u64> {
433        for (&key, bitmap) in &self.map {
434            let len = bitmap.len();
435            if len > n {
436                return Some(((key as u64) << 32) | bitmap.select(n as u32).unwrap() as u64);
437            }
438            n -= len;
439        }
440
441        None
442    }
443}
444
445impl Default for RoaringTreemap {
446    fn default() -> RoaringTreemap {
447        RoaringTreemap::new()
448    }
449}
450
451impl Clone for RoaringTreemap {
452    fn clone(&self) -> Self {
453        RoaringTreemap { map: self.map.clone() }
454    }
455
456    fn clone_from(&mut self, other: &Self) {
457        self.map.clone_from(&other.map);
458    }
459}