Implemented Filter::then (seanmonstar#878)

optimizasean · Sep 8, 2021 · a8dbdb6 · a8dbdb6
1 parent f490c69
commit a8dbdb6
Show file tree

Hide file tree

Showing 3 changed files with 133 additions and 4 deletions.
diff --git a/src/filter/and_then.rs b/src/filter/and_then.rs
@@ -34,7 +34,7 @@ where
 
 #[allow(missing_debug_implementations)]
 #[pin_project]
-pub struct AndThenFuture<T: Filter, F>
+pub struct AndThenFuture<T, F>
 where
     T: Filter,
     F: Func<T::Extract>,
@@ -96,12 +96,12 @@ where
                     self.set(State::Second(fut2));
                 }
                 StateProj::Second(second) => {
-                    let ex3 = match ready!(second.try_poll(cx)) {
+                    let ex2 = match ready!(second.try_poll(cx)) {
                         Ok(item) => Ok((item,)),
                         Err(err) => Err(From::from(err)),
                     };
                     self.set(State::Done);
-                    return Poll::Ready(ex3);
+                    return Poll::Ready(ex2);
                 }
                 StateProj::Done => panic!("polled after complete"),
             }

diff --git a/src/filter/mod.rs b/src/filter/mod.rs
@@ -7,6 +7,7 @@ mod or;
 mod or_else;
 mod recover;
 pub(crate) mod service;
+mod then;
 mod unify;
 mod untuple_one;
 mod wrap;
@@ -27,6 +28,7 @@ pub(crate) use self::map_err::MapErr;
 pub(crate) use self::or::Or;
 use self::or_else::OrElse;
 use self::recover::Recover;
+use self::then::Then;
 use self::unify::Unify;
 use self::untuple_one::UntupleOne;
 pub use self::wrap::wrap_fn;
@@ -197,13 +199,45 @@ pub trait Filter: FilterBase {
         }
     }
 
-    /// Composes this `Filter` with a function receiving the extracted value.
+    /// Composes this `Filter` with an async function receiving
+    /// the extracted value.
+    ///
+    /// The function should return some `Future` type.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use warp::Filter;
+    ///
+    /// // Map `/:id`
+    /// warp::path::param().then(|id: u64| async move {
+    ///   format!("Hello #{}", id)
+    /// });
+    /// ```
+    fn then<F>(self, fun: F) -> Then<Self, F>
+    where
+        Self: Sized,
+        F: Func<Self::Extract> + Clone,
+        F::Output: Future + Send,
+    {
+        Then {
+            filter: self,
+            callback: fun,
+        }
+    }
+
+    /// Composes this `Filter` with a fallible async function receiving
+    /// the extracted value.
     ///
     /// The function should return some `TryFuture` type.
     ///
     /// The `Error` type of the return `Future` needs be a `Rejection`, which
     /// means most futures will need to have their error mapped into one.
     ///
+    /// Rejections are meant to say "this filter didn't accept the request,
+    /// maybe another can". So for application-level errors, consider using
+    /// [`Filter::then`] instead.
+    ///
     /// # Example
     ///
     /// ```

diff --git a/src/filter/then.rs b/src/filter/then.rs
@@ -0,0 +1,95 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use futures::{ready, TryFuture};
+use pin_project::pin_project;
+
+use super::{Filter, FilterBase, Func, Internal};
+
+#[derive(Clone, Copy, Debug)]
+pub struct Then<T, F> {
+    pub(super) filter: T,
+    pub(super) callback: F,
+}
+
+impl<T, F> FilterBase for Then<T, F>
+where
+    T: Filter,
+    F: Func<T::Extract> + Clone + Send,
+    F::Output: Future + Send,
+{
+    type Extract = (<F::Output as Future>::Output,);
+    type Error = T::Error;
+    type Future = ThenFuture<T, F>;
+    #[inline]
+    fn filter(&self, _: Internal) -> Self::Future {
+        ThenFuture {
+            state: State::First(self.filter.filter(Internal), self.callback.clone()),
+        }
+    }
+}
+
+#[allow(missing_debug_implementations)]
+#[pin_project]
+pub struct ThenFuture<T, F>
+where
+    T: Filter,
+    F: Func<T::Extract>,
+    F::Output: Future + Send,
+{
+    #[pin]
+    state: State<T::Future, F>,
+}
+
+#[pin_project(project = StateProj)]
+enum State<T, F>
+where
+    T: TryFuture,
+    F: Func<T::Ok>,
+    F::Output: Future + Send,
+{
+    First(#[pin] T, F),
+    Second(#[pin] F::Output),
+    Done,
+}
+
+impl<T, F> Future for ThenFuture<T, F>
+where
+    T: Filter,
+    F: Func<T::Extract>,
+    F::Output: Future + Send,
+{
+    type Output = Result<(<F::Output as Future>::Output,), T::Error>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
+        self.project().state.poll(cx)
+    }
+}
+
+impl<T, F> Future for State<T, F>
+where
+    T: TryFuture,
+    F: Func<T::Ok>,
+    F::Output: Future + Send,
+{
+    type Output = Result<(<F::Output as Future>::Output,), T::Error>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
+        loop {
+            match self.as_mut().project() {
+                StateProj::First(first, second) => {
+                    let ex1 = ready!(first.try_poll(cx))?;
+                    let fut2 = second.call(ex1);
+                    self.set(State::Second(fut2));
+                }
+                StateProj::Second(second) => {
+                    let ex2 = (ready!(second.poll(cx)),);
+                    self.set(State::Done);
+                    return Poll::Ready(Ok(ex2));
+                }
+                StateProj::Done => panic!("polled after complete"),
+            }
+        }
+    }
+}