load-buffer 0.1.0

#![cfg_attr(not(feature = "std"), no_std)]

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

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

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, buf: &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, buf: &mut [u8]) -> Result<(), Self::Error> {
        self.read_exact(buf)
    }
}

/// 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,
{
    pub fn new_static(loader: Loader) -> Self {
        Self {
            loader,
            buffer: [0u8; N],
            offset: N,
        }
    }
}

#[cfg(feature = "alloc")]
impl<Loader> BufferedLoader<Vec<u8>, Loader>
where
    Loader: Load,
{
    pub fn new_dynamic(loader: Loader, buffer_size: usize) -> Self {
        let mut buffer = Vec::with_capacity(buffer_size);
        buffer.resize(buffer_size, 0);
        Self {
            loader,
            buffer,
            offset: buffer_size,
        }
    }

    pub fn resize(&mut self, new_size: usize) {
        let invalid = new_size > self.buffer.len();

        self.buffer.resize(new_size, 0);
        if invalid {
            self.offset = new_size;
        }
    }
}

impl<Buffer, Loader> BufferedLoader<Buffer, Loader>
where
    Buffer: AsRef<[u8]> + AsMut<[u8]> + Debug,
    Loader: Load,
{
    /// 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 buf: &mut [u8]) -> Result<(), Self::Error> {
        if buf.is_empty() {
            return Ok(());
        }

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

        let buffer_remaining = buffer_size - self.offset;

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

            let buffers = buf.len() / buffer_size;
            let remainder = buf.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 buf[..buffers_bytes])?;
                buf = &mut buf[buffers_bytes..];
            }

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