async-codec 0.4.1

//! The `Framed` type for `embrio_core::io::{Read, Write}` streams

use core::pin::Pin;

use futures::prelude::*;
use futures::ready;
use futures::task::{Context, Poll};

use embrio_core::io::{Read, Write};

use super::*;

struct Buffer<B> {
    data: B,
    len: usize,
}

impl<B> Buffer<B>
where
    B: AsMut<[u8]> + AsRef<[u8]>,
{
    fn len(&self) -> usize {
        self.len
    }

    fn resize(&mut self, new_size: usize) -> bool {
        if new_size > self.data.as_mut().len() {
            return false;
        }
        self.len = new_size;
        true
    }

    // Note: This is a O(n) operation rather than the usual ringbuffer style so as
    // to ensure that the data is always contiguous.
    fn consume(&mut self, bytes: usize) {
        self.data.as_mut().rotate_left(bytes);
        self.resize(self.len - bytes);
    }

    fn cap(&self) -> usize {
        self.data.as_ref().len()
    }

    fn unused_mut(&mut self) -> &mut [u8] {
        &mut self.data.as_mut()[self.len..]
    }

    fn as_mut(&mut self) -> &mut [u8] {
        &mut self.data.as_mut()[..self.len]
    }
    fn as_ref(&self) -> &[u8] {
        &self.data.as_ref()[..self.len]
    }
}

/// A wrapper around a byte stream that uses an `Encode + Decode` type to
/// produce a `Sink + Stream`
pub struct Framed<S, C, B> {
    stream: S,
    buffer: Buffer<B>,
    offset: usize,
    codec: C,
}

impl<S, C, B> Framed<S, C, B>
where
    B: AsMut<[u8]> + AsRef<[u8]>,
{
    /// Create a new framed stream from a byte stream and a codec.
    pub fn new(stream: S, codec: C, buffer: B) -> Self {
        Framed {
            stream,
            codec,
            buffer: Buffer {
                data: buffer,
                len: 0,
            },
            offset: 0,
        }
    }
}

/// Errors arising from reading a frame.
#[derive(Debug)]
pub enum ReadFrameError<I, E> {
    /// There was an error in the underlying stream.
    Io(I),
    /// There was an error decoding the frame.
    Decode(E),
    /// The buffer filled without finding a full frame.
    Overflow,
}

impl<S, C, B> Stream for Framed<S, C, B>
where
    S: Read,
    C: Decode,
    B: AsMut<[u8]> + AsRef<[u8]>,
{
    type Item = Result<C::Item, ReadFrameError<S::Error, C::Error>>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let (mut reader, codec, buffer) = unsafe {
            let this = self.get_unchecked_mut();
            (
                Pin::new_unchecked(&mut this.stream),
                &mut this.codec,
                &mut this.buffer,
            )
        };
        loop {
            let (consumed, res) = codec.decode(buffer.as_mut());
            buffer.consume(consumed);
            match res {
                DecodeResult::Ok(value) => {
                    return Poll::Ready(Some(Ok(value)));
                }
                DecodeResult::Err(e) => {
                    return Poll::Ready(Some(Err(ReadFrameError::Decode(e))));
                }
                DecodeResult::UnexpectedEnd => {}
            }
            let orig_len = buffer.len();
            // If we're here, the decoder needed more bytes, but we don't have
            // space to put them :(
            if buffer.unused_mut().len() == 0 {
                return Poll::Ready(Some(Err(ReadFrameError::Overflow)));
            }
            let bytes = match ready!(reader.as_mut().poll_read(cx, buffer.unused_mut()))
                .map_err(ReadFrameError::Io)
            {
                Ok(buf) => buf,
                Err(e) => return Poll::Ready(Some(Err(e))),
            };
            if bytes == 0 {
                return Poll::Ready(None);
            }
            buffer.resize(orig_len + bytes);
        }
    }
}

/// Errors arising from writing a frame to a stream
#[derive(Debug)]
pub enum WriteFrameError<S, E> {
    /// An error in the underlying stream
    Io(S),
    /// An error from encoding the frame
    Encode(E),
    /// End of stream while writing frame
    UnexpectedEnd,
    /// Message will not fit in the static output buffer
    Overflow,
}

