chunked_buffer/

lib.rs

#![no_std]

extern crate alloc;

use alloc::boxed::Box;
use alloc::collections::VecDeque;
use alloc::vec::Vec;
use core::cmp::min;
use core::ops::{Index, IndexMut};

/// The default chunk size
pub const DEFAULT_CHUNK_SIZE: usize = 256;

/// A deque style buffer that can be written to and read from.
/// See `GenericChunkedBuffer` for more details.
pub type ChunkedBuffer = GenericChunkedBuffer<DEFAULT_CHUNK_SIZE>;

/// A deque style buffer that can be written to and read from.
///
/// The buffer is composed of a series of fixed size chunks.
///
/// This structure is useful for memory constrained environments.  It limits the size of contiguous
/// allocations and incrementally releases memory as the buffer is consumed.
///
/// Example code:
/// ```
/// use chunked_buffer::ChunkedBuffer;
///
/// let mut buf = ChunkedBuffer::new();
/// buf.write(&[1, 2, 3]);
/// let mut dest = [0; 10];
/// let n = buf.read(&mut dest);
/// assert_eq!(n, 3);
/// assert_eq!(dest, [1, 2, 3, 0, 0, 0, 0, 0, 0, 0]);
/// ```
///
// Invariants:
// - `chunks` is never empty
// - all chunks are the same size
// - the read position is always in the first chunk
// - the write position is always in the last chunk
pub struct GenericChunkedBuffer<const CHUNK_SIZE: usize> {
    // the inclusive position of the first octet in the first chunk
    write_pos: usize,
    // the exclusive position of the last octet written in the last chunk
    read_pos: usize,
    chunks: VecDeque<Box<[u8; CHUNK_SIZE]>>,
}

impl<const CHUNK_SIZE: usize> GenericChunkedBuffer<CHUNK_SIZE> {
    /// Create a new, empty buffer
    pub fn new() -> Self {
        let mut s = GenericChunkedBuffer {
            write_pos: 0,
            read_pos: 0,
            chunks: VecDeque::new(),
        };
        s.chunks.push_back(Box::new([0; CHUNK_SIZE]));
        s
    }

    /// The number of bytes in the buffer
    pub fn len(&self) -> usize {
        (self.chunks.len() - 1) * CHUNK_SIZE + self.write_pos - self.read_pos
    }

    /// Returns true if the buffer is empty
    pub fn is_empty(&self) -> bool {
        self.chunks.len() == 1 && self.write_pos == self.read_pos
    }

    /// Consumes as many bytes as possible from the start of the buffer
    /// and writes them to `dest`
    pub fn read(&mut self, dest: &mut [u8]) -> usize {
        let mut nread = 0;
        while nread < dest.len() && !self.is_empty() {
            let chunk = &*self.chunks[0];
            let start = self.read_pos;
            let mut end = min(CHUNK_SIZE, start + dest.len() - nread);
            if self.chunks.len() == 1 {
                end = end.min(self.write_pos);
            }
            let n = end - start;
            dest[nread..nread + n].copy_from_slice(&chunk[start..end]);
            if end == CHUNK_SIZE {
                self.read_pos = 0;
                self.chunks.pop_front();
                // self.chunks cannot be empty here, because we must have created a new chunk
                // when we wrote the last byte
            } else {
                self.read_pos = end;
            }
            nread += n;
        }
        nread
    }

    /// pushes all the bytes in `src` to the end of the buffer
    pub fn write(&mut self, src: &[u8]) {
        let mut nwritten = 0;
        while nwritten < src.len() {
            let chunk = &mut **self.chunks.back_mut().unwrap();
            let start = self.write_pos;
            let end = min(CHUNK_SIZE, start + src.len() - nwritten);
            let n = end - start;
            chunk[start..end].copy_from_slice(&src[nwritten..nwritten + n]);
            nwritten += n;
            if end == CHUNK_SIZE {
                self.write_pos = 0;
                self.chunks.push_back(Box::new([0; CHUNK_SIZE]));
            } else {
                self.write_pos = end;
            }
        }
    }

    /// Returns an iterator over u8 contained in buffer
    pub fn iter(&self) -> Iter<CHUNK_SIZE> {
        Iter::new(self)
    }

    /// Returns an iterator over the chunks contained in buffer
    pub fn iter_chunks(&self) -> IterChunk<CHUNK_SIZE> {
        IterChunk::new(self)
    }

    /// Converts the buffer into a vector
    pub fn to_vec(&self) -> Vec<u8> {
        let mut vec = Vec::with_capacity(self.len());
        for chunk in self.iter_chunks() {
            vec.extend_from_slice(chunk);
        }
        vec
    }
}

impl<const CHUNK_SIZE: usize> Default for GenericChunkedBuffer<CHUNK_SIZE> {
    fn default() -> Self {
        Self::new()
    }
}

