Repository: tokio-rs/tokio
File: tokio-util/src/sync/mpsc.rs (Rust)

Code before:
pub struct PollSender<T> {
    sender: Option<Sender<T>>,
    state: State<T>,
    acquire: PollSenderFuture<T>,
}

// Creates a future for acquiring a permit from the underlying channel.  This is used to ensure
// there's capacity for a send to complete.
//
// By reusing the same async fn for both `Some` and `None`, we make sure every future passed to
// ReusableBoxFuture has the same underlying type, and hence the same size and alignment.
async fn make_acquire_future<T>(
    data: Option<Sender<T>>,
) -> Result<OwnedPermit<T>, PollSendError<T>> {
    match data {
        Some(sender) => sender
            .reserve_owned()
            .await
            .map_err(|_| PollSendError(None)),
        None => unreachable!("this future should not be pollable in this state"),
    }
}

#[derive(Debug)]
// TODO: This should be replace with a type_alias_impl_trait to eliminate `'static` and all the transmutes
struct PollSenderFuture<T>(ReusableBoxFuture<'static, Result<OwnedPermit<T>, PollSendError<T>>>);

impl<T> From<ReusableBoxFuture<'static, Result<OwnedPermit<T>, PollSendError<T>>>>
    for PollSenderFuture<T>
{
    fn from(v: ReusableBoxFuture<'static, Result<OwnedPermit<T>, PollSendError<T>>>) -> Self {
        Self(v)
    }
}

impl<T: Send> From<Option<Sender<T>>> for PollSenderFuture<T> {
    fn from(data: Option<Sender<T>>) -> Self {
        let v = ReusableBoxFuture::<'_, Result<OwnedPermit<T>, PollSendError<T>>>::new(
            make_acquire_future(data),
        );
        // This is safe because we're just transmuting the `'static` bound on the Box<dyn> to get around storage issue
        unsafe { mem::transmute(v) }
    }
}

impl<T: Send> PollSenderFuture<T> {
    fn get<'a>(&mut self) -> &mut ReusableBoxFuture<'a, Result<OwnedPermit<T>, PollSendError<T>>>
    where
        T: 'a,
    {
        // This is safe because we're just untransmuting the `'static` bound on the Box<dyn> to get around storage issue

Review comment (refactor):
You may be able to do this simpler by adding an `unsafe_set` method to `ReusableBoxFuture` that accepts any future, even if it isn't `Send` or `'static`.

Code after:
pub struct PollSender<T> {
    sender: Option<Sender<T>>,
    state: State<T>,
    acquire: PollSenderFuture<T>,
}

// Creates a future for acquiring a permit from the underlying channel.  This is used to ensure
// there's capacity for a send to complete.
//
// By reusing the same async fn for both `Some` and `None`, we make sure every future passed to
// ReusableBoxFuture has the same underlying type, and hence the same size and alignment.
async fn make_acquire_future<T>(
    data: Option<Sender<T>>,
) -> Result<OwnedPermit<T>, PollSendError<T>> {
    match data {
        Some(sender) => sender
            .reserve_owned()
            .await
            .map_err(|_| PollSendError(None)),
        None => unreachable!("this future should not be pollable in this state"),
    }
}

type InnerFuture<'a, T> = ReusableBoxFuture<'a, Result<OwnedPermit<T>, PollSendError<T>>>;

#[derive(Debug)]
// TODO: This should be replace with a type_alias_impl_trait to eliminate `'static` and all the transmutes
struct PollSenderFuture<T>(InnerFuture<'static, T>);

impl<T> PollSenderFuture<T> {
    /// Create with an empty inner future with no `Send` bound.
    fn empty() -> Self {
        // We don't use `make_acquire_future` here because our relaxed bounds on `T` are not
        // compatible with the transitive bounds required by `Sender<T>`.
        Self(ReusableBoxFuture::new(async { unreachable!() }))
    }
}

impl<T: Send> PollSenderFuture<T> {
    /// Create with an empty inner future.
    fn new() -> Self {
        let v = InnerFuture::new(make_acquire_future(None));
        // This is safe because `make_acquire_future(None)` is actually `'static`
        Self(unsafe { mem::transmute::<InnerFuture<'_, T>, InnerFuture<'static, T>>(v) })
    }

    /// Poll the inner future.
    fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<OwnedPermit<T>, PollSendError<T>>> {
        self.0.poll(cx)
    }