bincode-next 3.0.0-rc.2

//! This module contains writer-based structs and traits.
//!
//! Because `std::io::Write` is only limited to `std` and not `core`, we provide our own [Writer].
#![allow(unsafe_code)]

use crate::error::EncodeError;

/// Trait that indicates that a struct can be used as a destination to encode data too. This is used by [Encode]
///
/// [Encode]: ../trait.Encode.html
pub trait Writer {
    /// Write `bytes` to the underlying writer. Exactly `bytes.len()` bytes must be written, or else an error should be returned.
    ///
    /// # Errors
    ///
    /// Returns `EncodeError::UnexpectedEnd` if the writer does not have enough space.
    fn write(
        &mut self,
        bytes: &[u8],
    ) -> Result<(), EncodeError>;

    /// Write a single byte to the underlying writer.
    ///
    /// # Errors
    ///
    /// Returns `EncodeError::UnexpectedEnd` if the writer does not have enough space.
    #[inline(always)]
    fn write_u8(
        &mut self,
        value: u8,
    ) -> Result<(), EncodeError> {
        self.write(&[value])
    }
}

impl<T: Writer> Writer for &mut T {
    #[inline]
    fn write(
        &mut self,
        bytes: &[u8],
    ) -> Result<(), EncodeError> {
        (**self).write(bytes)
    }

    #[inline]
    fn write_u8(
        &mut self,
        value: u8,
    ) -> Result<(), EncodeError> {
        (**self).write_u8(value)
    }
}

/// A helper struct that implements `Writer` for a `&[u8]` slice.
///
/// ```
/// use bincode_next::enc::write::SliceWriter;
/// use bincode_next::enc::write::Writer;
///
/// let destination = &mut [0u8; 100];
/// let mut writer = SliceWriter::new(destination);
/// writer.write(&[1, 2, 3, 4, 5]).unwrap();
///
/// assert_eq!(writer.bytes_written(), 5);
/// assert_eq!(destination[0..6], [1, 2, 3, 4, 5, 0]);
/// ```
pub struct SliceWriter<'storage> {
    slice: &'storage mut [u8],
    original_length: usize,
}

impl<'storage> SliceWriter<'storage> {
    /// Create a new instance of `SliceWriter` with the given byte array.
    pub const fn new(bytes: &'storage mut [u8]) -> Self {
        let original = bytes.len();
        Self {
            slice: bytes,
            original_length: original,
        }
    }

    /// Return the amount of bytes written so far.
    #[must_use]
    pub const fn bytes_written(&self) -> usize {
        self.original_length - self.slice.len()
    }
}

impl Writer for SliceWriter<'_> {
    #[inline(always)]
    fn write(
        &mut self,
        bytes: &[u8],
    ) -> Result<(), EncodeError> {
        if bytes.len() > self.slice.len() {
            return crate::error::cold_encode_error_unexpected_end();
        }
        // SAFETY: `bytes.len() <= self.slice.len()` is checked above.
        // `copy_nonoverlapping` is safe since the pointers won't overlap and bounds are exact.
        // Advancing the pointer by `bytes.len()` is safe for the same reason.
        unsafe {
            core::ptr::copy_nonoverlapping(bytes.as_ptr(), self.slice.as_mut_ptr(), bytes.len());
            let ptr = self.slice.as_mut_ptr().add(bytes.len());
            let len = self.slice.len() - bytes.len();
            self.slice = core::slice::from_raw_parts_mut(ptr, len);
        }

        Ok(())
    }

    #[inline(always)]
    fn write_u8(
        &mut self,
        value: u8,
    ) -> Result<(), EncodeError> {
        if self.slice.is_empty() {
            return crate::error::cold_encode_error_unexpected_end();
        }
        // SAFETY: `!self.slice.is_empty()` is checked above.
        unsafe {
            *self.slice.as_mut_ptr() = value;
            let ptr = self.slice.as_mut_ptr().add(1);
            let len = self.slice.len() - 1;
            self.slice = core::slice::from_raw_parts_mut(ptr, len);
        }

        Ok(())
    }
}

/// A writer that counts how many bytes were written. This is useful for e.g. pre-allocating buffers before writing to them.
#[derive(Default)]
pub struct SizeWriter {
    /// the amount of bytes that were written so far
    pub bytes_written: usize,
}
impl Writer for SizeWriter {
    #[inline(always)]
    fn write(
        &mut self,
        bytes: &[u8],
    ) -> Result<(), EncodeError> {
        self.bytes_written += bytes.len();

        Ok(())
    }

    #[inline(always)]
    fn write_u8(
        &mut self,
        _: u8,
    ) -> Result<(), EncodeError> {
        self.bytes_written += 1;

        Ok(())
    }
}