From f9da61d00cdb5b481d8fee24be10d1ae4221f0da Mon Sep 17 00:00:00 2001
From: fakeshadow <24548779@qq.com>
Date: Wed, 27 Nov 2024 20:32:21 +0800
Subject: [PATCH] add nightly feature enabling async closure as middleware

---
 service/Cargo.toml                 |  1 +
 service/src/lib.rs                 |  7 +++++++
 service/src/middleware/async_fn.rs | 15 +++++++++++++++
 service/src/service/ext.rs         | 31 ++++++++++++++++++++++++++++++
 4 files changed, 54 insertions(+)

diff --git a/service/Cargo.toml b/service/Cargo.toml
index 5a54eead..6a0e5ace 100644
--- a/service/Cargo.toml
+++ b/service/Cargo.toml
@@ -15,6 +15,7 @@ workspace = true
 [features]
 alloc = []
 std = []
+nightly = []
 
 [dev-dependencies]
 xitca-unsafe-collection = "0.2.0"
diff --git a/service/src/lib.rs b/service/src/lib.rs
index 4659da90..4b824137 100644
--- a/service/src/lib.rs
+++ b/service/src/lib.rs
@@ -36,8 +36,15 @@
 //! ```
 #![no_std]
 #![forbid(unsafe_code)]
+#![cfg_attr(feature = "nightly", feature(async_closure))]
 
+#[cfg(not(feature = "nightly"))]
 mod async_fn;
+#[cfg(feature = "nightly")]
+mod async_fn {
+    pub use core::ops::AsyncFn;
+}
+
 mod service;
 
 pub mod middleware;
diff --git a/service/src/middleware/async_fn.rs b/service/src/middleware/async_fn.rs
index 56bc81f7..ec27eedb 100644
--- a/service/src/middleware/async_fn.rs
+++ b/service/src/middleware/async_fn.rs
@@ -22,6 +22,7 @@ where
     }
 }
 
+#[cfg(not(feature = "nightly"))]
 impl<S, Req, F, Res, Err> Service<Req> for PipelineT<S, F, marker::AsyncFn>
 where
     F: for<'s> async_fn::AsyncFn<(&'s S, Req), Output = Result<Res, Err>>,
@@ -34,3 +35,17 @@ where
         self.second.call((&self.first, req)).await
     }
 }
+
+#[cfg(feature = "nightly")]
+impl<S, Req, F, Res, Err> Service<Req> for PipelineT<S, F, marker::AsyncFn>
+where
+    F: for<'s> async_fn::AsyncFn(&'s S, Req) -> Result<Res, Err>,
+{
+    type Response = Res;
+    type Error = Err;
+
+    #[inline]
+    async fn call(&self, req: Req) -> Result<Self::Response, Self::Error> {
+        (self.second)(&self.first, req).await
+    }
+}
diff --git a/service/src/service/ext.rs b/service/src/service/ext.rs
index eb313d78..29bef279 100644
--- a/service/src/service/ext.rs
+++ b/service/src/service/ext.rs
@@ -17,6 +17,7 @@ pub trait ServiceExt<Arg>: Service<Arg> {
         PipelineT::new(self, build)
     }
 
+    #[cfg(not(feature = "nightly"))]
     /// Function version of [Self::enclosed] method.
     fn enclosed_fn<T, Req>(self, func: T) -> PipelineT<Self, middleware::AsyncFn<T>, marker::BuildEnclosed>
     where
@@ -26,6 +27,16 @@ pub trait ServiceExt<Arg>: Service<Arg> {
         self.enclosed(middleware::AsyncFn(func))
     }
 
+    #[cfg(feature = "nightly")]
+    /// Function version of [Self::enclosed] method.
+    fn enclosed_fn<T, Req, O>(self, func: T) -> PipelineT<Self, middleware::AsyncFn<T>, marker::BuildEnclosed>
+    where
+        T: for<'s> async_fn::AsyncFn(&'s Self::Response, Req) -> O + Clone,
+        Self: Sized,
+    {
+        self.enclosed(middleware::AsyncFn(func))
+    }
+
     /// Mutate `<<Self::Response as Service<Req>>::Future as Future>::Output` type with given
     /// closure.
     fn map<F, Res, ResMap>(self, mapper: F) -> PipelineT<Self, F, marker::BuildMap>
@@ -153,6 +164,26 @@ mod test {
         assert_eq!(res, "251");
     }
 
+    #[cfg(feature = "nightly")]
+    #[test]
+    fn enclosed_fn_async_closure() {
+        let res = fn_service(index)
+            .enclosed_fn(async |service, req| {
+                let res = service.call(req).now_or_panic()?;
+                assert_eq!(res, "996");
+                Ok::<_, ()>("251")
+            })
+            .call(())
+            .now_or_panic()
+            .unwrap()
+            .call("996")
+            .now_or_panic()
+            .ok()
+            .unwrap();
+
+        assert_eq!(res, "251");
+    }
+
     #[cfg(feature = "alloc")]
     #[test]
     fn enclosed_opt() {