impl<S, C, B> Sink<C::Item> for Framed<S, C, B>
where
    S: Write,
    C: Encode,
    B: AsMut<[u8]> + AsRef<[u8]>,
{
    type Error = WriteFrameError<S::Error, C::Error>;
    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let (buffer, offset, mut writer) = unsafe {
            let this = self.get_unchecked_mut();
            (
                &mut this.buffer,
                &mut this.offset,
                Pin::new_unchecked(&mut this.stream),
            )
        };
        loop {
            if *offset == buffer.len() {
                buffer.resize(0);
                *offset = 0;
                return Poll::Ready(Ok(()));
            }
            let written = ready!(writer
                .as_mut()
                .poll_write(cx, &buffer.as_ref()[*offset..])
                .map_err(WriteFrameError::Io))?;
            if written == 0 {
                return Poll::Ready(Err(WriteFrameError::UnexpectedEnd));
            }
            *offset += written;
        }
    }
    fn start_send(self: Pin<&mut Self>, item: C::Item) -> Result<(), Self::Error> {
        let (buffer, codec) = unsafe {
            let this = self.get_unchecked_mut();
            (&mut this.buffer, &mut this.codec)
        };
        buffer.resize(buffer.cap());
        codec.reset();
        match codec.encode(&item, buffer.as_mut()) {
            EncodeResult::Ok(len) => {
                buffer.resize(len);
                return Ok(());
            }
            EncodeResult::Err(e) => return Err(WriteFrameError::Encode(e)),
            EncodeResult::Overflow(_) => return Err(WriteFrameError::Overflow),
        }
    }
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        ready!(self.as_mut().poll_ready(cx))?;

        unsafe { Pin::new_unchecked(&mut self.get_unchecked_mut().stream) }
            .poll_flush(cx)
            .map_err(WriteFrameError::Io)
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        ready!(self.as_mut().poll_flush(cx))?;

        unsafe { Pin::new_unchecked(&mut self.get_unchecked_mut().stream) }
            .poll_close(cx)
            .map_err(WriteFrameError::Io)
    }
}

#[cfg(test)]
mod test {
    use std::str::Utf8Error;

    use futures::prelude::*;

    use super::*;

    use embrio_core::io;

    struct LineCodec;

    impl Encode for LineCodec {
        type Item = String;
        type Error = ();
        fn encode(&mut self, item: &String, buf: &mut [u8]) -> EncodeResult<Self::Error> {
            let needed = item.as_bytes().len() + 1;
            if buf.len() < needed {
                return EncodeResult::Overflow(needed);
            }
            buf[..item.as_bytes().len()].copy_from_slice(item.as_bytes());
            buf[item.len()] = b'\n';
            Ok(needed).into()
        }
    }

    impl Decode for LineCodec {
        type Item = String;
        type Error = Utf8Error;

        fn decode(&mut self, buf: &mut [u8]) -> (usize, DecodeResult<String, Utf8Error>) {
            let newline = match buf.iter().position(|b| *b == b'\n') {
                Some(idx) => idx,
                None => return (0, DecodeResult::UnexpectedEnd),
            };
            let string_bytes = &buf[..newline];
            (
                newline + 1,
                std::str::from_utf8(string_bytes).map(String::from).into(),
            )
        }
    }

    const SHAKESPEARE: &str = r#"Now is the winter of our discontent
Made glorious summer by this sun of York.
Some are born great, some achieve greatness
And some have greatness thrust upon them.
Friends, Romans, countrymen - lend me your ears!
I come not to praise Caesar, but to bury him.
The evil that men do lives after them
The good is oft interred with their bones.
                    It is a tale
Told by an idiot, full of sound and fury
Signifying nothing.
Ay me! For aught that I could ever read,
Could ever hear by tale or history,
The course of true love never did run smooth.
I have full cause of weeping, but this heart
Shall break into a hundred thousand flaws,
Or ere I'll weep.-O Fool, I shall go mad!
                    Each your doing,
So singular in each particular,
Crowns what you are doing in the present deed,
That all your acts are queens.
"#;

    fn assert_reader<R: io::Read>(r: R) -> R {
        r
    }

    #[async_std::test]
    async fn test_embrio_stream() {
        let reader = assert_reader(SHAKESPEARE.as_bytes());
        let mut framed = Framed::new(reader, LineCodec, vec![0u8; 1024]);
        let expected = SHAKESPEARE.lines().map(String::from).collect::<Vec<_>>();
        let mut actual = vec![];
        while let Some(frame) = framed.next().await.transpose().unwrap() {
            actual.push(frame);
        }
        assert_eq!(actual, expected);
    }
    #[async_std::test]
    async fn test_embrio_sink() {
        let frames = SHAKESPEARE.lines().map(String::from).collect::<Vec<_>>();
        let mut actual = vec![0u8; SHAKESPEARE.as_bytes().len()];
        {
            let writer = io::Cursor::new(&mut actual);
            let mut framed = Framed::new(writer, LineCodec, vec![0u8; 1024]);
            for frame in frames {
                framed.send(frame).await.unwrap();
            }
        }
        assert_eq!(std::str::from_utf8(&actual).unwrap(), SHAKESPEARE);
    }
}