Merge pull request #81 from Dr-Emann/multi_iter

Expose the ability to read multiple items from an iterator, for performance

Author: saulius

croaring/src/bitmap/iter.rs

@@ -3,11 +3,15 @@ use std::marker::PhantomData;
 
 use super::Bitmap;
 
+#[derive(Clone)]
 pub struct BitmapIterator<'a> {
     iterator: ffi::roaring_uint32_iterator_s,
     phantom: PhantomData<&'a Bitmap>,
 }
 
+unsafe impl Send for BitmapIterator<'_> {}
+unsafe impl Sync for BitmapIterator<'_> {}
+
 impl<'a> BitmapIterator<'a> {
     fn new(bitmap: &'a Bitmap) -> Self {
         let mut iterator = std::mem::MaybeUninit::uninit();
@@ -38,6 +42,64 @@ impl<'a> BitmapIterator<'a> {
     fn advance(&mut self) -> bool {
         unsafe { ffi::roaring_advance_uint32_iterator(&mut self.iterator) }
     }
+
+    /// Attempt to read many values from the iterator into `dst`
+    ///
+    /// Returns the number of items read from the iterator, may be `< dst.len()` iff
+    /// the iterator is exhausted.
+    ///
+    /// This can be much more efficient than repeated iteration.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use croaring::Bitmap;
+    ///
+    /// let mut bitmap: Bitmap = Bitmap::create();
+    /// bitmap.add_range(0..100);
+    /// bitmap.add(222);
+    /// bitmap.add(555);
+    ///
+    /// let mut buf = [0; 100];
+    /// let mut iter = bitmap.iter();
+    /// assert_eq!(iter.next_many(&mut buf), 100);
+    /// // Get the first 100 items, from the original range added
+    /// for (i, item) in buf.iter().enumerate() {
+    ///     assert_eq!(*item, i as u32);
+    /// }
+    /// // Calls to next_many() can be interleaved with calls to next()
+    /// assert_eq!(iter.next(), Some(222));
+    /// assert_eq!(iter.next_many(&mut buf), 1);
+    /// assert_eq!(buf[0], 555);
+    ///
+    /// assert_eq!(iter.next(), None);
+    /// assert_eq!(iter.next_many(&mut buf), 0);
+    /// ```
+    ///
+    /// ```
+    /// use croaring::Bitmap;
+    ///
+    /// fn print_by_chunks(bitmap: &Bitmap) {
+    ///     let mut buf = [0; 1024];
+    ///     let mut iter = bitmap.iter();
+    ///     loop {
+    ///         let n = iter.next_many(&mut buf);
+    ///         if n == 0 {
+    ///             break;
+    ///         }
+    ///         println!("{:?}", &buf[..n]);
+    ///     }
+    /// }
+    ///
+    /// # print_by_chunks(&Bitmap::of(&[1, 2, 8, 20, 1000]));
+    /// ```
+    #[inline]
+    pub fn next_many(&mut self, dst: &mut [u32]) -> usize {
+        let count: u32 = u32::try_from(dst.len()).unwrap_or(u32::MAX);
+        let result = unsafe { ffi::roaring_read_uint32_iterator(&mut self.iterator, dst.as_mut_ptr(), count)};
+        debug_assert!(result <= count);
+        result as usize
+    }
 }
 
 impl<'a> Iterator for BitmapIterator<'a> {

croaring/src/treemap/imp.rs

@@ -564,13 +564,21 @@ impl Treemap {
     pub fn to_vec(&self) -> Vec<u64> {
         let treemap_size: usize = self.cardinality().try_into().unwrap();
 
-        let mut buffer: Vec<u64> = Vec::with_capacity(treemap_size);
-
-        for (key, bitmap) in &self.map {
-            bitmap.iter().for_each(|bit| buffer.push(util::join(*key, bit)));
+        let mut result: Vec<u64> = Vec::with_capacity(treemap_size);
+        let mut buffer = [0; 1024];
+
+        for (&key, bitmap) in &self.map {
+            let mut iter = bitmap.iter();
+            loop {
+                let n = iter.next_many(&mut buffer);
+                if n == 0 {
+                    break;
+                }
+                result.extend(buffer[..n].iter().map(|&bit| util::join(key, bit)))
+            }
         }
 
-        buffer
+        result
     }
 
     /// Creates a new treemap from a slice of u64 integers