Skip to main content

compio_runtime/future/combinator/
personality.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use futures_util::Stream;
7use pin_project_lite::pin_project;
8
9use crate::{
10    future::Ext,
11    waker::{ExtWaker, with_ext},
12};
13
14pin_project! {
15    /// A future with a personality attached to it.
16    pub struct WithPersonality<F: ?Sized> {
17        personality: u16,
18        #[pin]
19        future: F,
20    }
21}
22
23impl<F> WithPersonality<F> {
24    /// Create a new [`WithPersonality`] future.
25    pub fn new(future: F, personality: u16) -> Self {
26        Self {
27            future,
28            personality,
29        }
30    }
31}
32
33impl<F: Future + ?Sized> Future for WithPersonality<F> {
34    type Output = F::Output;
35
36    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
37        let this = self.project();
38
39        with_ext(cx.waker(), |waker, ext: &Ext| {
40            let ext = ext.with_personality(*this.personality);
41            ExtWaker::new(waker, &ext).poll(this.future)
42        })
43    }
44}
45
46impl<F: Stream + ?Sized> Stream for WithPersonality<F> {
47    type Item = F::Item;
48
49    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
50        let this = self.project();
51
52        with_ext(cx.waker(), |waker, ext: &Ext| {
53            let ext = ext.with_personality(*this.personality);
54            ExtWaker::new(waker, &ext).poll_next(this.future)
55        })
56    }
57}