streams1.rs

use byte_pool::{Block, BytePool};
use byteorder::{BigEndian, ByteOrder};
use std::io::{Cursor, Result};

fn main() {
    let pool = BytePool::<Vec<u8>>::new();

    let source = vec![1; 1024 * 1024];
    let incoming_bytes = Cursor::new(&source);
    let mut dest = Vec::new();

    incoming_bytes
        .chunks(&pool, 1024)
        .map(|mut chunk| {
            // wrap the chunk into something
            chunk.insert(0, 0);
            chunk.push(255);
            chunk
        })
        .pipe(&mut dest)
        .unwrap();

    assert_eq!(dest.len(), source.len() + (source.len() / 1024) * 2);

    // Parse incoming data to `u32`
    let incoming_bytes = Cursor::new(&source);
    let u32s = incoming_bytes
        .chunks(&pool, 4)
        .map(|chunk| BigEndian::read_u32(&chunk))
        .collect::<Vec<u32>>();

    assert_eq!(u32s[0], u32::from_le_bytes([1, 1, 1, 1]));
}

pub trait ChunksExt: std::io::Read + Sized {
    fn chunks<'a>(self, pool: &'a BytePool<Vec<u8>>, chunk_size: usize) -> Chunks<'a, Self> {
        Chunks::new(pool, self, chunk_size)
    }
}

impl<T: std::io::Read> ChunksExt for T {}

pub trait PipeExt<'a>: Iterator<Item = Block<'a, Vec<u8>>> + Sized {
    fn pipe<T: std::io::Write>(mut self, mut dest: T) -> Result<()> {
        while let Some(val) = self.next() {
            dest.write_all(&val)?;
        }

        Ok(())
    }
}

impl<'a, T: Iterator<Item = Block<'a, Vec<u8>>>> PipeExt<'a> for T {}

#[derive(Debug)]
pub struct Chunks<'a, T: std::io::Read> {
    pool: &'a BytePool<Vec<u8>>,
    source: T,
    chunk_size: usize,
    current: Option<(Block<'a, Vec<u8>>, usize)>,
    done: bool,
}

impl<'a, T: std::io::Read> Chunks<'a, T> {
    pub fn new(pool: &'a BytePool<Vec<u8>>, source: T, chunk_size: usize) -> Self {
        Self {
            pool,
            source,
            chunk_size,
            current: None,
            done: false,
        }
    }
}

impl<'a, T: std::io::Read> Iterator for Chunks<'a, T> {
    type Item = Block<'a, Vec<u8>>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            return None;
        }

        let mut buf = self.pool.alloc(self.chunk_size);
        buf.truncate(self.chunk_size);
        let mut pos = 0;

        while pos < self.chunk_size {
            let read = self.source.read(&mut buf[pos..]).unwrap(); // TODO: handle error
            if read == 0 {
                self.done = true;
                break;
            }
            pos += read;
        }

        if pos == 0 {
            return None;
        }

        if pos + 1 < buf.len() {
            // shrink buffer
            buf.realloc(pos + 1);
        }

        Some(buf)
    }
}