genmarkov/src/chunking.rs

//! Stream related things
use super::*;
use std::{  
    task::{
	Poll,
	Context,
    },
    pin::Pin,
    marker::PhantomData,
};
use tokio::{
    io::{
	AsyncBufRead,
	AsyncRead,
    },
    prelude::*,
};
use futures::{
    stream::{
	Stream,
	StreamExt,
	Fuse,
    },
};
use pin_project::pin_project;

/// Converts a stream of byte-containing objects into an `AsyncRead` and `AsyncBufRead`er.
#[pin_project]
pub struct StreamReader<I, T>
where I: Stream<Item=T>
{
    #[pin]
    source: Fuse<I>,
    buffer: Vec<u8>,
}

impl<T, I> StreamReader<I, T>
where I: Stream<Item=T>,
      T: AsRef<[u8]>
{
    /// The current buffer
    pub fn buffer(&self) -> &[u8]
    {
	&self.buffer[..]
    }
    /// Consume into the original stream
    pub fn into_inner(self) -> I
    {
	self.source.into_inner()
    }
    /// Create a new instance with a buffer capacity
    pub fn with_capacity(source: I, cap: usize) -> Self
    {
	Self {
	    source: source.fuse(),
	    buffer: Vec::with_capacity(cap)
	}
    }
    /// Create a new instance from this stream
    pub fn new(source: I) -> Self
    {
	Self {
	    source: source.fuse(),
	    buffer: Vec::new(),
	}
    }
    /// Attempt to add to this buffer
    #[cold] fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<usize>
    {
	let this = self.project();
	match this.source.poll_next(cx) {
	    Poll::Ready(None) => Poll::Ready(0),
	    Poll::Ready(Some(buf)) if buf.as_ref().len() > 0 => {
		let buf = buf.as_ref();
		this.buffer.extend_from_slice(buf);
		Poll::Ready(buf.len())
	    },
	    _ => Poll::Pending,
	}
    }
}

impl<T: AsRef<[u8]>, I: Stream<Item=T>> AsyncRead for StreamReader<I,T>
{
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
	let this = self.project();
	if this.buffer.len() != 0 {
	    // We can fill the whole buffer, do it.
	    Poll::Ready(Ok(bytes::copy_slice(buf, this.buffer.drain(..buf.len()).as_slice())))
	} else  {
	    // Buffer is empty, try to fill it
	    match match this.source.poll_next(cx) {
		Poll::Ready(None) => Poll::Ready(0),
		Poll::Ready(Some(buf)) if buf.as_ref().len() > 0 => {
		    let buf = buf.as_ref();
		    this.buffer.extend_from_slice(buf);
		    Poll::Ready(buf.len())
		},
		_ => Poll::Pending,
	    } {
		Poll::Ready(0) => Poll::Ready(Ok(0)),
		Poll::Ready(x) => {
		    // x has been written
		    Poll::Ready(Ok(bytes::copy_slice(buf, this.buffer.drain(..x).as_slice())))
		},
		_ => Poll::Pending,
	    }	
	}
	
    }
    
}

impl<T: AsRef<[u8]>, I: Stream<Item=T>> AsyncBufRead for StreamReader<I,T>
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
	let this = self.project();
	if this.buffer.len() < 1 {
	    // Fetch more into buffer
	    match match this.source.poll_next(cx) {
		Poll::Ready(None) => Poll::Ready(0),
		Poll::Ready(Some(buf)) if buf.as_ref().len() > 0 => {
		    let buf = buf.as_ref();
		    this.buffer.extend_from_slice(buf);
		    Poll::Ready(buf.len())
		},
		_ => Poll::Pending,
	    } {
		Poll::Ready(0) => Poll::Ready(Ok(&[])), // should we return EOF error here?
		Poll::Ready(x) => Poll::Ready(Ok(&this.buffer[..x])),
		_ => Poll::Pending
	    }
	} else {
	    Poll::Ready(Ok(&this.buffer[..]))
	}
    }
    fn consume(self: Pin<&mut Self>, amt: usize) {
	self.project().buffer.drain(..amt);
    }
}

#[cfg(test)]
mod tests
{
    use super::*;
    use tokio::{
	sync::{
	    mpsc,
	},
    };
    #[tokio::test]
    async fn stream_of_vec()
    {
	let (mut tx, rx) = mpsc::channel(16);
	let sender = tokio::spawn(async move {
	    tx.send("Hello ").await.unwrap();
	    tx.send("world").await.unwrap();
	    tx.send("\n").await.unwrap();
	    tx.send("How ").await.unwrap();
	    tx.send("are ").await.unwrap();
	    tx.send("you").await.unwrap();
	});

	let mut reader = StreamReader::new(rx);

	let mut output = String::new();
	let mut read;
	while {read = reader.read_line(&mut output).await.expect("Failed to read"); read!=0} {
	    println!("Read: {}", read);
	}

	println!("Done: {:?}", output);
	sender.await.expect("Child panic");
	assert_eq!(&output[..], "Hello world\nHow are you");
    }
}

/// A stream that chunks its input.
#[pin_project]
pub struct ChunkingStream<S, T, Into=Vec<T>>
{
    #[pin] stream: Fuse<S>,
    buf: Vec<T>,
    cap: usize,
    _output: PhantomData<Into>,

    push_now: bool,
}


impl<S, T, Into>  ChunkingStream<S,T, Into>
where S: Stream<Item=T>,
      Into: From<Vec<T>>
{
    pub fn new(stream: S, sz: usize) -> Self
    {
	Self {
	    stream: stream.fuse(),
	    buf: Vec::with_capacity(sz),
	    cap: sz,
	    _output: PhantomData,
	    push_now: false,
	}
    }
    pub fn into_inner(self) -> S
    {
	self.stream.into_inner()
    }
    pub fn cap(&self) -> usize
    {
	self.cap
    }
    pub fn buffer(&self) -> &[T]
    {
	&self.buf[..]
    }

    pub fn get_ref(&self) -> &S
    {
	self.stream.get_ref()
    }

    pub fn get_mut(&mut self)-> &mut S
    {
	self.stream.get_mut()
    }

    /// Force the next read to send the buffer even if it's not full.
    ///
    /// # Note
    /// The buffer still won't send if it's empty.
    pub fn push_now(&mut self)
    {
	self.push_now= true;
    }

    /// Consume into the current held buffer
    pub fn into_buffer(self) -> Vec<T>
    {
	self.buf
    }

    /// Take the buffer now
    pub fn take_now(&mut self) -> Into
    {
	std::mem::replace(&mut self.buf, Vec::with_capacity(self.cap)).into()
    }
}

impl<S, T, Into> Stream for ChunkingStream<S,T, Into>
where S: Stream<Item=T>,
      Into: From<Vec<T>>
{
    type Item = Into;
    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
	while !(self.push_now && !self.buf.is_empty()) && self.buf.len() < self.cap {
	    // Buffer isn't full, keep filling
	    let this = self.as_mut().project();

	    match this.stream.poll_next(cx) {
		Poll::Ready(None) => {
		    // Stream is over
		    break;
		},
		Poll::Ready(Some(item)) => {
		    this.buf.push(item);
		},
		_ => return Poll::Pending,
	    }
	}
	debug!("Sending buffer of {} (cap {})", self.buf.len(), self.cap);
	// Buffer is full or we reach end of stream
	Poll::Ready(if self.buf.len() == 0 {
	    None
	} else {
	    let this = self.project();
	    *this.push_now = false;
	    let output = std::mem::replace(this.buf, Vec::with_capacity(*this.cap));
	    Some(output.into())
	})
    }
}