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

use core::future::Future;
use core::ops::{Deref, DerefMut};

use tokio::io::{AsyncRead, AsyncWrite};

/// Accepts connections on a transport
pub trait Accept {
    /// Transport-specific invocation context
    type Context: Send + Sync + 'static;

    /// Outgoing byte stream
    type Outgoing: AsyncWrite + Send + Sync + Unpin + 'static;

    /// Incoming byte stream
    type Incoming: AsyncRead + Send + Sync + Unpin + 'static;

    /// Accept a connection returning a pair of streams and connection context
    fn accept(
        &self,
    ) -> impl Future<Output = std::io::Result<(Self::Context, Self::Outgoing, Self::Incoming)>>;
}

/// Wrapper returned by [`AcceptExt::map_context`]
pub struct AcceptMapContext<T, F> {
    inner: T,
    f: F,
}

impl<T, F> Deref for AcceptMapContext<T, F> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<T, F> DerefMut for AcceptMapContext<T, F> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

/// Extension trait for [Accept]
pub trait AcceptExt: Accept + Sized {
    /// Maps [`Self::Context`](Accept::Context) to a type `T` using `F`
    fn map_context<T, F: Fn(Self::Context) -> T>(self, f: F) -> AcceptMapContext<Self, F> {
        AcceptMapContext { inner: self, f }
    }
}

impl<T: Accept> AcceptExt for T {}

impl<T, U, F> Accept for AcceptMapContext<T, F>
where
    T: Accept,
    U: Send + Sync + 'static,
    F: Fn(T::Context) -> U,
{
    type Context = U;
    type Outgoing = T::Outgoing;
    type Incoming = T::Incoming;

    async fn accept(&self) -> std::io::Result<(Self::Context, Self::Outgoing, Self::Incoming)> {
        (&self).accept().await
    }
}

impl<T, U, F> Accept for &AcceptMapContext<T, F>
where
    T: Accept,
    U: Send + Sync + 'static,
    F: Fn(T::Context) -> U,
{
    type Context = U;
    type Outgoing = T::Outgoing;
    type Incoming = T::Incoming;

    async fn accept(&self) -> std::io::Result<(Self::Context, Self::Outgoing, Self::Incoming)> {
        let (cx, tx, rx) = self.inner.accept().await?;
        Ok(((self.f)(cx), tx, rx))
    }
}