Skip to main content

compact_waitgroup/
with_worker_handle.rs

1use core::{
2    ops::Deref,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use derive_more::Into;
8use pin_project_lite::pin_project;
9
10use crate::{MonoWorkerHandle, WorkerHandle};
11
12pin_project! {
13    #[derive(Debug, Into)]
14    pub struct WithWorkerHandleFuture<F, H> {
15        #[pin]
16        inner: F,
17        worker_handle: H,
18    }
19}
20
21impl<F, H> Deref for WithWorkerHandleFuture<F, H> {
22    type Target = F;
23
24    fn deref(&self) -> &Self::Target {
25        &self.inner
26    }
27}
28
29impl<F, H> WithWorkerHandleFuture<F, H> {
30    pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
31        self.project().inner
32    }
33
34    pub fn worker_handle(&self) -> &H {
35        &self.worker_handle
36    }
37}
38
39pub trait WithWorkerHandle<H>: Sized {
40    fn with_worker_handle(self, handle: H) -> WithWorkerHandleFuture<Self, H>;
41}
42
43impl<F: Future, H: WorkerHandleType> WithWorkerHandle<H> for F {
44    fn with_worker_handle(self, handle: H) -> WithWorkerHandleFuture<Self, H> {
45        WithWorkerHandleFuture {
46            inner: self,
47            worker_handle: handle,
48        }
49    }
50}
51
52impl<F: Future, H: WorkerHandleType> Future for WithWorkerHandleFuture<F, H> {
53    type Output = F::Output;
54
55    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
56        self.project().inner.poll(cx)
57    }
58}
59
60trait WorkerHandleType {}
61
62impl WorkerHandleType for WorkerHandle {}
63impl WorkerHandleType for MonoWorkerHandle {}