Skip to main content

compact_waitgroup/
with_worker_handle.rs

1use core::{
2    ops::{Deref, DerefMut},
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> DerefMut for WithWorkerHandleFuture<F, H> {
30    fn deref_mut(&mut self) -> &mut Self::Target {
31        &mut self.inner
32    }
33}
34
35pub trait WithWorkerHandle<H>: Sized {
36    fn with_worker_handle(self, handle: H) -> WithWorkerHandleFuture<Self, H>;
37}
38
39impl<F: Future, H: WorkerHandleType> WithWorkerHandle<H> for F {
40    fn with_worker_handle(self, handle: H) -> WithWorkerHandleFuture<Self, H> {
41        WithWorkerHandleFuture {
42            inner: self,
43            worker_handle: handle,
44        }
45    }
46}
47
48impl<F: Future, H: WorkerHandleType> Future for WithWorkerHandleFuture<F, H> {
49    type Output = F::Output;
50
51    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
52        self.project().inner.poll(cx)
53    }
54}
55
56trait WorkerHandleType {}
57
58impl WorkerHandleType for WorkerHandle {}
59impl WorkerHandleType for MonoWorkerHandle {}