rexsgdata 0.12.0

//
// Copyright (c) 2021 RepliXio Ltd. All rights reserved.
// Use is subject to license terms.
//

#![allow(clippy::missing_safety_doc)]

use std::iter::once;
use std::iter::FromIterator;
use std::ops::AddAssign;
use std::ops::{Add, Range};

use bytes::Bytes;
use libc::iovec;
use serde::{Deserialize, Serialize};
use tinyvec::{tiny_vec, TinyVec};

use super::Iovec;
use super::SgList;

/// High Level scatter gather list wrapper for multiple data representation methods.
///
///# Safety: When creating SgData from Iovec / SgList, it is assumed that these
///#          strustures live longer than the created SgData lifetime.
///#          This is the user responsibility taking care of
///#          allocating/deallocating Iovec / SgList memory.
#[derive(Clone, Default, Serialize, Deserialize, Debug)]
pub struct SgData {
    data: TinyVec<[Bytes; 8]>,
}

impl SgData {
    /// Copies self into a Vector
    pub fn to_vec(&self) -> Vec<u8> {
        let mut vec = Vec::with_capacity(self.capacity());

        for slice in self.iter() {
            vec.extend_from_slice(slice)
        }
        vec
    }

    /// Moves the data contained in this `SgData` instance into `iovec`.
    /// Consumes this instance and creates a new one, based on the buffers
    /// described by `iovec` and `count`.
    #[must_use]
    pub unsafe fn drain_into(self, iovec: *mut iovec, count: usize) -> Self {
        let mut sglist = SgList::new(iovec, count);
        assert!(self.capacity() <= sglist.capacity());

        for (mut buf, range) in sglist.iter_slices_range_mut() {
            for slice in self.masked(range).iter() {
                buf[..slice.len()].copy_from_slice(slice);
                buf = &mut buf[slice.len()..]
            }
        }

        Self::from_sglist(sglist)
    }

    /// Applies the `mask` to the current object and creates a new `SgData` object which describes
    /// only masked region
    #[must_use]
    pub fn masked(&self, mask: Range<usize>) -> Self {
        let data = self
            .data
            .iter()
            .zip(self.iter_ranges())
            .filter_map(move |(bytes, range)| {
                let start = mask.start.checked_sub(range.start).unwrap_or_default();
                let end = range.end.min(mask.end).checked_sub(range.start)?;
                let range = start..end;
                if range.is_empty() {
                    None
                } else {
                    Some(bytes.slice(range))
                }
            })
            .collect();

        Self { data }
    }

    #[must_use]
    pub fn clone_masked(&self, mask: Range<usize>) -> Self {
        self.masked(mask)
    }

    pub unsafe fn from_sglist(sglist: SgList) -> Self {
        let data = sglist.iter_static_slices().map(Bytes::from).collect();

        Self { data }
    }

    /// Iterates over all the underlying byte slices of this `SgData`.
    pub fn iter(&self) -> impl Iterator<Item = &[u8]> {
        self.data.iter().map(AsRef::as_ref)
    }

    /// Iterate over `slices` `ranges`
    fn iter_ranges(&self) -> impl Iterator<Item = Range<usize>> + '_ {
        let mut offset = 0;
        self.iter().map(move |slice| {
            let start = offset;
            offset += slice.len();
            start..offset
        })
    }

    /// Get an iterator over `SgData`s bytes - useful for comparing values.
    fn iter_bytes(&self) -> impl Iterator<Item = &u8> {
        self.iter().flatten()
    }

    /// Returns number of elements in the underlying object, NOT the total number of bytes.
    pub fn size(&self) -> usize {
        self.data.len()
    }

    /// Returns the total number of bytes in the underlying object.
    pub fn capacity(&self) -> usize {
        self.data.iter().map(|item| item.len()).sum()
    }

    /// Checks this `SgData` for containing any byte slice
    pub fn is_empty(&self) -> bool {
        self.capacity() == 0
    }

    /// Create an `SgData` with specified capacity
    pub fn with_capacity(capacity: usize) -> Self {
        let ptr = std::mem::ManuallyDrop::new(Vec::<u8>::with_capacity(capacity)).as_mut_ptr();
        let vec = unsafe { Vec::from_raw_parts(ptr, capacity, capacity) };
        vec.into()
    }
}

