roaring/treemap/
iter.rs

1use alloc::collections::{btree_map, BTreeMap};
2use core::iter;
3use core::ops::Add;
4
5use super::util;
6use crate::bitmap::IntoIter as IntoIter32;
7use crate::bitmap::Iter as Iter32;
8use crate::{NonSortedIntegers, RoaringBitmap, RoaringTreemap};
9
10struct To64Iter<'a> {
11    hi: u32,
12    inner: Iter32<'a>,
13}
14
15impl To64Iter<'_> {
16    fn advance_to(&mut self, n: u32) {
17        self.inner.advance_to(n)
18    }
19
20    fn advance_back_to(&mut self, n: u32) {
21        self.inner.advance_back_to(n)
22    }
23}
24
25impl Iterator for To64Iter<'_> {
26    type Item = u64;
27    fn next(&mut self) -> Option<u64> {
28        self.inner.next().map(|n| util::join(self.hi, n))
29    }
30
31    fn size_hint(&self) -> (usize, Option<usize>) {
32        self.inner.size_hint()
33    }
34
35    #[inline]
36    fn fold<B, F>(self, init: B, mut f: F) -> B
37    where
38        Self: Sized,
39        F: FnMut(B, Self::Item) -> B,
40    {
41        self.inner.fold(init, move |b, lo| f(b, ((self.hi as u64) << 32) + (lo as u64)))
42    }
43}
44
45impl DoubleEndedIterator for To64Iter<'_> {
46    fn next_back(&mut self) -> Option<Self::Item> {
47        self.inner.next_back().map(|n| util::join(self.hi, n))
48    }
49
50    #[inline]
51    fn rfold<B, F>(self, init: B, mut f: F) -> B
52    where
53        Self: Sized,
54        F: FnMut(B, Self::Item) -> B,
55    {
56        self.inner.rfold(init, move |b, lo| f(b, ((self.hi as u64) << 32) + (lo as u64)))
57    }
58}
59
60fn to64iter(t: (u32, &RoaringBitmap)) -> To64Iter<'_> {
61    To64Iter { hi: t.0, inner: t.1.iter() }
62}
63
64struct To64IntoIter {
65    hi: u32,
66    inner: IntoIter32,
67}
68
69impl Iterator for To64IntoIter {
70    type Item = u64;
71    fn next(&mut self) -> Option<u64> {
72        self.inner.next().map(|n| util::join(self.hi, n))
73    }
74
75    #[inline]
76    fn fold<B, F>(self, init: B, mut f: F) -> B
77    where
78        Self: Sized,
79        F: FnMut(B, Self::Item) -> B,
80    {
81        self.inner.fold(init, move |b, lo| f(b, ((self.hi as u64) << 32) + (lo as u64)))
82    }
83}
84
85impl DoubleEndedIterator for To64IntoIter {
86    fn next_back(&mut self) -> Option<Self::Item> {
87        self.inner.next_back().map(|n| util::join(self.hi, n))
88    }
89
90    #[inline]
91    fn rfold<B, F>(self, init: B, mut f: F) -> B
92    where
93        Self: Sized,
94        F: FnMut(B, Self::Item) -> B,
95    {
96        self.inner.rfold(init, move |b, lo| f(b, ((self.hi as u64) << 32) + (lo as u64)))
97    }
98}
99
100fn to64intoiter(t: (u32, RoaringBitmap)) -> To64IntoIter {
101    To64IntoIter { hi: t.0, inner: t.1.into_iter() }
102}
103
104type InnerIntoIter = iter::FlatMap<
105    btree_map::IntoIter<u32, RoaringBitmap>,
106    To64IntoIter,
107    fn((u32, RoaringBitmap)) -> To64IntoIter,
108>;
109
110/// An iterator for `RoaringTreemap`.
111pub struct Iter<'a> {
112    outer: BitmapIter<'a>,
113    front: Option<To64Iter<'a>>,
114    back: Option<To64Iter<'a>>,
115}
116
117/// An iterator for `RoaringTreemap`.
118pub struct IntoIter {
119    inner: InnerIntoIter,
120    size_hint: u64,
121}
122
123impl Iter<'_> {
124    fn new(map: &'_ BTreeMap<u32, RoaringBitmap>) -> Iter<'_> {
125        let outer = BitmapIter::new(map);
126        Iter { outer, front: None, back: None }
127    }
128
129    /// Advance the iterator to the first position where the item has a value >= `n`
130    ///
131    /// # Examples
132    ///
133    /// ```rust
134    /// use roaring::RoaringTreemap;
135    /// use core::iter::FromIterator;
136    ///
137    /// let bitmap = (1..3).collect::<RoaringTreemap>();
138    /// let mut iter = bitmap.iter();
139    /// iter.advance_to(2);
140    ///
141    /// assert_eq!(iter.next(), Some(2));
142    /// assert_eq!(iter.next(), None);
143    /// ```
144    pub fn advance_to(&mut self, n: u64) {
145        let (key, index) = util::split(n);
146
147        self.outer.advance_to(key);
148
149        if self.front.is_none() {
150            let Some(next) = self.outer.next() else {
151                // if the current front iterator is empty or not yet initialized,
152                // but the outer bitmap iterator is empty, then consume the back
153                // iterator from the front if it is not also exhausted
154                if let Some(ref mut back) = self.back {
155                    back.advance_to(index);
156                }
157                return;
158            };
159            self.front = Some(to64iter(next));
160        }
161
162        if let Some(ref mut front) = self.front {
163            front.advance_to(index);
164        }
165    }
166
167    /// Advance the back of the iterator to the first position where the item has a value <= `n`
168    ///
169    /// # Examples
170    ///
171    /// ```rust
172    /// use roaring::RoaringTreemap;
173    /// use core::iter::FromIterator;
174    ///
175    /// let bitmap = (1..3).collect::<RoaringTreemap>();
176    /// let mut iter = bitmap.iter();
177    /// iter.advance_back_to(1);
178    ///
179    /// assert_eq!(iter.next_back(), Some(1));
180    /// assert_eq!(iter.next_back(), None);
181    /// ```
182    pub fn advance_back_to(&mut self, n: u64) {
183        let (key, index) = util::split(n);
184
185        self.outer.advance_back_to(key);
186
187        if self.back.is_none() {
188            let Some(next_back) = self.outer.next_back() else {
189                // if the current back iterator is empty or not yet initialized,
190                // but the outer bitmap iterator is empty, then consume the front
191                // iterator from the back if it is not also exhausted
192                if let Some(ref mut front) = self.front {
193                    front.advance_back_to(index);
194                }
195                return;
196            };
197            self.back = Some(to64iter(next_back));
198        }
199
200        if let Some(ref mut back) = self.back {
201            back.advance_back_to(index);
202        }
203    }
204}
205
206impl IntoIter {
207    fn new(map: BTreeMap<u32, RoaringBitmap>) -> IntoIter {
208        let size_hint = map.values().map(|r| r.len()).sum();
209        let i = map.into_iter().flat_map(to64intoiter as _);
210        IntoIter { inner: i, size_hint }
211    }
212}
213
214impl Iterator for Iter<'_> {
215    type Item = u64;
216
217    fn next(&mut self) -> Option<u64> {
218        if let Some(ref mut front) = &mut self.front {
219            if let Some(inner) = front.next() {
220                return Some(inner);
221            }
222        }
223
224        let Some(outer_next) = self.outer.next() else {
225            // if the current front iterator is empty or not yet initialized,
226            // but the outer bitmap iterator is empty, then consume the back
227            // iterator from the front if it is not also exhausted
228            if let Some(ref mut back) = &mut self.back {
229                if let Some(next) = back.next() {
230                    return Some(next);
231                }
232            }
233            return None;
234        };
235
236        self.front = Some(to64iter(outer_next));
237        self.next()
238    }
239
240    fn size_hint(&self) -> (usize, Option<usize>) {
241        let front_size_hint = self.front.as_ref().map_or(0, |f| f.size_hint().0);
242        let back_size_hint = self.back.as_ref().map_or(0, |b| b.size_hint().0);
243
244        let size_hint = front_size_hint
245            .saturating_add(back_size_hint)
246            .saturating_add(self.outer.remaining() as usize);
247
248        (size_hint, Some(size_hint))
249    }
250}
251
252impl DoubleEndedIterator for Iter<'_> {
253    fn next_back(&mut self) -> Option<Self::Item> {
254        if let Some(ref mut back) = &mut self.back {
255            if let Some(inner) = back.next_back() {
256                return Some(inner);
257            }
258        }
259
260        let Some(outer_next_back) = self.outer.next_back() else {
261            // if the current back iterator is empty or not yet initialized,
262            // but the outer bitmap iterator is empty, then consume the front
263            // iterator from the back if it is not also exhausted
264            if let Some(ref mut front) = &mut self.front {
265                if let Some(next_back) = front.next_back() {
266                    return Some(next_back);
267                }
268            }
269            return None;
270        };
271
272        self.back = Some(to64iter(outer_next_back));
273        self.next_back()
274    }
275}
276
277#[cfg(target_pointer_width = "64")]
278impl ExactSizeIterator for Iter<'_> {
279    fn len(&self) -> usize {
280        self.size_hint().0
281    }
282}
283
284impl Iterator for IntoIter {
285    type Item = u64;
286
287    fn next(&mut self) -> Option<u64> {
288        self.size_hint = self.size_hint.saturating_sub(1);
289        self.inner.next()
290    }
291
292    fn size_hint(&self) -> (usize, Option<usize>) {
293        if self.size_hint < usize::MAX as u64 {
294            (self.size_hint as usize, Some(self.size_hint as usize))
295        } else {
296            (usize::MAX, None)
297        }
298    }
299
300    #[inline]
301    fn fold<B, F>(self, init: B, f: F) -> B
302    where
303        Self: Sized,
304        F: FnMut(B, Self::Item) -> B,
305    {
306        self.inner.fold(init, f)
307    }
308}
309
310impl DoubleEndedIterator for IntoIter {
311    fn next_back(&mut self) -> Option<Self::Item> {
312        self.size_hint = self.size_hint.saturating_sub(1);
313        self.inner.next_back()
314    }
315
316    #[inline]
317    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
318    where
319        Fold: FnMut(Acc, Self::Item) -> Acc,
320    {
321        self.inner.rfold(init, fold)
322    }
323}
324
325#[cfg(target_pointer_width = "64")]
326impl ExactSizeIterator for IntoIter {
327    fn len(&self) -> usize {
328        self.size_hint as usize
329    }
330}
331
332impl RoaringTreemap {
333    /// Iterator over each value stored in the RoaringTreemap, guarantees values are ordered by
334    /// value.
335    ///
336    /// # Examples
337    ///
338    /// ```rust
339    /// use roaring::RoaringTreemap;
340    /// use core::iter::FromIterator;
341    ///
342    /// let bitmap = (1..3).collect::<RoaringTreemap>();
343    /// let mut iter = bitmap.iter();
344    ///
345    /// assert_eq!(iter.next(), Some(1));
346    /// assert_eq!(iter.next(), Some(2));
347    /// assert_eq!(iter.next(), None);
348    /// ```
349    pub fn iter(&'_ self) -> Iter<'_> {
350        Iter::new(&self.map)
351    }
352
353    /// Iterator over pairs of partition number and the corresponding RoaringBitmap.
354    /// The partition number is defined by the 32 most significant bits of the bit index.
355    ///
356    /// # Examples
357    ///
358    /// ```rust
359    /// use roaring::{RoaringBitmap, RoaringTreemap};
360    /// use core::iter::FromIterator;
361    ///
362    /// let original = (0..6000).collect::<RoaringTreemap>();
363    /// let mut bitmaps = original.bitmaps();
364    ///
365    /// assert_eq!(bitmaps.next(), Some((0, &(0..6000).collect::<RoaringBitmap>())));
366    /// assert_eq!(bitmaps.next(), None);
367    /// ```
368    pub fn bitmaps(&'_ self) -> BitmapIter<'_> {
369        BitmapIter::new(&self.map)
370    }
371
372    /// Construct a RoaringTreemap from an iterator of partition number and RoaringBitmap pairs.
373    /// The partition number is defined by the 32 most significant bits of the bit index.
374    /// Note that repeated partitions, if present, will replace previously set partitions.
375    ///
376    /// # Examples
377    ///
378    /// ```rust
379    /// use roaring::RoaringTreemap;
380    /// use core::iter::FromIterator;
381    ///
382    /// let original = (0..6000).collect::<RoaringTreemap>();
383    /// let clone = RoaringTreemap::from_bitmaps(original.bitmaps().map(|(p, b)| (p, b.clone())));
384    ///
385    /// assert_eq!(clone, original);
386    /// ```
387    pub fn from_bitmaps<I: IntoIterator<Item = (u32, RoaringBitmap)>>(iterator: I) -> Self {
388        RoaringTreemap { map: iterator.into_iter().collect() }
389    }
390}
391
392impl<'a> IntoIterator for &'a RoaringTreemap {
393    type Item = u64;
394    type IntoIter = Iter<'a>;
395
396    fn into_iter(self) -> Iter<'a> {
397        self.iter()
398    }
399}
400
401impl IntoIterator for RoaringTreemap {
402    type Item = u64;
403    type IntoIter = IntoIter;
404
405    fn into_iter(self) -> IntoIter {
406        IntoIter::new(self.map)
407    }
408}
409
410impl<const N: usize> From<[u64; N]> for RoaringTreemap {
411    fn from(arr: [u64; N]) -> Self {
412        RoaringTreemap::from_iter(arr)
413    }
414}
415
416impl FromIterator<u64> for RoaringTreemap {
417    fn from_iter<I: IntoIterator<Item = u64>>(iterator: I) -> RoaringTreemap {
418        let mut rb = RoaringTreemap::new();
419        rb.extend(iterator);
420        rb
421    }
422}
423
424impl<'a> FromIterator<&'a u64> for RoaringTreemap {
425    fn from_iter<I: IntoIterator<Item = &'a u64>>(iterator: I) -> RoaringTreemap {
426        let mut rb = RoaringTreemap::new();
427        rb.extend(iterator);
428        rb
429    }
430}
431
432impl Extend<u64> for RoaringTreemap {
433    fn extend<I: IntoIterator<Item = u64>>(&mut self, iterator: I) {
434        for value in iterator {
435            self.insert(value);
436        }
437    }
438}
439
440impl<'a> Extend<&'a u64> for RoaringTreemap {
441    fn extend<I: IntoIterator<Item = &'a u64>>(&mut self, iterator: I) {
442        for value in iterator {
443            self.insert(*value);
444        }
445    }
446}
447
448impl RoaringTreemap {
449    /// Create the set from a sorted iterator. Values must be sorted and deduplicated.
450    ///
451    /// The values of the iterator must be ordered and strictly greater than the greatest value
452    /// in the set. If a value in the iterator doesn't satisfy this requirement, it is not added
453    /// and the append operation is stopped.
454    ///
455    /// Returns `Ok` with the requested `RoaringTreemap`, `Err` with the number of elements
456    /// we tried to append before an error occurred.
457    ///
458    /// # Examples
459    ///
460    /// ```rust
461    /// use roaring::RoaringTreemap;
462    ///
463    /// let mut rb = RoaringTreemap::from_sorted_iter(0..10).unwrap();
464    ///
465    /// assert!(rb.iter().eq(0..10));
466    /// ```
467    pub fn from_sorted_iter<I: IntoIterator<Item = u64>>(
468        iterator: I,
469    ) -> Result<RoaringTreemap, NonSortedIntegers> {
470        let mut rt = RoaringTreemap::new();
471        rt.append(iterator).map(|_| rt)
472    }
473
474    /// Extend the set with a sorted iterator.
475    ///
476    /// The values of the iterator must be ordered and strictly greater than the greatest value
477    /// in the set. If a value in the iterator doesn't satisfy this requirement, it is not added
478    /// and the append operation is stopped.
479    ///
480    /// Returns `Ok` with the number of elements appended to the set, `Err` with
481    /// the number of elements we effectively appended before an error occurred.
482    ///
483    /// # Examples
484    ///
485    /// ```rust
486    /// use roaring::RoaringTreemap;
487    ///
488    /// let mut rb = RoaringTreemap::new();
489    /// rb.append(0..10);
490    ///
491    /// assert!(rb.iter().eq(0..10));
492    /// ```
493    pub fn append<I: IntoIterator<Item = u64>>(
494        &mut self,
495        iterator: I,
496    ) -> Result<u64, NonSortedIntegers> {
497        let mut iterator = iterator.into_iter();
498        let mut prev = match (iterator.next(), self.max()) {
499            (None, _) => return Ok(0),
500            (Some(first), Some(max)) if first <= max => {
501                return Err(NonSortedIntegers { valid_until: 0 })
502            }
503            (Some(first), _) => first,
504        };
505
506        // It is now guaranteed that so long as the values of the iterator are
507        // monotonically increasing they must also be the greatest in the set.
508
509        self.push_unchecked(prev);
510
511        let mut count = 1;
512        for value in iterator {
513            if value <= prev {
514                return Err(NonSortedIntegers { valid_until: count });
515            } else {
516                self.push_unchecked(value);
517                prev = value;
518                count += 1;
519            }
520        }
521
522        Ok(count)
523    }
524}
525
526/// An iterator of `RoaringBitmap`s for `RoaringTreemap`.
527pub struct BitmapIter<'a> {
528    treemap: &'a BTreeMap<u32, RoaringBitmap>,
529    range: btree_map::Range<'a, u32, RoaringBitmap>,
530    latest_front_idx: Option<u32>,
531    latest_back_idx: Option<u32>,
532}
533
534impl<'a> BitmapIter<'a> {
535    fn new(treemap: &'a BTreeMap<u32, RoaringBitmap>) -> Self {
536        let range = treemap.range(..);
537        Self { treemap, range, latest_back_idx: None, latest_front_idx: None }
538    }
539
540    fn advance_to(&mut self, new_front_idx: u32) {
541        match self.latest_back_idx {
542            Some(latest_back_idx) => match self.latest_front_idx {
543                Some(last_idx) if last_idx >= new_front_idx => {}
544                _ => {
545                    // if asked to advance to beyond the back iterator,
546                    // update the self.range iterator to be empty
547                    if new_front_idx >= latest_back_idx {
548                        self.range = self.treemap.range(0..1);
549                        self.range.next_back();
550                    } else {
551                        // otherwise shrink the remaining range from the front
552                        self.range = self.treemap.range(new_front_idx..latest_back_idx);
553                    }
554
555                    // self.range = self.treemap.range(new_front_idx..latest_back_idx);
556                }
557            },
558            None => match self.latest_front_idx {
559                Some(latest_idx) if latest_idx >= new_front_idx => {}
560                _ => {
561                    self.range = self.treemap.range(new_front_idx..);
562                }
563            },
564        }
565    }
566
567    fn advance_back_to(&mut self, new_back_idx: u32) {
568        match self.latest_front_idx {
569            Some(latest_front_idx) => match self.latest_back_idx {
570                // do nothing if asked to advance back to a higher index than the back is already at
571                Some(latest_back_idx) if latest_back_idx <= new_back_idx => {}
572                _ => {
573                    // if asked to advance back to beyond the front iterator,
574                    // update the self.range iterator to be empty
575                    if new_back_idx <= latest_front_idx {
576                        self.range = self.treemap.range(0..1);
577                        self.range.next_back();
578                    } else {
579                        // otherwise shrink the remaining range from the back
580                        self.range = self.treemap.range((latest_front_idx + 1)..new_back_idx);
581                    }
582                }
583            },
584            None => match self.latest_back_idx {
585                Some(latest_back_idx) if latest_back_idx <= new_back_idx => {}
586                _ => {
587                    self.range = self.treemap.range(..=new_back_idx);
588                }
589            },
590        }
591    }
592
593    fn remaining(&self) -> u64 {
594        let range = self.range.clone();
595        range.fold(0, |acc, (_, bitmap)| acc.add(bitmap.len()))
596    }
597}
598
599impl<'a> Iterator for BitmapIter<'a> {
600    type Item = (u32, &'a RoaringBitmap);
601
602    fn next(&mut self) -> Option<Self::Item> {
603        match self.range.next().map(|(&p, b)| (p, b)) {
604            None => {
605                self.latest_front_idx = None;
606                None
607            }
608            Some((next_idx, next_map)) => {
609                self.latest_front_idx = Some(next_idx);
610                Some((next_idx, next_map))
611            }
612        }
613    }
614
615    fn size_hint(&self) -> (usize, Option<usize>) {
616        self.range.size_hint()
617    }
618}
619
620impl FromIterator<(u32, RoaringBitmap)> for RoaringTreemap {
621    fn from_iter<I: IntoIterator<Item = (u32, RoaringBitmap)>>(iterator: I) -> RoaringTreemap {
622        Self::from_bitmaps(iterator)
623    }
624}
625
626impl DoubleEndedIterator for BitmapIter<'_> {
627    fn next_back(&mut self) -> Option<Self::Item> {
628        match self.range.next_back().map(|(&p, b)| (p, b)) {
629            None => {
630                self.latest_back_idx = None;
631                None
632            }
633            Some((next_back_idx, next_back_map)) => {
634                self.latest_back_idx = Some(next_back_idx);
635                Some((next_back_idx, next_back_map))
636            }
637        }
638    }
639}
roaring/treemap/iter.rs

roaring/treemap/
iter.rs