impl<const CHUNK_SIZE: usize> Index<usize> for GenericChunkedBuffer<CHUNK_SIZE> {
    type Output = u8;

    fn index(&self, index: usize) -> &Self::Output {
        if self.is_empty() {
            panic!("indexed into an empty buffer");
        }
        if index >= self.len() {
            panic!("out of bounds access");
        }
        let pos = index % CHUNK_SIZE;
        let index = index / CHUNK_SIZE;
        &self.chunks[index][pos]
    }
}

impl<const CHUNK_SIZE: usize> IndexMut<usize> for GenericChunkedBuffer<CHUNK_SIZE> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        if self.is_empty() {
            panic!("indexed into an empty buffer");
        }
        if index >= self.len() {
            panic!("out of bounds access");
        }
        let pos = index % CHUNK_SIZE;
        let index = index / CHUNK_SIZE;
        &mut self.chunks[index][pos]
    }
}

pub struct Iter<'a, const CHUNK_SIZE: usize> {
    nread: usize,
    index: usize,
    read_pos: usize,
    buf: &'a GenericChunkedBuffer<CHUNK_SIZE>,
}

impl<'a, const CHUNK_SIZE: usize> Iter<'a, CHUNK_SIZE> {
    fn new(buf: &'a GenericChunkedBuffer<CHUNK_SIZE>) -> Self {
        Iter {
            nread: 0,
            index: 0,
            read_pos: buf.read_pos,
            buf,
        }
    }
}

impl<const CHUNK_SIZE: usize> Iterator for Iter<'_, CHUNK_SIZE> {
    type Item = u8;
    fn next(&mut self) -> Option<Self::Item> {
        if self.nread == self.buf.len() {
            None
        } else {
            let byte = self.buf.chunks[self.index][self.read_pos];
            self.nread += 1;
            self.read_pos += 1;
            if self.read_pos == CHUNK_SIZE {
                self.read_pos = 0;
                self.index += 1;
            }
            Some(byte)
        }
    }
}

pub struct IterChunk<'a, const CHUNK_SIZE: usize> {
    index: usize,
    buf: &'a GenericChunkedBuffer<CHUNK_SIZE>,
}

impl<'a, const CHUNK_SIZE: usize> IterChunk<'a, CHUNK_SIZE> {
    fn new(buf: &'a GenericChunkedBuffer<CHUNK_SIZE>) -> Self {
        IterChunk { index: 0, buf }
    }
}