impl From<Vec<u8>> for SgData {
    fn from(vec: Vec<u8>) -> Self {
        Self {
            data: tiny_vec![vec.into()],
        }
    }
}

impl From<Vec<Vec<u8>>> for SgData {
    fn from(sgvec: Vec<Vec<u8>>) -> Self {
        Self {
            data: sgvec.into_iter().map(Bytes::from).collect(),
        }
    }
}

impl From<Iovec> for SgData {
    fn from(iov: Iovec) -> Self {
        Self {
            data: tiny_vec![iov.as_static_slice().into()],
        }
    }
}

impl FromIterator<u8> for SgData {
    fn from_iter<I>(iter: I) -> Self
    where
        I: IntoIterator<Item = u8>,
    {
        Self {
            data: tiny_vec![iter.into_iter().collect()],
        }
    }
}

impl<I: Into<Self>> FromIterator<I> for SgData {
    fn from_iter<U>(iter: U) -> Self
    where
        U: IntoIterator<Item = I>,
    {
        let mut iter = iter.into_iter().map(Into::into);
        let mut data: Self = iter.next().unwrap_or_default();
        data.extend(iter);
        data
    }
}

impl Extend<u8> for SgData {
    fn extend<T>(&mut self, iter: T)
    where
        T: IntoIterator<Item = u8>,
    {
        self.data.push(Bytes::from_iter(iter))
    }
}

impl Extend<Self> for SgData {
    fn extend<T>(&mut self, iter: T)
    where
        T: IntoIterator<Item = Self>,
    {
        for item in iter {
            self.data.extend(item.data)
        }
    }
}

impl Add for SgData {
    type Output = Self;

    fn add(mut self, rhs: Self) -> Self::Output {
        self += rhs;
        self
    }
}

impl AddAssign for SgData {
    fn add_assign(&mut self, rhs: Self) {
        self.extend(once(rhs));
    }
}

impl slog::Value for SgData {
    fn serialize(
        &self,
        _record: &slog::Record<'_>,
        key: slog::Key,
        serializer: &mut dyn slog::Serializer,
    ) -> slog::Result {
        serializer.emit_str(key, "SgData<>")
    }
}

