navy-nvim-rs 0.0.18

//! Encoding msgpack-rpc messages to Neovim.
use rmp::encode::{
    write_array_len, write_bool, write_nil, write_sint, write_str, write_u32, write_uint8,
};
use rmpv::{ValueRef, encode::write_value_ref};
use std::io;
use tokio::io::{AsyncWrite, AsyncWriteExt};

const MSG_TYPE_REQUEST: u8 = 0;
const MSG_TYPE_RESPONSE: u8 = 1;
const MSG_TYPE_NOTIFICATION: u8 = 2;

/// State reused while encoding msgpack-rpc messages to a stream.
pub struct Encoder<W> {
    writer: W,
    buffer: Vec<u8>,
}

impl<W> Encoder<W>
where
    W: AsyncWrite + Send + Unpin + 'static,
{
    #[inline]
    #[must_use]
    pub fn new(writer: W) -> Self {
        Self {
            writer,
            buffer: Vec::new(),
        }
    }

    #[cfg(test)]
    #[inline]
    #[must_use]
    pub(crate) fn writer(&self) -> &W {
        &self.writer
    }

    pub async fn encode_request<A: EncodeArgs>(
        &mut self,
        msgid: u32,
        method: &str,
        args: A,
    ) -> io::Result<()> {
        self.buffer.clear();
        write_array_len(&mut self.buffer, 4)?;
        write_uint8(&mut self.buffer, MSG_TYPE_REQUEST)?;
        write_u32(&mut self.buffer, msgid)?;
        write_str(&mut self.buffer, method)?;
        write_array_len(&mut self.buffer, A::NUM_ARGS)?;
        args.encode_args(&mut self.buffer)?;

        self.writer.write_all(&self.buffer).await?;
        self.writer.flush().await?;
        Ok(())
    }

    pub async fn encode_notify<A: EncodeArgs>(&mut self, method: &str, args: A) -> io::Result<()> {
        self.buffer.clear();
        write_array_len(&mut self.buffer, 3)?;
        write_uint8(&mut self.buffer, MSG_TYPE_NOTIFICATION)?;
        write_str(&mut self.buffer, method)?;
        write_array_len(&mut self.buffer, A::NUM_ARGS)?;
        args.encode_args(&mut self.buffer)?;

        self.writer.write_all(&self.buffer).await?;
        self.writer.flush().await?;
        Ok(())
    }

    /// Encode an error response with a string error and nil result.
    pub async fn encode_error_response<E: Encode>(
        &mut self,
        msgid: u32,
        error: E,
    ) -> io::Result<()> {
        self.buffer.clear();
        write_array_len(&mut self.buffer, 4)?;
        write_uint8(&mut self.buffer, MSG_TYPE_RESPONSE)?;
        write_u32(&mut self.buffer, msgid)?;
        error.encode(&mut self.buffer)?;
        write_nil(&mut self.buffer)?;

        self.writer.write_all(&self.buffer).await?;
        self.writer.flush().await?;
        Ok(())
    }
}

pub trait Encode: Copy {
    fn encode(self, buf: &mut Vec<u8>) -> io::Result<()>;
}
impl Encode for &ValueRef<'_> {
    #[inline]
    fn encode(self, buf: &mut Vec<u8>) -> io::Result<()> {
        Ok(write_value_ref(buf, self)?)
    }
}
impl Encode for &str {
    #[inline]
    fn encode(self, buf: &mut Vec<u8>) -> io::Result<()> {
        Ok(write_str(buf, self)?)
    }
}
impl Encode for bool {
    #[inline]
    fn encode(self, buf: &mut Vec<u8>) -> io::Result<()> {
        write_bool(buf, self)
    }
}
impl Encode for i64 {
    #[inline]
    fn encode(self, buf: &mut Vec<u8>) -> io::Result<()> {
        write_sint(buf, self)?;
        Ok(())
    }
}

