roc/crates/compiler/serialize/src/bytes.rs
2022-10-12 16:42:11 -05:00

385 lines
10 KiB
Rust

use std::{
borrow::Borrow,
io::{self, Write},
};
use roc_collections::{MutMap, VecMap};
pub fn serialize_slice<T: Copy>(
slice: &[T],
writer: &mut impl Write,
written: usize,
) -> io::Result<usize> {
let alignment = std::mem::align_of::<T>();
let padding_bytes = next_multiple_of(written, alignment) - written;
for _ in 0..padding_bytes {
writer.write_all(&[0])?;
}
let bytes_slice = unsafe { slice_as_bytes(slice) };
writer.write_all(bytes_slice)?;
Ok(written + padding_bytes + bytes_slice.len())
}
pub fn deserialize_slice<T: Copy>(bytes: &[u8], length: usize, mut offset: usize) -> (&[T], usize) {
let alignment = std::mem::align_of::<T>();
let size = std::mem::size_of::<T>();
offset = next_multiple_of(offset, alignment);
let byte_length = length * size;
let byte_slice = &bytes[offset..][..byte_length];
let slice = unsafe { std::slice::from_raw_parts(byte_slice.as_ptr() as *const T, length) };
(slice, offset + byte_length)
}
pub fn deserialize_vec<T: Clone + Copy>(
bytes: &[u8],
length: usize,
offset: usize,
) -> (Vec<T>, usize) {
let (slice, offset) = deserialize_slice(bytes, length, offset);
(slice.to_vec(), offset)
}
#[derive(Copy, Clone)]
struct VecSlice<T> {
pub start: u32,
pub length: u16,
_marker: std::marker::PhantomData<T>,
}
impl<T> VecSlice<T> {
const fn len(&self) -> usize {
self.length as usize
}
fn indices(&self) -> std::ops::Range<usize> {
self.start as usize..(self.start as usize + self.length as usize)
}
fn extend_new(vec: &mut Vec<T>, it: impl IntoIterator<Item = T>) -> Self {
let start = vec.len();
vec.extend(it);
let end = vec.len();
Self {
start: start as u32,
length: (end - start) as u16,
_marker: Default::default(),
}
}
}
pub fn serialize_slice_of_slices<'a, T, U>(
slice_of_slices: &[U],
writer: &mut impl Write,
written: usize,
) -> io::Result<usize>
where
T: 'a + Copy,
U: 'a + Borrow<[T]> + Sized,
{
let mut item_buf: Vec<T> = Vec::new();
let mut serialized_slices: Vec<VecSlice<T>> = Vec::new();
for slice in slice_of_slices {
let slice = VecSlice::extend_new(&mut item_buf, slice.borrow().iter().copied());
serialized_slices.push(slice);
}
let written = serialize_slice(&serialized_slices, writer, written)?;
serialize_slice(&item_buf, writer, written)
}
pub fn deserialize_slice_of_slices<T, Container>(
bytes: &[u8],
length: usize,
offset: usize,
) -> (Vec<Container>, usize)
where
T: Copy,
Container: From<Vec<T>>,
{
let (serialized_slices, offset) = deserialize_slice::<VecSlice<T>>(bytes, length, offset);
let (vars_slice, offset) = {
let total_items = serialized_slices.iter().map(|s| s.len()).sum();
deserialize_slice::<T>(bytes, total_items, offset)
};
let mut slice_of_slices = Vec::with_capacity(length);
for slice in serialized_slices {
let deserialized_slice = &vars_slice[slice.indices()];
slice_of_slices.push(deserialized_slice.to_vec().into())
}
(slice_of_slices, offset)
}
pub fn serialize_map<K: Clone, V: Clone, W: Write>(
map: &MutMap<K, V>,
ser_keys: fn(&[K], &mut W, usize) -> io::Result<usize>,
ser_values: fn(&[V], &mut W, usize) -> io::Result<usize>,
writer: &mut W,
written: usize,
) -> io::Result<usize> {
let keys = map.keys().cloned().collect::<Vec<_>>();
let values = map.values().cloned().collect::<Vec<_>>();
let written = ser_keys(keys.as_slice(), writer, written)?;
let written = ser_values(values.as_slice(), writer, written)?;
Ok(written)
}
#[allow(clippy::type_complexity)]
pub fn deserialize_map<K, V>(
bytes: &[u8],
deser_keys: fn(&[u8], usize, usize) -> (Vec<K>, usize),
deser_values: fn(&[u8], usize, usize) -> (Vec<V>, usize),
length: usize,
offset: usize,
) -> (MutMap<K, V>, usize)
where
K: Clone + std::hash::Hash + Eq,
V: Clone,
{
let (keys, offset) = deser_keys(bytes, length, offset);
let (values, offset) = deser_values(bytes, length, offset);
(
MutMap::from_iter((keys.iter().cloned()).zip(values.iter().cloned())),
offset,
)
}
pub fn serialize_vec_map<K, V, W: Write>(
map: &VecMap<K, V>,
ser_keys: fn(&[K], &mut W, usize) -> io::Result<usize>,
ser_values: fn(&[V], &mut W, usize) -> io::Result<usize>,
writer: &mut W,
written: usize,
) -> io::Result<usize> {
let (keys, values) = map.unzip_slices();
let written = ser_keys(keys, writer, written)?;
let written = ser_values(values, writer, written)?;
Ok(written)
}
#[allow(clippy::type_complexity)]
pub fn deserialize_vec_map<K, V>(
bytes: &[u8],
deser_keys: fn(&[u8], usize, usize) -> (Vec<K>, usize),
deser_values: fn(&[u8], usize, usize) -> (Vec<V>, usize),
length: usize,
offset: usize,
) -> (VecMap<K, V>, usize)
where
K: PartialEq,
{
let (keys, offset) = deser_keys(bytes, length, offset);
let (values, offset) = deser_values(bytes, length, offset);
(unsafe { VecMap::zip(keys, values) }, offset)
}
unsafe fn slice_as_bytes<T>(slice: &[T]) -> &[u8] {
let ptr = slice.as_ptr();
let byte_length = std::mem::size_of::<T>() * slice.len();
std::slice::from_raw_parts(ptr as *const u8, byte_length)
}
// TODO check on https://github.com/rust-lang/rust/issues/88581 some time in the future
pub const fn next_multiple_of(lhs: usize, rhs: usize) -> usize {
match lhs % rhs {
0 => lhs,
r => lhs + (rhs - r),
}
}
#[cfg(test)]
mod test {
use roc_collections::{MutMap, VecMap, VecSet};
use super::{
deserialize_map, deserialize_slice, deserialize_slice_of_slices, deserialize_vec,
deserialize_vec_map, serialize_map, serialize_slice, serialize_slice_of_slices,
serialize_vec_map,
};
#[test]
fn serde_empty_slice() {
let mut buf = vec![];
serialize_slice(&[] as &[u8], &mut buf, 0).unwrap();
assert!(buf.is_empty());
let (out, size) = deserialize_slice::<u8>(&buf, 0, 0);
assert!(out.is_empty());
assert_eq!(size, 0);
}
#[test]
fn serde_slice() {
let input: &[u64] = &[15u64, 23, 37, 89];
let mut buf = vec![];
serialize_slice(input, &mut buf, 0).unwrap();
assert!(!buf.is_empty());
let (out, size) = deserialize_slice::<u64>(&buf, 4, 0);
assert_eq!(out, input);
assert_eq!(size, 4 * 8);
}
#[test]
fn serde_vec() {
let input: &[u64] = &[15u64, 23, 37, 89];
let mut buf = vec![];
serialize_slice(input, &mut buf, 0).unwrap();
assert!(!buf.is_empty());
let (out, size) = deserialize_vec::<u64>(&buf, 4, 0);
assert_eq!(out, input);
assert_eq!(size, buf.len());
}
#[test]
fn serde_empty_slice_of_slices() {
let input: &[&[u64]] = &[];
let mut buf = vec![];
serialize_slice_of_slices(input, &mut buf, 0).unwrap();
assert!(buf.is_empty());
let (out, size) = deserialize_slice_of_slices::<u64, Vec<_>>(&buf, 0, 0);
assert!(out.is_empty());
assert_eq!(size, 0);
}
#[test]
fn serde_slice_of_slices() {
let input: &[&[u64]] = &[&[15, 23, 47], &[61, 72], &[85, 91]];
let mut buf = vec![];
serialize_slice_of_slices(input, &mut buf, 0).unwrap();
assert!(!buf.is_empty());
let (out, size) = deserialize_slice_of_slices::<u64, Vec<_>>(&buf, 3, 0);
assert_eq!(out, input);
assert_eq!(size, buf.len());
}
#[test]
fn serde_slice_of_slices_into_vec_set() {
let input: &[&[u64]] = &[&[15, 23, 47], &[61, 72], &[85, 91]];
let mut buf = vec![];
serialize_slice_of_slices(input, &mut buf, 0).unwrap();
assert!(!buf.is_empty());
let (out, size) = deserialize_slice_of_slices::<u64, VecSet<_>>(&buf, 3, 0);
assert_eq!(size, buf.len());
let mut out = out.into_iter();
assert_eq!(out.next().unwrap().into_vec(), &[15, 23, 47]);
assert_eq!(out.next().unwrap().into_vec(), &[61, 72]);
assert_eq!(out.next().unwrap().into_vec(), &[85, 91]);
assert!(out.next().is_none());
}
#[test]
fn serde_empty_map() {
let input: MutMap<u64, u64> = Default::default();
let mut buf = vec![];
serialize_map(&input, serialize_slice, serialize_slice, &mut buf, 0).unwrap();
assert!(buf.is_empty());
let (out, size) = deserialize_map::<u64, u64>(&buf, deserialize_vec, deserialize_vec, 0, 0);
assert!(out.is_empty());
assert_eq!(size, 0);
}
#[test]
fn serde_map() {
let mut input: MutMap<u64, Vec<u64>> = Default::default();
input.insert(51, vec![15, 23, 37]);
input.insert(39, vec![17, 91, 43]);
input.insert(82, vec![90, 35, 76]);
let mut buf = vec![];
serialize_map(
&input,
serialize_slice,
serialize_slice_of_slices,
&mut buf,
0,
)
.unwrap();
assert!(!buf.is_empty());
let (out, size) = deserialize_map::<u64, Vec<u64>>(
&buf,
deserialize_vec,
deserialize_slice_of_slices,
3,
0,
);
assert_eq!(out, input);
assert_eq!(size, buf.len());
}
#[test]
fn serde_empty_vec_map() {
let input: VecMap<u64, u64> = Default::default();
let mut buf = vec![];
serialize_vec_map(&input, serialize_slice, serialize_slice, &mut buf, 0).unwrap();
assert!(buf.is_empty());
let (out, size) =
deserialize_vec_map::<u64, u64>(&buf, deserialize_vec, deserialize_vec, 0, 0);
assert!(out.is_empty());
assert_eq!(size, 0);
}
#[test]
fn serde_vec_map() {
let mut input: VecMap<u64, Vec<u64>> = Default::default();
input.insert(51, vec![15, 23, 37]);
input.insert(39, vec![17, 91, 43]);
input.insert(82, vec![90, 35, 76]);
let mut buf = vec![];
serialize_vec_map(
&input,
serialize_slice,
serialize_slice_of_slices,
&mut buf,
0,
)
.unwrap();
assert!(!buf.is_empty());
let (out, size) = deserialize_vec_map::<u64, Vec<u64>>(
&buf,
deserialize_vec,
deserialize_slice_of_slices,
3,
0,
);
assert_eq!(out.unzip_slices(), input.unzip_slices());
assert_eq!(size, buf.len());
}
}