impl<'a, const CHUNK_SIZE: usize> Iterator for IterChunk<'a, CHUNK_SIZE> {
    type Item = &'a [u8];
    fn next(&mut self) -> Option<Self::Item> {
        if self.index == self.buf.chunks.len() {
            None
        } else {
            let begin = if self.index == 0 {
                self.buf.read_pos
            } else {
                0
            };
            let end = if self.index + 1 == self.buf.chunks.len() {
                self.buf.write_pos
            } else {
                CHUNK_SIZE
            };
            let slice = &self.buf.chunks[self.index][begin..end];
            self.index += 1;
            if slice.is_empty() {
                None
            } else {
                Some(slice)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    type Buf = GenericChunkedBuffer<4>;

    #[test]
    fn test() {
        let mut buf = Buf::new();
        assert_eq!(buf.len(), 0);
        assert!(buf.is_empty());

        let mut dest = [0; 10];
        assert_eq!(buf.read(&mut dest), 0);

        buf.write(&[1, 2, 3]);
        assert_eq!(buf.len(), 3);
        assert!(!buf.is_empty());

        let mut dest = [0; 10];
        assert_eq!(buf.read(&mut dest), 3);
        assert_eq!(dest, [1, 2, 3, 0, 0, 0, 0, 0, 0, 0]);
        assert_eq!(buf.len(), 0);
        assert!(buf.is_empty());

        buf.write(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        assert_eq!(buf.len(), 10);
        assert!(!buf.is_empty());

        let mut dest = [0; 9];
        assert_eq!(buf.read(&mut dest), 9);
        assert_eq!(dest, [1, 2, 3, 4, 5, 6, 7, 8, 9]);
        assert_eq!(buf.len(), 1);

        let mut dest = [0; 9];
        assert_eq!(buf.read(&mut dest), 1);
        assert_eq!(dest, [10, 0, 0, 0, 0, 0, 0, 0, 0]);
        assert_eq!(buf.len(), 0);
        assert!(buf.is_empty());

        buf.write(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        assert_eq!(buf.len(), 10);
        assert!(!buf.is_empty());

        let mut dest = [0; 5];
        assert_eq!(buf.read(&mut dest), 5);
        assert_eq!(dest, [1, 2, 3, 4, 5]);
        assert_eq!(buf.len(), 5);
        assert!(!buf.is_empty());

        let mut dest = [0; 5];
        assert_eq!(buf.read(&mut dest), 5);
        assert_eq!(dest, [6, 7, 8, 9, 10]);
        assert_eq!(buf.len(), 0);
        assert!(buf.is_empty());

        // at chunk boundary
        let mut dest = [0; 5];
        buf.write(&[99]);
        assert_eq!(buf.len(), 1);
        assert_eq!(buf.read(&mut dest), 1);
        assert_eq!(dest, [99, 0, 0, 0, 0]);
    }
    #[test]
    fn test_iterator() {
        // empty buffer
        let mut buf = Buf::new();
        let mut i = buf.iter();
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // empty write
        buf = Buf::new();
        buf.write(&[]);
        i = buf.iter();
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // first chunk not filled
        buf = Buf::new();
        buf.write(&[1, 2, 3]);
        i = buf.iter();
        for count in 1..4 {
            assert_eq!(i.next(), Some(count));
        }
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // first chunk filled
        buf = Buf::new();
        buf.write(&[1, 2, 3, 4]);
        i = buf.iter();
        for count in 1..5 {
            assert_eq!(i.next(), Some(count));
        }
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // third chunk not filled
        buf = Buf::new();
        buf.write(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        i = buf.iter();
        for count in 1..11 {
            assert_eq!(i.next(), Some(count));
        }
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // third chunk filled
        buf = Buf::new();
        buf.write(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
        i = buf.iter();
        for count in 1..13 {
            assert_eq!(i.next(), Some(count));
        }
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // fourth chunk not filled
        buf = Buf::new();
        buf.write(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]);
        i = buf.iter();
        for count in 1..14 {
            assert_eq!(i.next(), Some(count));
        }
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
    }
    #[test]
    fn test_iter_chunks() {
        // empty buffer
        let mut buf = Buf::new();
        let mut i = buf.iter_chunks();
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // empty write
        buf = Buf::new();
        buf.write(&[]);
        i = buf.iter_chunks();
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // first chunk not filled
        buf = Buf::new();
        buf.write(&[1, 2, 3]);
        i = buf.iter_chunks();
        assert_eq!(i.next(), Some(&[1, 2, 3][..]));
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // first chunk filled
        buf = Buf::new();
        buf.write(&[1, 2, 3, 4]);
        i = buf.iter_chunks();
        assert_eq!(i.next(), Some(&[1, 2, 3, 4][..]));
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // third chunk not filled
        buf = Buf::new();
        buf.write(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        i = buf.iter_chunks();
        assert_eq!(i.next(), Some(&[1, 2, 3, 4][..]));
        assert_eq!(i.next(), Some(&[5, 6, 7, 8][..]));
        assert_eq!(i.next(), Some(&[9, 10][..]));
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // third chunk filled
        buf = Buf::new();
        buf.write(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
        i = buf.iter_chunks();
        assert_eq!(i.next(), Some(&[1, 2, 3, 4][..]));
        assert_eq!(i.next(), Some(&[5, 6, 7, 8][..]));
        assert_eq!(i.next(), Some(&[9, 10, 11, 12][..]));
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);

        // fourth chunk not filled
        buf = Buf::new();
        buf.write(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]);
        i = buf.iter_chunks();
        assert_eq!(i.next(), Some(&[1, 2, 3, 4][..]));
        assert_eq!(i.next(), Some(&[5, 6, 7, 8][..]));
        assert_eq!(i.next(), Some(&[9, 10, 11, 12][..]));
        assert_eq!(i.next(), Some(&[13][..]));
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
        assert_eq!(i.next(), None);
    }
    #[test]
    fn test_index() {
        let mut buf = Buf::new();
        buf.write(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]);
        for i in 0..14 {
            assert_eq!(buf[i], i as u8);
        }
    }
    #[test]
    #[should_panic]
    fn test_index_panic_on_empty() {
        let buf = Buf::new();
        buf[0];
    }
    #[test]
    #[should_panic]
    fn test_index_out_of_bounds() {
        let mut buf = Buf::new();
        buf.write(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]);
        buf[13];
    }
    #[test]
    fn test_index_mut() {
        let mut buf = Buf::new();
        buf.write(&[0u8; 14]);
        for i in 0..14 {
            buf[i] = i as u8;
        }
        let mut want = Buf::new();
        want.write(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]);
        assert_eq!(buf.to_vec(), want.to_vec());
    }
    #[test]
    #[should_panic]
    fn test_index_mut_panic_on_empty() {
        let mut buf = Buf::new();
        buf[0] = 3;
    }
    #[test]
    #[should_panic]
    fn test_index_mut_out_of_bounds() {
        let mut buf = Buf::new();
        buf.write(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]);
        buf[13] = 3;
    }
}