load-buffer 1.0.0

//!
//! ```no_run
//! # use load_buffer::*;
//! # #[derive(Default)]
//! # struct Loader;
//! # impl Load for Loader {
//! #   type Error = core::convert::Infallible;
//! #   fn load(&mut self, destination: &mut [u8]) -> Result<(), Self::Error> { unimplemented!() }
//! # }
//! // Static, backed by the array size you give it:
//! let mut loader: BufferedLoader<[u8; 64], _> =
//!     BufferedLoader::new_static(Loader::default());
//!
//! // Dynamic, backed by alloc::vec::Vec<u8>:
//! let mut loader = BufferedLoader::new_dynamic(Loader::default(), 64);
//!
//! // Dynamic can be resized.  This invalidates the buffer.
//! loader.resize(268435456);
//!
//! // Heap, backed by alloc::boxed::Box<[u8]>:
//! let mut loader = BufferedLoader::new_heap(Loader::default(), 268435456);
//!
//! // Heap can not be resized, but will not fill the stack for huge sizes.
//! // You can also supply your own buffer type, as long as it implements
//! // AsRef<[u8]> + AsMut<[u8]> + Debug.
//! ```
#![cfg_attr(not(feature = "std"), no_std)]

#[cfg(feature = "alloc")]
extern crate alloc;

#[cfg(feature = "alloc")]
use alloc::vec;

#[cfg(feature = "alloc")]
use alloc::vec::Vec;

#[cfg(feature = "alloc")]
use alloc::boxed::Box;

use core::fmt::Debug;

/// Like [`std::io::Read`], but only supporting [`read_exact`], and with an
/// arbitrary error type.
///
/// [`std::io::Read`]: https://doc.rust-lang.org/std/io/trait.Read.html
/// [`read_exact`]: https://doc.rust-lang.org/std/io/trait.Read.html#method.read_exact
pub trait Load {
    type Error;

    fn load(&mut self, destination: &mut [u8]) -> Result<(), Self::Error>;
}

#[cfg(feature = "std")]
impl<T> Load for T
where
    T: std::io::Read,
{
    type Error = std::io::Error;
    fn load(&mut self, destination: &mut [u8]) -> Result<(), Self::Error> {
        self.read_exact(destination)
    }
}

/// A buffering wrapper around a [`Load`].
///
/// This fills some of the simpler niches of a [`std::io::BufReader`] for no_std
/// use-cases, and allowing a static buffer.
///
/// [`std::io::BufReader`]: https://doc.rust-lang.org/std/io/struct.BufReader.html
/// [`Load`]: trait.Load.html
#[derive(Clone, Debug)]
pub struct BufferedLoader<Buffer, Loader>
where
    Buffer: AsRef<[u8]> + AsMut<[u8]> + Debug,
    Loader: Load,
{
    loader: Loader,
    buffer: Buffer,

    /// The amount of buffer that has been read already.
    offset: usize,
}

impl<const N: usize, Loader> BufferedLoader<[u8; N], Loader>
where
    Loader: Load,
{
    /// Creates a new stack-based BufferedLoader.  Does not allocate.
    pub fn new_static(loader: Loader) -> Self {
        Self::new([0u8; N], loader)
    }
}

#[cfg(feature = "alloc")]
impl<Loader> BufferedLoader<Vec<u8>, Loader>
where
    Loader: Load,
{
    /// Creates a new heap-based BufferedLoader.  Allocates only once.
    pub fn new_dynamic(loader: Loader, buffer_size: usize) -> Self {
        Self::new(vec![0u8; buffer_size], loader)
    }

    pub fn resize(&mut self, new_size: usize) {
        self.buffer.resize(new_size, 0);
        self.offset = new_size;
    }
}

#[cfg(feature = "alloc")]
impl<Loader> BufferedLoader<Box<[u8]>, Loader>
where
    Loader: Load,
{
    /// Creates a new heap-based BufferedLoader.  Allocates only once.
    pub fn new_heap(loader: Loader, buffer_size: usize) -> Self {
        Self::new(vec![0u8; buffer_size].into_boxed_slice(), loader)
    }
}

impl<Buffer, Loader> BufferedLoader<Buffer, Loader>
where
    Buffer: AsRef<[u8]> + AsMut<[u8]> + Debug,
    Loader: Load,
{
    pub fn new(buffer: Buffer, loader: Loader) -> Self {
        let offset = buffer.as_ref().len();
        Self {
            loader,
            buffer,
            offset,
        }
    }

    /// Drop what's remaining in the buffer and give a mutable reference to the
    /// inner loader, so it can be seeked or otherwise manipulated.
    pub fn loader(&mut self) -> &mut Loader {
        self.offset = self.buffer.as_ref().len();
        &mut self.loader
    }
}

