Skip to main content

ps_promise/methods/
poll.rs

1use std::task::{Context, Poll};
2
3use crate::{Promise, PromiseRejection};
4
5impl<T, E> Promise<T, E>
6where
7    T: Unpin,
8    E: PromiseRejection,
9{
10    /// Attempts to advance this [`Promise`] using the provided execution [`Context`].
11    ///
12    /// This performs exactly one poll of the underlying future.
13    pub fn poll(&mut self, cx: &mut Context<'_>) {
14        let Self::Pending(future) = self else {
15            return;
16        };
17
18        match future.as_mut().poll(cx) {
19            Poll::Ready(Ok(value)) => *self = Self::Resolved(value),
20            Poll::Ready(Err(err)) => *self = Self::Rejected(err),
21            Poll::Pending => {}
22        }
23    }
24}
25
26#[cfg(test)]
27mod tests {
28    use std::{
29        future::Future,
30        pin::Pin,
31        sync::{
32            atomic::{AtomicBool, AtomicUsize, Ordering},
33            Arc,
34        },
35        task::{Context, Poll, Waker},
36    };
37
38    use crate::{Promise, PromiseRejection};
39
40    #[derive(Debug, Clone, PartialEq, Eq)]
41    enum E {
42        AlreadyConsumed,
43        Fail,
44    }
45
46    impl PromiseRejection for E {
47        fn already_consumed() -> Self {
48            Self::AlreadyConsumed
49        }
50    }
51
52    fn cx() -> Context<'static> {
53        Context::from_waker(Waker::noop())
54    }
55
56    #[test]
57    fn resolves_ready_future() {
58        let mut promise: Promise<i32, E> = Promise::new(async { Ok(42) });
59        promise.poll(&mut cx());
60        match promise {
61            Promise::Resolved(v) => assert_eq!(v, 42),
62            other => panic!("expected Resolved(42), got {other:?}"),
63        }
64    }
65
66    #[test]
67    fn rejects_ready_future() {
68        let mut promise: Promise<(), E> = Promise::new(async { Err(E::Fail) });
69        promise.poll(&mut cx());
70        match promise {
71            Promise::Rejected(E::Fail) => {}
72            other => panic!("expected Rejected(Fail), got {other:?}"),
73        }
74    }
75
76    #[test]
77    fn stays_pending() {
78        struct Never;
79
80        impl Future for Never {
81            type Output = Result<(), E>;
82            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
83                Poll::Pending
84            }
85        }
86
87        let mut promise: Promise<(), E> = Promise::new(Never);
88        promise.poll(&mut cx());
89        assert!(promise.is_pending());
90    }
91
92    #[test]
93    fn resolved_is_identity() {
94        let mut promise: Promise<i32, E> = Promise::resolve(99);
95        promise.poll(&mut cx());
96        match promise {
97            Promise::Resolved(v) => assert_eq!(v, 99),
98            other => panic!("expected Resolved(99), got {other:?}"),
99        }
100    }
101
102    #[test]
103    fn rejected_is_identity() {
104        let mut promise: Promise<(), E> = Promise::reject(E::Fail);
105        promise.poll(&mut cx());
106        match promise {
107            Promise::Rejected(E::Fail) => {}
108            other => panic!("expected Rejected(Fail), got {other:?}"),
109        }
110    }
111
112    #[test]
113    fn consumed_is_identity() {
114        let mut promise: Promise<(), E> = Promise::Consumed;
115        promise.poll(&mut cx());
116        assert!(promise.is_consumed());
117    }
118
119    #[test]
120    fn pending_then_resolves() {
121        struct Delayed {
122            ready: Arc<AtomicBool>,
123        }
124
125        impl Future for Delayed {
126            type Output = Result<i32, E>;
127            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
128                if self.ready.load(Ordering::Relaxed) {
129                    Poll::Ready(Ok(7))
130                } else {
131                    Poll::Pending
132                }
133            }
134        }
135
136        let ready = Arc::new(AtomicBool::new(false));
137
138        let mut promise: Promise<i32, E> = Promise::new(Delayed {
139            ready: ready.clone(),
140        });
141
142        promise.poll(&mut cx());
143        assert!(promise.is_pending());
144
145        ready.store(true, Ordering::Relaxed);
146
147        promise.poll(&mut cx());
148        match promise {
149            Promise::Resolved(v) => assert_eq!(v, 7),
150            other => panic!("expected Resolved(7), got {other:?}"),
151        }
152    }
153
154    #[test]
155    fn pending_then_rejects() {
156        struct Delayed {
157            ready: Arc<AtomicBool>,
158        }
159
160        impl Future for Delayed {
161            type Output = Result<(), E>;
162            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
163                if self.ready.load(Ordering::Relaxed) {
164                    Poll::Ready(Err(E::Fail))
165                } else {
166                    Poll::Pending
167                }
168            }
169        }
170
171        let ready = Arc::new(AtomicBool::new(false));
172
173        let mut promise: Promise<(), E> = Promise::new(Delayed {
174            ready: ready.clone(),
175        });
176
177        promise.poll(&mut cx());
178        assert!(promise.is_pending());
179
180        ready.store(true, Ordering::Relaxed);
181
182        promise.poll(&mut cx());
183        match promise {
184            Promise::Rejected(E::Fail) => {}
185            other => panic!("expected Rejected(Fail), got {other:?}"),
186        }
187    }
188
189    #[test]
190    fn waker_is_forwarded() {
191        struct StoreWaker {
192            woken: Arc<AtomicBool>,
193        }
194
195        impl Future for StoreWaker {
196            type Output = Result<(), E>;
197            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
198                let waker = cx.waker().clone();
199                let flag = self.woken.clone();
200                waker.wake();
201                flag.store(true, Ordering::Relaxed);
202                Poll::Pending
203            }
204        }
205
206        let woken = Arc::new(AtomicBool::new(false));
207
208        let mut promise: Promise<(), E> = Promise::new(StoreWaker {
209            woken: woken.clone(),
210        });
211        promise.poll(&mut cx());
212
213        assert!(
214            woken.load(Ordering::Relaxed),
215            "waker should have been invoked"
216        );
217    }
218
219    #[test]
220    fn polls_exactly_once() {
221        struct Counter {
222            count: Arc<AtomicUsize>,
223        }
224
225        impl Future for Counter {
226            type Output = Result<(), E>;
227            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
228                self.count.fetch_add(1, Ordering::Relaxed);
229                Poll::Pending
230            }
231        }
232
233        let count = Arc::new(AtomicUsize::new(0));
234
235        let mut promise: Promise<(), E> = Promise::new(Counter {
236            count: count.clone(),
237        });
238        promise.poll(&mut cx());
239
240        assert_eq!(count.load(Ordering::Relaxed), 1);
241    }
242}