warp/filter/
or_else.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4
5use futures_util::{ready, TryFuture};
6use pin_project::pin_project;
7
8use super::{Filter, FilterBase, Func, Internal};
9use crate::reject::IsReject;
10use crate::route;
11
12#[derive(Clone, Copy, Debug)]
13pub struct OrElse<T, F> {
14    pub(super) filter: T,
15    pub(super) callback: F,
16}
17
18impl<T, F> FilterBase for OrElse<T, F>
19where
20    T: Filter,
21    F: Func<T::Error> + Clone + Send,
22    F::Output: TryFuture<Ok = T::Extract> + Send,
23    <F::Output as TryFuture>::Error: IsReject,
24{
25    type Extract = <F::Output as TryFuture>::Ok;
26    type Error = <F::Output as TryFuture>::Error;
27    type Future = OrElseFuture<T, F>;
28    #[inline]
29    fn filter(&self, _: Internal) -> Self::Future {
30        let idx = route::with(|route| route.matched_path_index());
31        OrElseFuture {
32            state: State::First(self.filter.filter(Internal), self.callback.clone()),
33            original_path_index: PathIndex(idx),
34        }
35    }
36}
37
38#[allow(missing_debug_implementations)]
39#[pin_project]
40pub struct OrElseFuture<T: Filter, F>
41where
42    T: Filter,
43    F: Func<T::Error>,
44    F::Output: TryFuture<Ok = T::Extract> + Send,
45{
46    #[pin]
47    state: State<T, F>,
48    original_path_index: PathIndex,
49}
50
51#[pin_project(project = StateProj)]
52enum State<T, F>
53where
54    T: Filter,
55    F: Func<T::Error>,
56    F::Output: TryFuture<Ok = T::Extract> + Send,
57{
58    First(#[pin] T::Future, F),
59    Second(#[pin] F::Output),
60    Done,
61}
62
63#[derive(Copy, Clone)]
64struct PathIndex(usize);
65
66impl PathIndex {
67    fn reset_path(&self) {
68        route::with(|route| route.reset_matched_path_index(self.0));
69    }
70}
71
72impl<T, F> Future for OrElseFuture<T, F>
73where
74    T: Filter,
75    F: Func<T::Error>,
76    F::Output: TryFuture<Ok = T::Extract> + Send,
77{
78    type Output = Result<<F::Output as TryFuture>::Ok, <F::Output as TryFuture>::Error>;
79
80    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
81        loop {
82            let pin = self.as_mut().project();
83            let (err, second) = match pin.state.project() {
84                StateProj::First(first, second) => match ready!(first.try_poll(cx)) {
85                    Ok(ex) => return Poll::Ready(Ok(ex)),
86                    Err(err) => (err, second),
87                },
88                StateProj::Second(second) => {
89                    let ex2 = ready!(second.try_poll(cx));
90                    self.set(OrElseFuture {
91                        state: State::Done,
92                        ..*self
93                    });
94                    return Poll::Ready(ex2);
95                }
96                StateProj::Done => panic!("polled after complete"),
97            };
98
99            pin.original_path_index.reset_path();
100            let fut2 = second.call(err);
101            self.set(OrElseFuture {
102                state: State::Second(fut2),
103                ..*self
104            });
105        }
106    }
107}