1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169

cfg_async! {
use std::{
    io::{Result as IoResult, Write},
    pin::Pin,
    task::{Context, Poll},
};

use futures_lite::{AsyncWrite, AsyncWriteExt};

/// An asynchronous encoder.

pub struct AsyncEncoder<W> {
    pub(crate) writer: W,
    pos: usize,
    data: Vec<u8>,
}

impl<W> std::fmt::Debug for AsyncEncoder<W> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AsyncEncoder").finish()
    }
}

impl<W> Clone for AsyncEncoder<W>
where
    W: Clone,
{
    fn clone(&self) -> Self {
        Self {
            writer: self.writer.clone(),
            pos: 0,
            data: vec![],
        }
    }
}

impl<W> Write for AsyncEncoder<W>
where
    W: Write + Send + Sync,
{
    fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
        self.writer.write(buf)
    }

    fn flush(&mut self) -> IoResult<()> {
        self.writer.flush()
    }
}

impl<W> AsyncEncoder<W>
where
    W: Write + Send + Sync,
{
    /// If the wrapped writer is synchronous, you can use this method to encode the message to it.

    pub fn encode_sync<M>(&mut self, msg: M) -> IoResult<()>
    where
        M: crate::Encodable + Send + Sync,
    {
        msg.encode(&mut self.data)?;
        let data = &self.data[self.pos..];

        self.writer.write_all(data)?;
        self.writer.flush()?;

        self.data.clear();
        self.pos = 0;
        Ok(())
    }
}

impl<W> AsyncEncoder<W>
where
    W: AsyncWrite + Send + Sync + Unpin,
{
    /// Create a new Encoder over this [futures_lite::AsyncWrite] instance

    pub fn new(writer: W) -> Self {
        Self {
            writer,
            pos: 0,
            data: Vec::with_capacity(1024),
        }
    }

    /// Get the inner [futures_lite::AsyncWrite] instance out

    ///

    /// This writes and flushes any buffered data before it consumes self.

    pub async fn into_inner(mut self) -> IoResult<W> {
        if self.data.is_empty() {
            return Ok(self.writer);
        }

        let data = std::mem::take(&mut self.data);
        self.writer.write_all(&data).await?;
        self.writer.flush().await?;
        Ok(self.writer)
    }

    /// Encode this [Encodable](crate::Encodable) message to the writer.

    ///

    /// This flushes the data before returning

    pub async fn encode<M>(&mut self, msg: M) -> IoResult<()>
    where
        M: crate::Encodable + Send + Sync,
        W: Unpin,
    {
        msg.encode(&mut self.data)?;
        let data = &self.data[self.pos..];

        self.writer.write_all(data).await?;
        self.writer.flush().await?;

        self.data.clear();
        self.pos = 0;
        Ok(())
    }
}

impl<W> AsyncWrite for AsyncEncoder<W>
where
    W: AsyncWrite + Unpin + Send + Sync,
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        ctx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<IoResult<usize>> {
        let mut this = self.as_mut();
        let writer = &mut this.writer;
        futures_lite::pin!(writer);
        writer.poll_write(ctx, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<IoResult<()>> {
        let mut this = self.as_mut();
        let writer = &mut this.writer;
        futures_lite::pin!(writer);
        writer.poll_flush(ctx)
    }

    fn poll_close(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<IoResult<()>> {
        let mut this = self.as_mut();
        let writer = &mut this.writer;
        futures_lite::pin!(writer);
        writer.poll_close(ctx)
    }
}

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

    #[test]
    fn encoder_async() {
        let fut = async move {
            let mut output = vec![];
            {
                let mut encoder = AsyncEncoder::new(&mut output);

                encoder.encode(join("#museun")).await.unwrap();
                encoder.encode(join("#shaken_bot")).await.unwrap();
            }

            let s = std::str::from_utf8(&output).unwrap();
            assert_eq!(s, "JOIN #museun\r\nJOIN #shaken_bot\r\n");
        };
        futures_lite::future::block_on(fut);
    }
}
}