use crate::serialization::{Frozen, Native, Portable};
use crate::{bitmap, Treemap};
use crate::{Bitmap, JvmLegacy};
use std::collections::BTreeMap;
use std::io;
use std::io::Write as _;
use byteorder::{BigEndian, NativeEndian, ReadBytesExt, WriteBytesExt};
use std::mem::size_of;
pub trait Serializer {
fn serialize_into_writer<W>(treemap: &Treemap, dst: W) -> io::Result<usize>
where
W: io::Write;
fn serialize_into<'a>(treemap: &Treemap, dst: &'a mut Vec<u8>) -> &'a [u8];
fn get_serialized_size_in_bytes(treemap: &Treemap) -> usize;
}
pub trait Deserializer {
fn try_deserialize(buffer: &[u8]) -> Option<(Treemap, usize)>;
}
fn serialize_impl<'a, S>(treemap: &Treemap, dst: &'a mut Vec<u8>) -> &'a [u8]
where
S: bitmap::Serializer,
{
let start_idx = dst.len();
let map_len = u64::try_from(treemap.map.len()).unwrap();
dst.extend_from_slice(&map_len.to_ne_bytes());
treemap.map.iter().for_each(|(&key, bitmap)| {
dst.extend_from_slice(&key.to_ne_bytes());
let prev_len = dst.len();
let serialized_slice = bitmap.serialize_into::<S>(dst);
let serialized_len = serialized_slice.len();
let serialized_range = serialized_slice.as_ptr_range();
debug_assert_eq!(prev_len + serialized_len, dst.len());
debug_assert_eq!(serialized_range.end, dst.as_ptr_range().end);
});
&dst[start_idx..]
}
fn serialize_writer_impl<S, W>(treemap: &Treemap, dst: W) -> io::Result<usize>
where
S: bitmap::Serializer,
W: io::Write,
{
let mut dst = OffsetTrackingWriter::new(dst);
let map_len = u64::try_from(treemap.map.len()).unwrap();
dst.write_u64::<NativeEndian>(map_len)?;
let mut buf = Vec::new();
for (&key, bitmap) in &treemap.map {
dst.write_u32::<NativeEndian>(key)?;
let bitmap_serialized = bitmap.serialize_into::<S>(&mut buf);
dst.write_all(bitmap_serialized)?;
buf.clear();
}
Ok(dst.bytes_written)
}
fn size_in_bytes_impl<S>(treemap: &Treemap) -> usize
where
S: bitmap::Serializer,
{
let overhead = size_of::<u64>() + treemap.map.len() * size_of::<u32>();
let total_sizes = treemap
.map
.values()
.map(Bitmap::get_serialized_size_in_bytes::<S>)
.sum::<usize>();
overhead + total_sizes
}
fn deserialize_impl<S>(mut buffer: &[u8]) -> Option<(Treemap, usize)>
where
S: bitmap::Serializer + bitmap::Deserializer,
{
let start_len = buffer.len();
let map_len = buffer.read_u64::<NativeEndian>().ok()?;
let mut map = BTreeMap::new();
for _ in 0..map_len {
let key = buffer.read_u32::<NativeEndian>().ok()?;
let bitmap = Bitmap::try_deserialize::<S>(buffer)?;
buffer = &buffer[bitmap.get_serialized_size_in_bytes::<S>()..];
map.insert(key, bitmap);
}
Some((Treemap { map }, start_len - buffer.len()))
}
impl Serializer for Portable {
fn serialize_into_writer<W>(treemap: &Treemap, dst: W) -> io::Result<usize>
where
W: io::Write,
{
serialize_writer_impl::<Self, W>(treemap, dst)
}
fn serialize_into<'a>(treemap: &Treemap, dst: &'a mut Vec<u8>) -> &'a [u8] {
serialize_impl::<Self>(treemap, dst)
}
fn get_serialized_size_in_bytes(treemap: &Treemap) -> usize {
size_in_bytes_impl::<Self>(treemap)
}
}
impl Deserializer for Portable {
fn try_deserialize(buffer: &[u8]) -> Option<(Treemap, usize)> {
deserialize_impl::<Self>(buffer)
}
}
impl Serializer for Native {
fn serialize_into_writer<W>(treemap: &Treemap, dst: W) -> io::Result<usize>
where
W: io::Write,
{
serialize_writer_impl::<Self, W>(treemap, dst)
}
fn serialize_into<'a>(treemap: &Treemap, dst: &'a mut Vec<u8>) -> &'a [u8] {
serialize_impl::<Self>(treemap, dst)
}
fn get_serialized_size_in_bytes(treemap: &Treemap) -> usize {
size_in_bytes_impl::<Self>(treemap)
}
}
impl Deserializer for Native {
fn try_deserialize(buffer: &[u8]) -> Option<(Treemap, usize)> {
deserialize_impl::<Self>(buffer)
}
}
const FROZEN_BITMAP_METADATA_SIZE: usize = size_of::<usize>() + size_of::<u32>();
impl Serializer for Frozen {
fn serialize_into_writer<W>(treemap: &Treemap, dst: W) -> io::Result<usize>
where
W: io::Write,
{
const FULL_PADDING: [u8; Frozen::MAX_PADDING] = [0; Frozen::MAX_PADDING];
let mut dst = OffsetTrackingWriter::new(dst);
let map_size = u64::try_from(treemap.map.len()).unwrap();
dst.write_all(&u64::to_ne_bytes(map_size))?;
let mut buf = Vec::new();
for (&key, bitmap) in &treemap.map {
let bitmap_serialized = bitmap.serialize_into::<Self>(&mut buf);
let required_padding =
Self::required_padding(dst.bytes_written + FROZEN_BITMAP_METADATA_SIZE);
dst.write_all(&FULL_PADDING[..required_padding])?;
dst.write_all(&usize::to_ne_bytes(bitmap_serialized.len()))?;
dst.write_all(&u32::to_ne_bytes(key))?;
debug_assert_eq!(dst.bytes_written % Self::REQUIRED_ALIGNMENT, 0);
dst.write_all(bitmap_serialized)?;
buf.clear();
}
Ok(dst.bytes_written)
}
fn serialize_into<'a>(treemap: &Treemap, dst: &'a mut Vec<u8>) -> &'a [u8] {
let len = Self::get_serialized_size_in_bytes(treemap);
let mut offset = dst.len();
if dst.capacity() < dst.len() + len
|| Self::required_padding(dst.as_ptr() as usize + offset) != 0
{
dst.reserve(len.checked_add(Self::MAX_PADDING).unwrap());
let extra_offset = Self::required_padding(dst.as_ptr() as usize + offset);
offset = offset.checked_add(extra_offset).unwrap();
dst.resize(offset, 0);
}
let total_len = offset.checked_add(len).unwrap();
debug_assert!(dst.capacity() >= total_len);
let map_size = u64::try_from(treemap.map.len()).unwrap();
dst.extend_from_slice(&map_size.to_ne_bytes());
treemap.map.iter().for_each(|(&key, bitmap)| {
let end_with_metadata = dst.as_ptr_range().end as usize + FROZEN_BITMAP_METADATA_SIZE;
let extra_padding = Self::required_padding(end_with_metadata);
dst.resize(dst.len() + extra_padding, 0);
let frozen_size_in_bytes: usize = bitmap.get_serialized_size_in_bytes::<Self>();
dst.extend_from_slice(&frozen_size_in_bytes.to_ne_bytes());
dst.extend_from_slice(&key.to_ne_bytes());
let before_bitmap_serialize = dst.as_ptr_range().end;
let serialized_slice = bitmap.serialize_into::<Self>(dst);
debug_assert_eq!(before_bitmap_serialize, serialized_slice.as_ptr());
debug_assert_eq!(serialized_slice.as_ptr_range().end, dst.as_ptr_range().end);
});
&dst[offset..]
}
fn get_serialized_size_in_bytes(treemap: &Treemap) -> usize {
let mut result = size_of::<u64>();
for bitmap in treemap.map.values() {
result += FROZEN_BITMAP_METADATA_SIZE;
result += Self::required_padding(result);
result += bitmap.get_serialized_size_in_bytes::<Self>();
}
result
}
}
impl Serializer for JvmLegacy {
fn serialize_into_writer<W>(treemap: &Treemap, dst: W) -> io::Result<usize>
where
W: io::Write,
{
let mut dst = OffsetTrackingWriter::new(dst);
dst.write_u8(0)?;
let bitmap_count: u32 = treemap.map.len().try_into().unwrap();
dst.write_u32::<BigEndian>(bitmap_count)?;
let mut buf = Vec::new();
for (&key, bitmap) in &treemap.map {
dst.write_u32::<BigEndian>(key)?;
let bitmap_serialized = bitmap.serialize_into::<Portable>(&mut buf);
dst.write_all(bitmap_serialized)?;
buf.clear();
}
Ok(dst.bytes_written)
}
fn serialize_into<'a>(treemap: &Treemap, dst: &'a mut Vec<u8>) -> &'a [u8] {
let start_idx = dst.len();
dst.write_u8(0).unwrap();
let bitmap_count: u32 = treemap.map.len().try_into().unwrap();
dst.write_u32::<BigEndian>(bitmap_count).unwrap();
treemap.map.iter().for_each(|(&key, bitmap)| {
dst.write_u32::<BigEndian>(key).unwrap();
bitmap.serialize_into::<Portable>(dst);
});
&dst[start_idx..]
}
fn get_serialized_size_in_bytes(treemap: &Treemap) -> usize {
let overhead = size_of::<u8>() + size_of::<u32>() + size_of::<u32>() * treemap.map.len();
let total_sizes = treemap
.map
.values()
.map(Bitmap::get_serialized_size_in_bytes::<Portable>)
.sum::<usize>();
overhead + total_sizes
}
}
impl Deserializer for JvmLegacy {
fn try_deserialize(mut buffer: &[u8]) -> Option<(Treemap, usize)> {
let start_len = buffer.len();
let _is_signed = buffer.read_u8().ok()?;
let bitmap_count = buffer.read_u32::<BigEndian>().ok()?;
let mut map = BTreeMap::new();
for _ in 0..bitmap_count {
let key = buffer.read_u32::<BigEndian>().ok()?;
let bitmap = Bitmap::try_deserialize::<Portable>(buffer)?;
buffer = &buffer[bitmap.get_serialized_size_in_bytes::<Portable>()..];
map.insert(key, bitmap);
}
Some((Treemap { map }, start_len - buffer.len()))
}
}
struct OffsetTrackingWriter<W> {
writer: W,
bytes_written: usize,
}
impl<W> OffsetTrackingWriter<W> {
pub fn new(writer: W) -> Self {
Self {
writer,
bytes_written: 0,
}
}
}
impl<W: io::Write> io::Write for OffsetTrackingWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let written = self.writer.write(buf)?;
self.bytes_written += written;
Ok(written)
}
fn flush(&mut self) -> io::Result<()> {
self.writer.flush()
}
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
self.writer.write_all(buf)?;
self.bytes_written += buf.len();
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn example_treemap() -> Treemap {
Treemap::from_iter([1, 2, 3, 4, 5, u64::from(u32::MAX), u64::MAX])
}
fn smoke_test_ser<S: Serializer>(expected_len: usize) {
let treemap = example_treemap();
assert_eq!(treemap.get_serialized_size_in_bytes::<S>(), expected_len);
let mut buf = Vec::new();
let serialized = treemap.serialize_into::<S>(&mut buf);
assert_eq!(serialized.len(), expected_len);
let mut writer = Vec::new();
assert_eq!(
treemap.serialize_into_writer::<S, _>(&mut writer).unwrap(),
expected_len,
);
assert_eq!(serialized, writer);
}
#[test]
fn smoke_portable() {
smoke_test_ser::<Portable>(70);
}
#[test]
fn smoke_native() {
smoke_test_ser::<Native>(54);
}
#[test]
fn smoke_frozen() {
smoke_test_ser::<Frozen>(107);
}
#[test]
fn smoke_jvm_legacy() {
smoke_test_ser::<JvmLegacy>(67);
}
}