impl PartialEq for SgData {
    /// This method tests for `self` and `other` values to be equal, and is used
    /// by `==`.
    fn eq(&self, other: &Self) -> bool {
        self.capacity() == other.capacity()
            && self
                .iter_bytes()
                .zip(other.iter_bytes())
                .all(|(self_byte, other_byte)| self_byte == other_byte)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bincode::{deserialize, serialize};
    use libc::{c_void, iovec};
    use std::iter::once;
    use std::mem;

    fn _assert_impls() {
        fn assert_send<T: Send>() {}
        fn assert_sync<T: Sync>() {}
        fn assert_clone<T: Clone>() {}
        fn assert_copy<T: Copy>() {}

        assert_send::<SgData>();
        assert_sync::<SgData>();
    }

    // NB - never use this code outside of the tests - it leaks memory
    fn vec_into_iovec(mut vec: Vec<u8>) -> iovec {
        let len = vec.len();
        let base = vec.as_mut_ptr();
        mem::forget(vec);
        iovec {
            iov_base: base as *mut c_void,
            iov_len: len,
        }
    }

    // NB - never use this code outside of the tests - it leaks memory
    fn create_sglist(sgvec: Vec<Vec<u8>>) -> SgList {
        let vec = sgvec.into_iter().map(vec_into_iovec).collect::<Vec<_>>();
        let len = vec.len();
        let iov = vec.as_ptr();
        mem::forget(vec);
        SgList::new(iov, len)
    }

    //Drain data into sgvec - add compare values
    unsafe fn drain_data_compare(data_a: &SgData, sgvec: Vec<Vec<u8>>) {
        //Drain data into sgvec - add compare values
        let (base, count) = create_sglist(sgvec).into_inner();

        let drained_data = data_a.clone().drain_into(base as *mut libc::iovec, count);

        assert_eq!(data_a, &drained_data);
    }

    // asserts value-equality of to SgData instances.
    fn eq_data<T>(data_a: &SgData, data_b: &SgData, mask: T)
    where
        T: Into<Option<Range<usize>>>,
    {
        let mask = if let Some(mask) = mask.into() {
            mask
        } else {
            0..usize::MAX
        };

        let copy_a = data_a.masked(mask.clone());
        let copy_b = data_a.masked(mask.clone());
        assert_eq!(copy_a, copy_b);

        let ser_a = serialize(&copy_a).unwrap();
        let deser_a: SgData = deserialize(ser_a.as_slice()).unwrap();
        let ser_b = serialize(&copy_b).unwrap();
        let deser_b: SgData = deserialize(ser_b.as_slice()).unwrap();
        assert_eq!(deser_a, deser_b);

        unsafe {
            drain_data_compare(data_a, partition_vec(&vec![42; data_a.capacity()]));

            // create a vector of vector where one is of size 1
            drain_data_compare(data_a, vec![vec![42; data_a.capacity() - 1], vec![44; 1]]);

            // create a vector of vector where one is of size 1
            drain_data_compare(data_a, vec![vec![44; 1], vec![42; data_a.capacity() - 1]]);

            //Drain data into masked sgvec - add compare values
            drain_data_compare(
                &data_a.masked(mask.clone()),
                partition_vec(&vec![43; data_a.masked(mask.clone()).capacity()]),
            );

            assert_eq!(data_a.masked(mask.clone()), data_b.masked(mask.clone()));

            let tmp = data_b.masked(mask.clone());
            let data: SgData = tmp.iter().map(Iovec::from).collect();

            assert_eq!(data, data_a.masked(mask));
            drop(tmp)
        }
    }

    // given two value-equal mutable instances of SgData.
    // Test that they are (not)-eqaul after mutation.
    fn compare_data(mut data_a: SgData, mut data_b: SgData) {
        // Generate two sg-datas of different types with equal byte-wise values.
        eq_data(&data_a, &data_b, None);

        let capacity = data_a.capacity();

        // Define an arbitray data representation for sg-data
        let vec_def = gen_def_data(25, 25);

        let assign_a = build_data(&vec_def);
        let assign_b = build_data(&vec_def);

        data_a.add_assign(assign_a);

        assert_ne!(data_a, data_b);

        eq_data(&data_a, &data_b, 0..capacity);

        data_b.add_assign(assign_b);

        assert_eq!(data_a, data_b);

        let mut vec = data_a.to_vec();
        if !vec.is_empty() {
            vec[0] = if u8::MAX == vec[0] { 0 } else { vec[0] + 1 };

            let data_c = SgData::from(vec);
            assert_ne!(data_a, data_c);

            eq_data(&data_a, &data_c, 1..data_a.capacity());

            assert_ne!(data_b, data_c);
            eq_data(&data_b, &data_c, 1..data_b.capacity());
        }

        let sgvec = vec![vec![0_u8; data_a.capacity()]];
        let (base, count) = create_sglist(sgvec).into_inner();

        let drained_data = unsafe { data_a.clone().drain_into(base as *mut libc::iovec, count) };
        assert_eq!(data_a, drained_data);
    }

    // Given a vector defining a Sg-data - pairs of values X size, create a random SgData instance.
    fn build_data(data_def: &[u8]) -> SgData {
        use rand::{thread_rng, Rng};
        let mut rng = thread_rng();
        partition_vec(data_def)
            .into_iter()
            .map(|vec| match rng.gen_range(0..3) {
                0 => SgData::from(Iovec::from(vec_into_iovec(vec))),
                1 => SgData::from(vec),
                _ => unsafe { SgData::from_sglist(create_sglist(partition_vec(&vec))) },
            })
            .collect()
    }

    fn gen_def_data(num_of_entries: usize, max_entry_capacity: usize) -> Vec<u8> {
        use rand::{distributions::Uniform, random, thread_rng, Rng};
        (0..num_of_entries)
            .flat_map(|_| {
                let value = if thread_rng().gen_bool(0.2) {
                    0
                } else {
                    random::<u8>()
                };
                let size = thread_rng().sample::<usize, _>(Uniform::new(1, max_entry_capacity + 1));

                vec![value; size]
            })
            .collect()
    }

    fn partition_vec(vec: &[u8]) -> Vec<Vec<u8>> {
        use rand::{thread_rng, Rng};
        let indices: Vec<usize> = (0..vec.len())
            .filter(|&index| index == 0 || thread_rng().gen_bool(0.02))
            .collect();

        indices
            .iter()
            .zip(indices.iter().skip(1).chain(once(&vec.len())))
            .map(|(&start, &end)| vec[start..end].into())
            .collect()
    }

    #[test]
    fn compare_gen() {
        let num_of_entries_iter = [25, 50, 20].to_vec().into_iter();
        let max_entry_capacity_iter = num_of_entries_iter.clone().rev();

        for (num_of_entries, max_entry_capacity) in num_of_entries_iter.zip(max_entry_capacity_iter)
        {
            let data_def = gen_def_data(num_of_entries, max_entry_capacity);
            compare_data(build_data(&data_def), build_data(&data_def));
        }
    }

    #[test]
    fn compare_types_zeros() {
        // Define an arbitray data representation for sg-data
        let vec_def: Vec<(u8, usize)> = vec![(0, 44), (0, 511), (0, 44)];

        // define the data0 to be a vector of vectors
        let data_a = SgData::from(
            vec_def
                .iter()
                .map(|(val, size)| vec![*val; *size])
                .collect::<Vec<_>>(),
        );

        let vec = vec_def
            .iter()
            .flat_map(|(val, size)| vec![*val; *size])
            .collect::<Vec<_>>();

        let iovec = Iovec::from(iovec {
            iov_base: vec.as_ptr() as *mut c_void,
            iov_len: vec.len(),
        });

        let data_b = SgData::from(iovec);

        compare_data(data_a, data_b);
    }

    #[test]
    fn combine() {
        let vec_def: Vec<(u8, usize)> = vec![(75, 44), (144, 80), (77, 44)];

        // define the data0 to be a vector of vectors
        let data0: SgData = SgData::from(
            vec_def
                .iter()
                .map(|(val, size)| vec![*val; *size])
                .collect::<Vec<_>>(),
        );

        let mut data_combined = data0.masked(0..usize::MAX);

        let vec1 = vec_def
            .iter()
            .flat_map(|(val, size)| vec![*val; *size])
            .collect::<Vec<_>>();

        // define the data1 to be a direct vector
        let data1: SgData = SgData::from(vec1);

        let data1_elements = data1.masked(0..usize::MAX);
        data_combined.extend(Some(data1_elements));

        let mut vec_copied = data0.to_vec();
        vec_copied.extend(data1.to_vec());
        assert_eq!(data_combined.to_vec(), vec_copied);
    }

    #[cfg(feature = "nightly_features")]
    mod benches {
        use super::*;
        use test::Bencher;

        #[bench]
        fn direct_serialize(b: &mut Bencher) {
            let data: SgData = SgData::from(vec![0x45_u8; 8192]);

            b.iter(|| serialize(&data).unwrap())
        }

        #[bench]
        fn direct_deserialize(b: &mut Bencher) {
            let data: SgData = SgData::from(vec![0x45_u8; 8192]);
            let buf = serialize(&data).unwrap();

            b.iter(|| deserialize::<SgData>(&buf).unwrap());
        }

        #[bench]
        fn sglist_serialize(b: &mut Bencher) {
            let sgvec = vec![vec![0x45_u8; 4096]; 2];
            let data: SgData = unsafe { SgData::from_sglist(create_sglist(sgvec)) };

            b.iter(|| serialize(&data).unwrap())
        }

        #[bench]
        fn sglist_deserialize(b: &mut Bencher) {
            let sgvec = vec![vec![0x45_u8; 4096]; 2];
            let data: SgData = unsafe { SgData::from_sglist(create_sglist(sgvec)) };
            let buf = serialize(&data).unwrap();

            b.iter(|| deserialize::<SgData>(&buf).unwrap());
        }

        #[bench]
        fn iovec_serialize(b: &mut Bencher) {
            let data: SgData = vec![vec![0x46; 4096]; 2]
                .into_iter()
                .map(vec_into_iovec)
                .map(Iovec::from)
                .collect();

            b.iter(|| serialize(&data).unwrap())
        }

        #[bench]
        fn iovec_deserialize(b: &mut Bencher) {
            let data: SgData = vec![vec![0x46; 4096]; 2]
                .into_iter()
                .map(vec_into_iovec)
                .map(Iovec::from)
                .collect();
            let buf = serialize(&data).unwrap();

            b.iter(|| deserialize::<SgData>(&buf).unwrap());
        }
    }
}