croaring/treemap/
iter.rs

1use super::util;
2use crate::bitmap::BitmapIterator;
3use crate::{Bitmap, Treemap};
4use alloc::collections::btree_map;
5use core::iter;
6
7struct To64Iter<'a> {
8    key: u32,
9    iterator: BitmapIterator<'a>,
10}
11
12impl<'a> Iterator for To64Iter<'a> {
13    type Item = u64;
14
15    fn next(&mut self) -> Option<u64> {
16        self.iterator.next().map(|n| util::join(self.key, n))
17    }
18}
19
20fn to64iter<'a>((key, bitmap): (&'a u32, &'a Bitmap)) -> To64Iter<'a> {
21    assert!(!bitmap.is_empty(), "empty bitmap at {key}");
22    To64Iter {
23        key: *key,
24        iterator: bitmap.iter(),
25    }
26}
27
28type InnerIter<'a> = iter::FlatMap<
29    btree_map::Iter<'a, u32, Bitmap>,
30    To64Iter<'a>,
31    fn((&'a u32, &'a Bitmap)) -> To64Iter<'a>,
32>;
33
34/// Iterator over values stored in the treemap
35///
36/// Values are ordered in ascending order
37pub struct TreemapIterator<'a> {
38    iter: InnerIter<'a>,
39}
40
41impl<'a> TreemapIterator<'a> {
42    fn new(treemap: &'a Treemap) -> Self {
43        let iter = treemap.map.iter().flat_map(to64iter as _);
44
45        TreemapIterator { iter }
46    }
47}
48
49impl<'a> Iterator for TreemapIterator<'a> {
50    type Item = u64;
51
52    fn next(&mut self) -> Option<u64> {
53        self.iter.next()
54    }
55}
56
57impl Treemap {
58    /// Returns an iterator over each value stored in the bitmap.
59    /// Returned values are ordered in ascending order.
60    ///
61    /// # Examples
62    ///
63    /// ```
64    /// use std::u64;
65    /// use croaring::Treemap;
66    ///
67    /// let mut treemap = Treemap::new();
68    /// treemap.add(4);
69    /// treemap.add(3);
70    /// treemap.add(2);
71    /// treemap.add(2);
72    /// treemap.add(u64::MAX);
73    /// let mut iterator = treemap.iter();
74    ///
75    /// assert_eq!(iterator.next(), Some(2));
76    /// assert_eq!(iterator.next(), Some(3));
77    /// assert_eq!(iterator.next(), Some(4));
78    /// assert_eq!(iterator.next(), Some(u64::MAX));
79    /// assert_eq!(iterator.next(), None);
80    /// ```
81    #[must_use]
82    pub fn iter(&self) -> TreemapIterator<'_> {
83        TreemapIterator::new(self)
84    }
85}
86
87impl FromIterator<u64> for Treemap {
88    /// Convenience method for creating treemap from an iterator.
89    ///
90    /// # Examples
91    ///
92    /// ```
93    /// use std::{u32, u64};
94    /// use croaring::Treemap;
95    ///
96    /// let treemap: Treemap = (1..3).chain(u64::from(u32::MAX)+1..u64::from(u32::MAX)+10).collect();
97    ///
98    /// assert!(!treemap.is_empty());
99    /// assert!(treemap.contains(1));
100    /// assert!(treemap.contains(2));
101    /// assert!(treemap.contains(u64::from(u32::MAX)+1));
102    /// assert!(treemap.contains(u64::from(u32::MAX)+5));
103    /// assert_eq!(treemap.cardinality(), 11);
104    /// ```
105    fn from_iter<I: IntoIterator<Item = u64>>(iter: I) -> Self {
106        let mut result = Self::new();
107        result.extend(iter);
108        result
109    }
110}
111
112impl Extend<u64> for Treemap {
113    fn extend<T: IntoIterator<Item = u64>>(&mut self, iter: T) {
114        // Potentially reduce outer map lookups by optimistically
115        // assuming that adjacent values will belong to the same inner bitmap.
116        let mut last_bitmap: Option<(u32, &mut Bitmap)> = None;
117        for item in iter {
118            let (high, low) = util::split(item);
119            if let Some((last_high, ref mut last_bitmap)) = last_bitmap {
120                if last_high == high {
121                    last_bitmap.add(low);
122                    continue;
123                }
124            }
125            let bitmap = self.get_or_create(high);
126            bitmap.add(low);
127            last_bitmap = Some((high, bitmap));
128        }
129    }
130}
croaring/treemap/iter.rs

croaring/treemap/
iter.rs
