Make the generated code less complicated

crates/twirp-build/src/lib.rs

@@ -43,37 +43,24 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
 where
     T: {service_name} + Send + Sync + 'static,
 {{
-    twirp::Router::new()"#,
+    twirp::details::TwirpRouterBuilder::new(api)"#,
         )
         .unwrap();
         for m in &service.methods {
             let uri = &m.proto_name;
             let rust_method_name = &m.name;
             writeln!(
                 buf,
-                r#"        .route(
-            "/{uri}",
-            twirp::details::post(
-                |twirp::details::State(api): twirp::details::State<std::sync::Arc<T>>,
-                 req: twirp::details::Request| async move {{
-                    twirp::server::handle_request(
-                        req,
-                        move |req| async move {{
-                            api.{rust_method_name}(req).await
-                        }},
-                    )
-                    .await
-                }},
-            ),
-        )"#,
+                r#"        .route("/{uri}", |api: std::sync::Arc<T>| move |req| async move {{
+            api.{rust_method_name}(req).await
+        }})"#,
             )
             .unwrap();
         }
         writeln!(
             buf,
             r#"
-        .with_state(api)
-        .fallback(twirp::server::not_found_handler)
+        .build()
 }}"#
         )
         .unwrap();

crates/twirp/src/details.rs

@@ -1,10 +1,51 @@
 //! Undocumented features that are public for use in generated code (see `twirp-build`).
 
-#[doc(hidden)]
-pub use axum::extract::{Request, State};
+use std::future::Future;
 
-#[doc(hidden)]
-pub use axum::routing::post;
+use axum::extract::{Request, State};
+use axum::Router;
 
-#[doc(hidden)]
-pub use axum::response::Response;
+use crate::{server, TwirpErrorResponse};
+
+/// Builder object used by generated code to build a router.
+pub struct TwirpRouterBuilder<S> {
+    service: S,
+    router: Router<S>,
+}
+
+impl<S> TwirpRouterBuilder<S>
+where
+    S: Clone + Send + Sync + 'static,
+{
+    pub fn new(service: S) -> Self {
+        TwirpRouterBuilder {
+            service,
+            router: Router::new(),
+        }
+    }
+
+    pub fn route<F, G, Fut, RequestMessage, ResponseMessage>(self, url: &str, f: F) -> Self
+    where
+        F: Fn(S) -> G + Clone + Send + 'static,
+        G: FnOnce(RequestMessage) -> Fut + Clone + Sync + Send + 'static,
+        Fut: Future<Output = Result<ResponseMessage, TwirpErrorResponse>> + Send,
+        RequestMessage: prost::Message + Default + serde::de::DeserializeOwned,
+        ResponseMessage: prost::Message + serde::Serialize,
+    {
+        TwirpRouterBuilder {
+            service: self.service,
+            router: self.router.route(
+                url,
+                axum::routing::post(move |State(api): State<S>, req: Request| async move {
+                    server::handle_request(req, f(api)).await
+                }),
+            ),
+        }
+    }
+
+    pub fn build(self) -> axum::Router {
+        self.router
+            .fallback(crate::server::not_found_handler)
+            .with_state(self.service)
+    }
+}

crates/twirp/src/test.rs

@@ -11,6 +11,7 @@ use serde::de::DeserializeOwned;
 use tokio::task::JoinHandle;
 use tokio::time::Instant;
 
+use crate::details::TwirpRouterBuilder;
 use crate::server::Timings;
 use crate::{error, Client, Result, TwirpErrorResponse};
 
@@ -28,35 +29,14 @@ pub fn test_api_router() -> Router {
     let api = Arc::new(TestAPIServer {});
 
     // NB: This part would be generated
-    let test_router = crate::Router::new()
-        .route(
-            "/Ping",
-            crate::details::post(
-                |crate::details::State(api): crate::details::State<Arc<TestAPIServer>>,
-                 req: crate::details::Request| async move {
-                    crate::server::handle_request(
-                        req,
-                        move |req| async move { api.ping(req).await },
-                    )
-                    .await
-                },
-            ),
-        )
-        .route(
-            "/Boom",
-            crate::details::post(
-                |crate::details::State(api): crate::details::State<Arc<TestAPIServer>>,
-                 req: crate::details::Request| async move {
-                    crate::server::handle_request(
-                        req,
-                        move |req| async move { api.boom(req).await },
-                    )
-                    .await
-                },
-            ),
-        )
-        .fallback(crate::server::not_found_handler)
-        .with_state(api);
+    let test_router = TwirpRouterBuilder::new(api)
+        .route("/Ping", |api: Arc<TestAPIServer>| {
+            move |req| async move { api.ping(req).await }
+        })
+        .route("/Boom", |api: Arc<TestAPIServer>| {
+            move |req| async move { api.boom(req).await }
+        })
+        .build();
 
     axum::Router::new()
         .nest("/twirp/test.TestAPI", test_router)