1 | use core::pin::Pin; |
2 | use core::task::{Context, Poll}; |
3 | use futures_core::future::{FusedFuture, Future}; |
4 | use futures_core::stream::{FusedStream, Stream}; |
5 | #[cfg (feature = "sink" )] |
6 | use futures_sink::Sink; |
7 | |
8 | /// Combines two different futures, streams, or sinks having the same associated types into a single type. |
9 | /// |
10 | /// This is useful when conditionally choosing between two distinct future types: |
11 | /// |
12 | /// ```rust |
13 | /// use futures::future::Either; |
14 | /// |
15 | /// # futures::executor::block_on(async { |
16 | /// let cond = true; |
17 | /// |
18 | /// let fut = if cond { |
19 | /// Either::Left(async move { 12 }) |
20 | /// } else { |
21 | /// Either::Right(async move { 44 }) |
22 | /// }; |
23 | /// |
24 | /// assert_eq!(fut.await, 12); |
25 | /// # }) |
26 | /// ``` |
27 | #[derive (Debug, Clone)] |
28 | pub enum Either<A, B> { |
29 | /// First branch of the type |
30 | Left(/* #[pin] */ A), |
31 | /// Second branch of the type |
32 | Right(/* #[pin] */ B), |
33 | } |
34 | |
35 | impl<A, B> Either<A, B> { |
36 | /// Convert `Pin<&Either<A, B>>` to `Either<Pin<&A>, Pin<&B>>`, |
37 | /// pinned projections of the inner variants. |
38 | pub fn as_pin_ref(self: Pin<&Self>) -> Either<Pin<&A>, Pin<&B>> { |
39 | // SAFETY: We can use `new_unchecked` because the `inner` parts are |
40 | // guaranteed to be pinned, as they come from `self` which is pinned. |
41 | unsafe { |
42 | match *Pin::get_ref(self) { |
43 | Either::Left(ref inner) => Either::Left(Pin::new_unchecked(inner)), |
44 | Either::Right(ref inner) => Either::Right(Pin::new_unchecked(inner)), |
45 | } |
46 | } |
47 | } |
48 | |
49 | /// Convert `Pin<&mut Either<A, B>>` to `Either<Pin<&mut A>, Pin<&mut B>>`, |
50 | /// pinned projections of the inner variants. |
51 | pub fn as_pin_mut(self: Pin<&mut Self>) -> Either<Pin<&mut A>, Pin<&mut B>> { |
52 | // SAFETY: `get_unchecked_mut` is fine because we don't move anything. |
53 | // We can use `new_unchecked` because the `inner` parts are guaranteed |
54 | // to be pinned, as they come from `self` which is pinned, and we never |
55 | // offer an unpinned `&mut A` or `&mut B` through `Pin<&mut Self>`. We |
56 | // also don't have an implementation of `Drop`, nor manual `Unpin`. |
57 | unsafe { |
58 | match *Pin::get_unchecked_mut(self) { |
59 | Either::Left(ref mut inner) => Either::Left(Pin::new_unchecked(inner)), |
60 | Either::Right(ref mut inner) => Either::Right(Pin::new_unchecked(inner)), |
61 | } |
62 | } |
63 | } |
64 | } |
65 | |
66 | impl<A, B, T> Either<(T, A), (T, B)> { |
67 | /// Factor out a homogeneous type from an either of pairs. |
68 | /// |
69 | /// Here, the homogeneous type is the first element of the pairs. |
70 | pub fn factor_first(self) -> (T, Either<A, B>) { |
71 | match self { |
72 | Either::Left((x: T, a: A)) => (x, Either::Left(a)), |
73 | Either::Right((x: T, b: B)) => (x, Either::Right(b)), |
74 | } |
75 | } |
76 | } |
77 | |
78 | impl<A, B, T> Either<(A, T), (B, T)> { |
79 | /// Factor out a homogeneous type from an either of pairs. |
80 | /// |
81 | /// Here, the homogeneous type is the second element of the pairs. |
82 | pub fn factor_second(self) -> (Either<A, B>, T) { |
83 | match self { |
84 | Either::Left((a: A, x: T)) => (Either::Left(a), x), |
85 | Either::Right((b: B, x: T)) => (Either::Right(b), x), |
86 | } |
87 | } |
88 | } |
89 | |
90 | impl<T> Either<T, T> { |
91 | /// Extract the value of an either over two equivalent types. |
92 | pub fn into_inner(self) -> T { |
93 | match self { |
94 | Either::Left(x: T) => x, |
95 | Either::Right(x: T) => x, |
96 | } |
97 | } |
98 | } |
99 | |
100 | impl<A, B> Future for Either<A, B> |
101 | where |
102 | A: Future, |
103 | B: Future<Output = A::Output>, |
104 | { |
105 | type Output = A::Output; |
106 | |
107 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
108 | match self.as_pin_mut() { |
109 | Either::Left(x: Pin<&mut A>) => x.poll(cx), |
110 | Either::Right(x: Pin<&mut B>) => x.poll(cx), |
111 | } |
112 | } |
113 | } |
114 | |
115 | impl<A, B> FusedFuture for Either<A, B> |
116 | where |
117 | A: FusedFuture, |
118 | B: FusedFuture<Output = A::Output>, |
119 | { |
120 | fn is_terminated(&self) -> bool { |
121 | match self { |
122 | Either::Left(x: &A) => x.is_terminated(), |
123 | Either::Right(x: &B) => x.is_terminated(), |
124 | } |
125 | } |
126 | } |
127 | |
128 | impl<A, B> Stream for Either<A, B> |
129 | where |
130 | A: Stream, |
131 | B: Stream<Item = A::Item>, |
132 | { |
133 | type Item = A::Item; |
134 | |
135 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { |
136 | match self.as_pin_mut() { |
137 | Either::Left(x: Pin<&mut A>) => x.poll_next(cx), |
138 | Either::Right(x: Pin<&mut B>) => x.poll_next(cx), |
139 | } |
140 | } |
141 | |
142 | fn size_hint(&self) -> (usize, Option<usize>) { |
143 | match self { |
144 | Either::Left(x: &A) => x.size_hint(), |
145 | Either::Right(x: &B) => x.size_hint(), |
146 | } |
147 | } |
148 | } |
149 | |
150 | impl<A, B> FusedStream for Either<A, B> |
151 | where |
152 | A: FusedStream, |
153 | B: FusedStream<Item = A::Item>, |
154 | { |
155 | fn is_terminated(&self) -> bool { |
156 | match self { |
157 | Either::Left(x: &A) => x.is_terminated(), |
158 | Either::Right(x: &B) => x.is_terminated(), |
159 | } |
160 | } |
161 | } |
162 | |
163 | #[cfg (feature = "sink" )] |
164 | impl<A, B, Item> Sink<Item> for Either<A, B> |
165 | where |
166 | A: Sink<Item>, |
167 | B: Sink<Item, Error = A::Error>, |
168 | { |
169 | type Error = A::Error; |
170 | |
171 | fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { |
172 | match self.as_pin_mut() { |
173 | Either::Left(x) => x.poll_ready(cx), |
174 | Either::Right(x) => x.poll_ready(cx), |
175 | } |
176 | } |
177 | |
178 | fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> { |
179 | match self.as_pin_mut() { |
180 | Either::Left(x) => x.start_send(item), |
181 | Either::Right(x) => x.start_send(item), |
182 | } |
183 | } |
184 | |
185 | fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { |
186 | match self.as_pin_mut() { |
187 | Either::Left(x) => x.poll_flush(cx), |
188 | Either::Right(x) => x.poll_flush(cx), |
189 | } |
190 | } |
191 | |
192 | fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { |
193 | match self.as_pin_mut() { |
194 | Either::Left(x) => x.poll_close(cx), |
195 | Either::Right(x) => x.poll_close(cx), |
196 | } |
197 | } |
198 | } |
199 | |
200 | #[cfg (feature = "io" )] |
201 | #[cfg (feature = "std" )] |
202 | mod if_std { |
203 | use super::*; |
204 | |
205 | use core::pin::Pin; |
206 | use core::task::{Context, Poll}; |
207 | use futures_io::{ |
208 | AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, IoSlice, IoSliceMut, Result, SeekFrom, |
209 | }; |
210 | |
211 | impl<A, B> AsyncRead for Either<A, B> |
212 | where |
213 | A: AsyncRead, |
214 | B: AsyncRead, |
215 | { |
216 | fn poll_read( |
217 | self: Pin<&mut Self>, |
218 | cx: &mut Context<'_>, |
219 | buf: &mut [u8], |
220 | ) -> Poll<Result<usize>> { |
221 | match self.as_pin_mut() { |
222 | Either::Left(x) => x.poll_read(cx, buf), |
223 | Either::Right(x) => x.poll_read(cx, buf), |
224 | } |
225 | } |
226 | |
227 | fn poll_read_vectored( |
228 | self: Pin<&mut Self>, |
229 | cx: &mut Context<'_>, |
230 | bufs: &mut [IoSliceMut<'_>], |
231 | ) -> Poll<Result<usize>> { |
232 | match self.as_pin_mut() { |
233 | Either::Left(x) => x.poll_read_vectored(cx, bufs), |
234 | Either::Right(x) => x.poll_read_vectored(cx, bufs), |
235 | } |
236 | } |
237 | } |
238 | |
239 | impl<A, B> AsyncWrite for Either<A, B> |
240 | where |
241 | A: AsyncWrite, |
242 | B: AsyncWrite, |
243 | { |
244 | fn poll_write( |
245 | self: Pin<&mut Self>, |
246 | cx: &mut Context<'_>, |
247 | buf: &[u8], |
248 | ) -> Poll<Result<usize>> { |
249 | match self.as_pin_mut() { |
250 | Either::Left(x) => x.poll_write(cx, buf), |
251 | Either::Right(x) => x.poll_write(cx, buf), |
252 | } |
253 | } |
254 | |
255 | fn poll_write_vectored( |
256 | self: Pin<&mut Self>, |
257 | cx: &mut Context<'_>, |
258 | bufs: &[IoSlice<'_>], |
259 | ) -> Poll<Result<usize>> { |
260 | match self.as_pin_mut() { |
261 | Either::Left(x) => x.poll_write_vectored(cx, bufs), |
262 | Either::Right(x) => x.poll_write_vectored(cx, bufs), |
263 | } |
264 | } |
265 | |
266 | fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> { |
267 | match self.as_pin_mut() { |
268 | Either::Left(x) => x.poll_flush(cx), |
269 | Either::Right(x) => x.poll_flush(cx), |
270 | } |
271 | } |
272 | |
273 | fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> { |
274 | match self.as_pin_mut() { |
275 | Either::Left(x) => x.poll_close(cx), |
276 | Either::Right(x) => x.poll_close(cx), |
277 | } |
278 | } |
279 | } |
280 | |
281 | impl<A, B> AsyncSeek for Either<A, B> |
282 | where |
283 | A: AsyncSeek, |
284 | B: AsyncSeek, |
285 | { |
286 | fn poll_seek( |
287 | self: Pin<&mut Self>, |
288 | cx: &mut Context<'_>, |
289 | pos: SeekFrom, |
290 | ) -> Poll<Result<u64>> { |
291 | match self.as_pin_mut() { |
292 | Either::Left(x) => x.poll_seek(cx, pos), |
293 | Either::Right(x) => x.poll_seek(cx, pos), |
294 | } |
295 | } |
296 | } |
297 | |
298 | impl<A, B> AsyncBufRead for Either<A, B> |
299 | where |
300 | A: AsyncBufRead, |
301 | B: AsyncBufRead, |
302 | { |
303 | fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<&[u8]>> { |
304 | match self.as_pin_mut() { |
305 | Either::Left(x) => x.poll_fill_buf(cx), |
306 | Either::Right(x) => x.poll_fill_buf(cx), |
307 | } |
308 | } |
309 | |
310 | fn consume(self: Pin<&mut Self>, amt: usize) { |
311 | match self.as_pin_mut() { |
312 | Either::Left(x) => x.consume(amt), |
313 | Either::Right(x) => x.consume(amt), |
314 | } |
315 | } |
316 | } |
317 | } |
318 | |