Prefer reading the content store by store

Author: Kerollmops

src/bitmap/ops_with_serialized.rs

@@ -51,7 +51,7 @@ impl RoaringBitmap {
 
     fn intersection_with_serialized_impl<R, A, AErr, B, BErr>(
         &self,
-        mut other: R,
+        mut reader: R,
         a: A,
         b: B,
     ) -> io::Result<RoaringBitmap>
@@ -64,9 +64,9 @@ impl RoaringBitmap {
     {
         // First read the cookie to determine which version of the format we are reading
         let (size, has_offsets, has_run_containers) = {
-            let cookie = other.read_u32::<LittleEndian>()?;
+            let cookie = reader.read_u32::<LittleEndian>()?;
             if cookie == SERIAL_COOKIE_NO_RUNCONTAINER {
-                (other.read_u32::<LittleEndian>()? as usize, true, false)
+                (reader.read_u32::<LittleEndian>()? as usize, true, false)
             } else if (cookie as u16) == SERIAL_COOKIE {
                 let size = ((cookie >> 16) + 1) as usize;
                 (size, size >= NO_OFFSET_THRESHOLD, true)
@@ -78,7 +78,7 @@ impl RoaringBitmap {
         // Read the run container bitmap if necessary
         let run_container_bitmap = if has_run_containers {
             let mut bitmap = vec![0u8; (size + 7) / 8];
-            other.read_exact(&mut bitmap)?;
+            reader.read_exact(&mut bitmap)?;
             Some(bitmap)
         } else {
             None
@@ -90,54 +90,75 @@ impl RoaringBitmap {
 
         // Read the container descriptions
         let mut description_bytes = vec![0u8; size * DESCRIPTION_BYTES];
-        other.read_exact(&mut description_bytes)?;
-        let mut description_bytes: Vec<[u16; 2]> = pod_collect_to_vec(&description_bytes);
-        description_bytes.iter_mut().for_each(|[ref mut k, ref mut c]| {
-            *k = u16::from_le(*k);
-            *c = u16::from_le(*c);
-        });
+        reader.read_exact(&mut description_bytes)?;
+        let mut description_bytes = &description_bytes[..];
 
         if has_offsets {
             let mut offsets = vec![0u8; size * OFFSET_BYTES];
-            other.read_exact(&mut offsets)?;
-            drop(offsets); // Not useful when deserializing into memory
+            reader.read_exact(&mut offsets)?;
+            drop(offsets); // We could use these offsets but we are lazy
         }
 
         let mut containers = Vec::with_capacity(size);
-        for container in &self.containers {
-            let (i, key, cardinality) =
-                match description_bytes.binary_search_by_key(&container.key, |[k, _]| *k) {
-                    Ok(index) => {
-                        let [key, cardinality] = description_bytes[index];
-                        (index, key, u64::from(cardinality) + 1)
-                    }
-                    Err(_) => continue,
-                };
+
+        // Read each container
+        for i in 0..size {
+            let key = description_bytes.read_u16::<LittleEndian>()?;
+            let container = match self.containers.binary_search_by_key(&key, |c| c.key) {
+                Ok(index) => self.containers.get(index),
+                Err(_) => None,
+            };
+            let cardinality = u64::from(description_bytes.read_u16::<LittleEndian>()?) + 1;
 
             // If the run container bitmap is present, check if this container is a run container
             let is_run_container =
                 run_container_bitmap.as_ref().map_or(false, |bm| bm[i / 8] & (1 << (i % 8)) != 0);
 
             let mut store = if is_run_container {
-                todo!("support run containers")
+                let runs = reader.read_u16::<LittleEndian>()?;
+                let mut intervals = vec![[0, 0]; runs as usize];
+                reader.read_exact(cast_slice_mut(&mut intervals))?;
+                if container.is_none() {
+                    continue;
+                }
+                intervals.iter_mut().for_each(|[s, len]| {
+                    *s = u16::from_le(*s);
+                    *len = u16::from_le(*len);
+                });
+
+                let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum();
+                let mut store = Store::with_capacity(cardinality);
+                intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> {
+                    let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?;
+                    store.insert_range(RangeInclusive::new(s, end));
+                    Ok(())
+                })?;
+                store
             } else if cardinality <= ARRAY_LIMIT {
                 let mut values = vec![0; cardinality as usize];
-                other.read_exact(cast_slice_mut(&mut values))?;
+                reader.read_exact(cast_slice_mut(&mut values))?;
+                if container.is_none() {
+                    continue;
+                }
                 values.iter_mut().for_each(|n| *n = u16::from_le(*n));
                 let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                 Store::Array(array)
             } else {
                 let mut values = Box::new([0; BITMAP_LENGTH]);
-                other.read_exact(cast_slice_mut(&mut values[..]))?;
+                reader.read_exact(cast_slice_mut(&mut values[..]))?;
+                if container.is_none() {
+                    continue;
+                }
                 values.iter_mut().for_each(|n| *n = u64::from_le(*n));
                 let bitmap = b(cardinality, values)
                     .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                 Store::Bitmap(bitmap)
             };
 
-            store &= &container.store;
-
-            containers.push(Container { key, store });
+            if let Some(container) = container {
+                store &= &container.store;
+                containers.push(Container { key, store });
+            }
         }
 
         Ok(RoaringBitmap { containers })