pub trait EncodeArgs {
    const NUM_ARGS: u32;
    fn encode_args(self, buf: &mut Vec<u8>) -> io::Result<()>;
}
impl<T: Encode> EncodeArgs for T {
    const NUM_ARGS: u32 = 1;
    #[inline]
    fn encode_args(self, buf: &mut Vec<u8>) -> io::Result<()> {
        self.encode(buf)
    }
}
macro_rules! impl_encode_tuple_args {
    ($len:expr, $($arg:ident),*) => {
        impl<$($arg: Encode,)*> EncodeArgs for ($($arg,)*) {
            const NUM_ARGS: u32 = $len;
            #[inline]
            fn encode_args(self, _buf: &mut Vec<u8>) -> io::Result<()> {
                #[allow(non_snake_case)]
                let ($($arg,)*) = self;
                $(
                    $arg.encode(_buf)?;
                )*
                Ok(())
            }
        }
    };
}
impl_encode_tuple_args!(0,);
impl_encode_tuple_args!(2, T, U);
impl_encode_tuple_args!(3, T, U, V);
impl_encode_tuple_args!(4, T, U, V, W);

#[cfg(test)]
mod tests {
    use super::*;
    use rmpv::{Value, decode::read_value};
    use tokio::io::Sink;

    fn encoder() -> Encoder<Sink> {
        Encoder::new(tokio::io::sink())
    }

    fn decoded_value(bytes: &[u8]) -> Value {
        let mut input = bytes;
        let value = read_value(&mut input).unwrap();
        assert!(input.is_empty());
        value
    }

    #[tokio::test]
    async fn encode_single_arg_matches_request_encoding() {
        let mut encoder = encoder();
        encoder
            .encode_request(7, "nvim_input", "<C-D>")
            .await
            .unwrap();

        assert_eq!(
            decoded_value(&encoder.buffer),
            Value::from(vec![
                Value::from(0),
                Value::from(7),
                Value::from("nvim_input"),
                Value::from(vec![Value::from("<C-D>")]),
            ])
        );
    }

    #[tokio::test]
    async fn encode_single_arg_matches_notification_encoding() {
        let mut encoder = encoder();
        encoder.encode_notify("nvim_input", "<C-D>").await.unwrap();

        assert_eq!(
            decoded_value(&encoder.buffer),
            Value::from(vec![
                Value::from(2),
                Value::from("nvim_input"),
                Value::from(vec![Value::from("<C-D>")]),
            ])
        );
    }

    #[tokio::test]
    async fn encode_message_matches_value_ref_request_encoding() {
        let cmd = ValueRef::Map(vec![
            (ValueRef::from("cmd"), ValueRef::from("echo")),
            (
                ValueRef::from("args"),
                ValueRef::Array(vec![ValueRef::from("hello")]),
            ),
        ]);
        let opts = ValueRef::Map(vec![(ValueRef::from("output"), ValueRef::Boolean(true))]);

        let mut encoder = encoder();
        encoder
            .encode_request(7, "nvim_cmd", (&cmd, &opts))
            .await
            .unwrap();

        let params = Value::from(vec![cmd.to_owned(), opts.to_owned()]);

        assert_eq!(
            decoded_value(&encoder.buffer),
            Value::from(vec![
                Value::from(0),
                Value::from(7),
                Value::from("nvim_cmd"),
                params,
            ])
        );
    }

    #[tokio::test]
    async fn encode_message_matches_integer_notification_encoding() {
        let mut encoder = encoder();
        encoder
            .encode_notify("nvim_ui_try_resize", (120_i64, 40_i64))
            .await
            .unwrap();

        assert_eq!(
            decoded_value(&encoder.buffer),
            Value::from(vec![
                Value::from(2),
                Value::from("nvim_ui_try_resize"),
                Value::from(vec![Value::from(120), Value::from(40)]),
            ])
        );
    }

    #[tokio::test]
    async fn encode_error_response_matches_response_encoding() {
        let mut encoder = encoder();
        encoder
            .encode_error_response(9, "Not implemented")
            .await
            .unwrap();

        assert_eq!(
            decoded_value(&encoder.buffer),
            Value::from(vec![
                Value::from(1),
                Value::from(9),
                Value::from("Not implemented"),
                Value::Nil,
            ])
        );
    }
}