use bytemuck::cast_slice_mut;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::convert::{Infallible, TryFrom};
use std::error::Error;
use std::io;
use super::container::Container;
use crate::bitmap::store::{ArrayStore, BitmapStore, Store};
use crate::RoaringBitmap;
const SERIAL_COOKIE_NO_RUNCONTAINER: u32 = 12346;
const SERIAL_COOKIE: u16 = 12347;
impl RoaringBitmap {
pub fn serialized_size(&self) -> usize {
let container_sizes: usize = self
.containers
.iter()
.map(|container| match container.store {
Store::Array(ref values) => 8 + values.len() as usize * 2,
Store::Bitmap(..) => 8 + 8 * 1024,
})
.sum();
8 + container_sizes
}
pub fn serialize_into<W: io::Write>(&self, mut writer: W) -> io::Result<()> {
writer.write_u32::<LittleEndian>(SERIAL_COOKIE_NO_RUNCONTAINER)?;
writer.write_u32::<LittleEndian>(self.containers.len() as u32)?;
for container in &self.containers {
writer.write_u16::<LittleEndian>(container.key)?;
writer.write_u16::<LittleEndian>((container.len() - 1) as u16)?;
}
let mut offset = 8 + 8 * self.containers.len() as u32;
for container in &self.containers {
writer.write_u32::<LittleEndian>(offset)?;
match container.store {
Store::Array(ref values) => {
offset += values.len() as u32 * 2;
}
Store::Bitmap(..) => {
offset += 8 * 1024;
}
}
}
for container in &self.containers {
match container.store {
Store::Array(ref values) => {
for &value in values.iter() {
writer.write_u16::<LittleEndian>(value)?;
}
}
Store::Bitmap(ref bits) => {
for &value in bits.as_array() {
writer.write_u64::<LittleEndian>(value)?;
}
}
}
}
Ok(())
}
pub fn deserialize_from<R: io::Read>(reader: R) -> io::Result<RoaringBitmap> {
RoaringBitmap::deserialize_from_impl(reader, ArrayStore::try_from, BitmapStore::try_from)
}
pub fn deserialize_unchecked_from<R: io::Read>(reader: R) -> io::Result<RoaringBitmap> {
RoaringBitmap::deserialize_from_impl::<R, _, Infallible, _, Infallible>(
reader,
|values| Ok(ArrayStore::from_vec_unchecked(values)),
|len, values| Ok(BitmapStore::from_unchecked(len, values)),
)
}
fn deserialize_from_impl<R, A, AErr, B, BErr>(
mut reader: R,
a: A,
b: B,
) -> io::Result<RoaringBitmap>
where
R: io::Read,
A: Fn(Vec<u16>) -> Result<ArrayStore, AErr>,
AErr: Error + Send + Sync + 'static,
B: Fn(u64, Box<[u64; 1024]>) -> Result<BitmapStore, BErr>,
BErr: Error + Send + Sync + 'static,
{
let (size, has_offsets) = {
let cookie = reader.read_u32::<LittleEndian>()?;
if cookie == SERIAL_COOKIE_NO_RUNCONTAINER {
(reader.read_u32::<LittleEndian>()? as usize, true)
} else if (cookie as u16) == SERIAL_COOKIE {
return Err(io::Error::new(io::ErrorKind::Other, "run containers are unsupported"));
} else {
return Err(io::Error::new(io::ErrorKind::Other, "unknown cookie value"));
}
};
if size > u16::MAX as usize + 1 {
return Err(io::Error::new(io::ErrorKind::Other, "size is greater than supported"));
}
let mut description_bytes = vec![0u8; size * 4];
reader.read_exact(&mut description_bytes)?;
let mut description_bytes = &description_bytes[..];
if has_offsets {
let mut offsets = vec![0u8; size * 4];
reader.read_exact(&mut offsets)?;
drop(offsets); }
let mut containers = Vec::with_capacity(size);
for _ in 0..size {
let key = description_bytes.read_u16::<LittleEndian>()?;
let len = u64::from(description_bytes.read_u16::<LittleEndian>()?) + 1;
let store = if len <= 4096 {
let mut values = vec![0; len as usize];
reader.read_exact(cast_slice_mut(&mut values))?;
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; 1024]);
reader.read_exact(cast_slice_mut(&mut values[..]))?;
values.iter_mut().for_each(|n| *n = u64::from_le(*n));
let bitmap =
b(len, values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
Store::Bitmap(bitmap)
};
containers.push(Container { key, store });
}
Ok(RoaringBitmap { containers })
}
}
#[cfg(test)]
mod test {
use crate::RoaringBitmap;
use proptest::prelude::*;
proptest! {
#[test]
fn test_serialization(
bitmap in RoaringBitmap::arbitrary(),
) {
let mut buffer = Vec::new();
bitmap.serialize_into(&mut buffer).unwrap();
prop_assert_eq!(bitmap, RoaringBitmap::deserialize_from(buffer.as_slice()).unwrap());
}
}
}