1use std::future::Future;
30use std::pin::Pin;
31use std::task::{Context, Poll};
32
33pub fn select_two<F1, F2, T1, T2>(future1: F1, future2: F2) -> SelectTwo<F1, F2>
53where
54 F1: Future<Output = T1> + Unpin,
55 F2: Future<Output = T2> + Unpin,
56{
57 SelectTwo {
58 future1: Some(future1),
59 future2: Some(future2),
60 }
61}
62
63pub struct SelectTwo<F1, F2> {
66 future1: Option<F1>,
67 future2: Option<F2>,
68}
69
70impl<F1, F2, T1, T2> Future for SelectTwo<F1, F2>
71where
72 F1: Future<Output = T1> + Unpin,
73 F2: Future<Output = T2> + Unpin,
74{
75 type Output = SelectTwoOutput<T1, T2>;
76
77 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
78 if let Some(future1) = &mut self.future1 {
81 match Pin::new(future1).poll(cx) {
82 Poll::Ready(value) => {
83 self.future1 = None;
84 return Poll::Ready(SelectTwoOutput::First(value));
85 },
86 Poll::Pending => {},
87 }
88 }
89
90 if let Some(future2) = &mut self.future2 {
93 match Pin::new(future2).poll(cx) {
94 Poll::Ready(value) => {
95 self.future2 = None;
96 return Poll::Ready(SelectTwoOutput::Second(value));
97 },
98 Poll::Pending => {},
99 }
100 }
101
102 Poll::Pending
103 }
104}
105
106#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub enum SelectTwoOutput<T1, T2> {
110 First(T1),
113 Second(T2),
116}
117
118pub fn select_multiple<F>(futures: Vec<F>) -> SelectMultiple<F>
124where
125 F: Future + Unpin,
126{
127 SelectMultiple { futures }
128}
129
130pub struct SelectMultiple<F> {
133 futures: Vec<F>,
134}
135
136impl<F> Future for SelectMultiple<F>
137where
138 F: Future + Unpin,
139{
140 type Output = SelectMultipleOutput<F::Output>;
141
142 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
143 for (index, future) in self.futures.iter_mut().enumerate() {
146 match Pin::new(future).poll(cx) {
147 Poll::Ready(value) => {
148 return Poll::Ready(SelectMultipleOutput::Completed(index, value));
149 },
150 Poll::Pending => {},
151 }
152 }
153
154 Poll::Pending
155 }
156}
157
158#[derive(Debug, Clone, Copy, PartialEq, Eq)]
161pub enum SelectMultipleOutput<T> {
162 Completed(usize, T),
165}
166
167pub struct SelectMultipleBinding<T> {
170 pub index: usize,
173
174 pub output: Option<T>,
177}
178
179#[cfg(test)]
180mod tests {
181 use super::*;
182 use std::task::{Context, Waker};
183
184 #[test]
185 fn test_select_two_output_first() {
186 let output: SelectTwoOutput<i32, i32> = SelectTwoOutput::First(42);
187 assert!(matches!(output, SelectTwoOutput::First(42)));
188 }
189
190 #[test]
191 fn test_select_two_output_second() {
192 let output: SelectTwoOutput<i32, i32> = SelectTwoOutput::Second(100);
193 assert!(matches!(output, SelectTwoOutput::Second(100)));
194 }
195
196 #[test]
197 fn test_select_multiple_output() {
198 let output = SelectMultipleOutput::Completed(0, 42);
199 assert!(matches!(output, SelectMultipleOutput::Completed(0, 42)));
200 }
201
202 #[test]
203 fn test_ready_future() {
204 struct Ready;
207
208 impl Future for Ready {
209 type Output = i32;
210
211 fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
212 Poll::Ready(42)
213 }
214 }
215
216 let noop_waker = Waker::noop();
217 let mut context = Context::from_waker(&noop_waker);
218
219 let mut future = Box::pin(Ready);
220 assert!(matches!(Pin::new(&mut future).poll(&mut context), Poll::Ready(42)));
221 }
222}