impl<Buffer, Loader> Load for BufferedLoader<Buffer, Loader>
where
    Buffer: AsRef<[u8]> + AsMut<[u8]> + Debug,
    Loader: Load,
{
    type Error = Loader::Error;

    /// Efficiently fill the destination buffer, calling the underlying
    /// [`Load::load`] as few times as possible.
    ///
    /// [`Load::Load`]: trait.Load.html#method.load
    fn load(&mut self, mut destination: &mut [u8]) -> Result<(), Self::Error> {
        if destination.is_empty() {
            return Ok(());
        }

        let buffer = self.buffer.as_mut();
        let buffer_size = buffer.len();

        let buffer_remaining = buffer_size - self.offset;

        if destination.len() <= buffer_remaining {
            // There are enough bytes left in the buffer to consume without
            // reading.
            let end = self.offset + destination.len();
            destination.copy_from_slice(&buffer[self.offset..end]);
            self.offset = end;
        } else {
            // Empty what's remaining in the buffer.
            if buffer_remaining > 0 {
                destination[..buffer_remaining].copy_from_slice(&buffer[self.offset..buffer_size]);
                let copied = buffer_size - self.offset;
                destination = &mut destination[copied..];
            }

            let buffers = destination.len() / buffer_size;
            let remainder = destination.len() % buffer_size;

            // Load full-sized chunks directly to the destination, bypassing
            // the buffer.
            if buffers > 0 {
                let buffers_bytes = buffers * buffer_size;
                self.loader.load(&mut destination[..buffers_bytes])?;

                destination = &mut destination[buffers_bytes..];
            }

            // Fill the buffer for the remainder, if there is any.
            if remainder > 0 {
                self.loader.load(buffer)?;
                destination.copy_from_slice(&buffer[..remainder]);
                self.offset = remainder;
            } else {
                // We have emptied the remaining buffer, so mark this empty.
                self.offset = buffer_size;
            }
        }
        Ok(())
    }
}

#[cfg(test)]
pub mod test {
    use core::convert::Infallible;

    use super::*;

    #[derive(Default)]
    struct Increment(u8);

    impl Load for Increment {
        type Error = Infallible;

        fn load(&mut self, destination: &mut [u8]) -> Result<(), Self::Error> {
            for each in destination {
                *each = self.0;
                self.0 = self.0.wrapping_add(1);
            }
            Ok(())
        }
    }

    #[test]
    fn test_static_buffered_loader() {
        let mut loader: BufferedLoader<[u8; 8], _> =
            BufferedLoader::new_static(Increment::default());
        let mut buf = [0u8; 4];
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [0, 1, 2, 3]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [4, 5, 6, 7]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [8, 9, 10, 11]);

        let mut buf = [0u8; 12];
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]);

        let mut buf = [0u8; 3];

        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [24, 25, 26]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [27, 28, 29]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [30, 31, 32]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [33, 34, 35]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [36, 37, 38]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [39, 40, 41]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [42, 43, 44]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [45, 46, 47]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [48, 49, 50]);
    }

    #[test]
    fn test_dynamic_buffered_loader() {
        let mut loader = BufferedLoader::new_dynamic(Increment::default(), 8);
        let mut buf = [0u8; 4];
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [0, 1, 2, 3]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [4, 5, 6, 7]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [8, 9, 10, 11]);

        let mut buf = [0u8; 12];
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]);

        let mut buf = [0u8; 3];

        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [24, 25, 26]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [27, 28, 29]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [30, 31, 32]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [33, 34, 35]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [36, 37, 38]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [39, 40, 41]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [42, 43, 44]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [45, 46, 47]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [48, 49, 50]);
    }

    #[test]
    fn test_boxed_slice_buffered_loader() {
        let mut loader = BufferedLoader::new_heap(Increment::default(), 8);
        let mut buf = [0u8; 4];
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [0, 1, 2, 3]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [4, 5, 6, 7]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [8, 9, 10, 11]);

        let mut buf = [0u8; 12];
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]);

        let mut buf = [0u8; 3];

        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [24, 25, 26]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [27, 28, 29]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [30, 31, 32]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [33, 34, 35]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [36, 37, 38]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [39, 40, 41]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [42, 43, 44]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [45, 46, 47]);
        loader.load(&mut buf).unwrap();
        assert_eq!(buf, [48, 49, 50]);
    }
}