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}