navy-nvim-rs 0.0.24

use std::{
    io::{self, Error},
    marker::PhantomData,
    sync::atomic::{AtomicU32, Ordering},
};

use parking_lot::Mutex as SyncMutex;
use tokio::io::AsyncWrite;
use tokio::sync::{Mutex, oneshot};

use crate::{
    error::{MessageIdNotFound, NeovimError, ResponseReceiveError},
    rpc::{
        decode::Decode,
        encode::{Encode, EncodeArgs, Encoder},
        payload::ResponsePayload,
    },
};

type RequestCallback = oneshot::Sender<ResponsePayload>;

#[derive(Default)]
struct PendingRequests {
    // `Vec` is okay here because its size is 1 in most cases.
    callbacks: SyncMutex<Vec<(u32, RequestCallback)>>,
}

impl PendingRequests {
    #[inline]
    fn clear(&self) {
        self.callbacks.lock().clear();
    }

    #[inline]
    fn add(&self, id: u32, callback: RequestCallback) {
        self.callbacks.lock().push((id, callback));
    }

    #[inline]
    fn take(&self, msgid: u32) -> io::Result<RequestCallback> {
        let mut callbacks = self.callbacks.lock();
        if let Some(pos) = callbacks.iter().position(|(id, _)| *id == msgid) {
            Ok(callbacks.remove(pos).1)
        } else {
            Err(Error::other(MessageIdNotFound::new(msgid)))
        }
    }
}

/// Shared state for sending API calls and matching their responses.
pub(crate) struct ApiCaller<W> {
    encoder: Mutex<Encoder<W>>,
    pending: PendingRequests,
    next_msgid: AtomicU32,
}

impl<W> ApiCaller<W>
where
    W: AsyncWrite + Unpin,
{
    #[inline]
    pub(crate) fn new(writer: W) -> Self {
        Self {
            encoder: Mutex::new(Encoder::new(writer)),
            pending: PendingRequests::default(),
            next_msgid: AtomicU32::new(0),
        }
    }

    pub(crate) async fn request<A: EncodeArgs>(
        &self,
        method: &'static str,
        args: A,
    ) -> io::Result<ResponsePayload> {
        let msgid = self.next_msgid.fetch_add(1, Ordering::Relaxed);
        let (sender, receiver) = oneshot::channel();

        self.pending.add(msgid, sender);
        self.encoder
            .lock()
            .await
            .encode_request(msgid, method, args)
            .await?;

        match receiver.await {
            Ok(response) => Ok(response),
            Err(err) => Err(Error::other(ResponseReceiveError::new(method, err))),
        }
    }

    #[inline]
    pub(crate) async fn notify<A: EncodeArgs>(
        &self,
        method: &'static str,
        args: A,
    ) -> io::Result<()> {
        self.encoder.lock().await.encode_notify(method, args).await
    }

    #[inline]
    pub(crate) fn call<'a, A>(&'a self, method: &'static str, args: A) -> ApiCall<'a, W, A>
    where
        A: EncodeArgs,
    {
        ApiCall {
            method,
            caller: self,
            args,
        }
    }

    #[inline]
    pub(crate) fn call_and_decode<'a, R, A>(
        &'a self,
        method: &'static str,
        args: A,
    ) -> ApiCallAndDecode<'a, W, R, A>
    where
        A: EncodeArgs,
        R: for<'de> Decode<'de>,
    {
        ApiCallAndDecode {
            method,
            caller: self,
            args,
            _return_type: PhantomData,
        }
    }

    #[inline]
    pub(crate) fn callback(&self, message_id: u32) -> io::Result<RequestCallback> {
        self.pending.take(message_id)
    }

    #[inline]
    pub(crate) async fn respond(
        &self,
        message_id: u32,
        response: Result<impl Encode, impl Encode>,
    ) -> io::Result<()> {
        let mut encoder = self.encoder.lock().await;
        match response {
            Ok(result) => encoder.encode_result_response(message_id, &result).await,
            Err(error) => encoder.encode_error_response(message_id, &error).await,
        }
    }

    #[inline]
    pub(crate) fn clear(&self) {
        self.pending.clear();
    }
}

pub struct ApiCall<'a, W, A> {
    method: &'static str,
    caller: &'a ApiCaller<W>,
    args: A,
}

impl<'a, W, A> ApiCall<'a, W, A>
where
    W: AsyncWrite + Unpin,
    A: EncodeArgs,
{
    #[inline]
    pub async fn request(self) -> io::Result<ResponsePayload> {
        self.caller.request(self.method, self.args).await
    }

    #[inline]
    pub async fn notify(self) -> io::Result<()> {
        self.caller.notify(self.method, self.args).await
    }
}

pub struct ApiCallAndDecode<'a, W, R, A> {
    method: &'static str,
    caller: &'a ApiCaller<W>,
    args: A,
    _return_type: PhantomData<fn() -> R>,
}

impl<'a, W, R, A> ApiCallAndDecode<'a, W, R, A>
where
    W: AsyncWrite + Unpin,
    A: EncodeArgs,
    R: for<'de> Decode<'de>,
{
    #[inline]
    pub async fn request(self) -> io::Result<R> {
        self.caller
            .request(self.method, self.args)
            .await?
            .decode()?
            .map_err(|error| Error::other(NeovimError::new(self.method, error)))
    }

    #[inline]
    pub async fn request_raw(self) -> io::Result<ResponsePayload> {
        self.caller.request(self.method, self.args).await
    }

    #[inline]
    pub async fn notify(self) -> io::Result<()> {
        self.caller.notify(self.method, self.args).await
    }
}

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

    #[test]
    fn pending_requests_takes_callback_by_message_id() {
        let pending = PendingRequests::default();
        let (callback1, _receiver1) = oneshot::channel();
        let (callback2, _receiver2) = oneshot::channel();
        let (callback3, _receiver3) = oneshot::channel();
        pending.add(1, callback1);
        pending.add(2, callback2);
        pending.add(3, callback3);

        pending.take(2).unwrap();
        let ids = pending
            .callbacks
            .lock()
            .iter()
            .map(|(id, _)| *id)
            .collect::<Vec<_>>();
        assert_eq!(ids, [1, 3]);
    }

    #[test]
    fn pending_requests_reports_unknown_message_id() {
        let err = PendingRequests::default().take(17).unwrap_err();
        let err = err
            .get_ref()
            .and_then(|err| err.downcast_ref::<MessageIdNotFound>())
            .expect("expected MessageIdNotFound");
        assert_eq!(err.msgid, 17);
    }

    #[tokio::test]
    async fn pending_requests_clear_closes_all_callbacks() {
        let pending = PendingRequests::default();
        let (callback1, receiver1) = oneshot::channel();
        let (callback2, receiver2) = oneshot::channel();
        pending.add(1, callback1);
        pending.add(2, callback2);

        pending.clear();

        assert!(receiver1.await.is_err());
        assert!(receiver2.await.is_err());
        assert!(pending.callbacks.lock().is_empty());
    }
}