From f9ccc27b3990cc4b8f3088b02952909bff3c6447 Mon Sep 17 00:00:00 2001
From: Yotam Ofek <yotam.ofek@gmail.com>
Date: Thu, 2 Jan 2025 20:59:00 +0000
Subject: [PATCH] build: refactor server trait gen to dedup some code

---
 tonic-build/src/server.rs | 126 +++++++++++++-------------------------
 1 file changed, 44 insertions(+), 82 deletions(-)

diff --git a/tonic-build/src/server.rs b/tonic-build/src/server.rs
index e2d0aacd9..498c31d78 100644
--- a/tonic-build/src/server.rs
+++ b/tonic-build/src/server.rs
@@ -257,92 +257,54 @@ fn generate_trait_methods<T: Service>(
             quote!(&self)
         };
 
-        let method = match (
-            method.client_streaming(),
-            method.server_streaming(),
-            generate_default_stubs,
-        ) {
-            (false, false, true) => {
-                quote! {
-                    #method_doc
-                    async fn #name(#self_param, request: tonic::Request<#req_message>)
-                        -> std::result::Result<tonic::Response<#res_message>, tonic::Status> {
-                        Err(tonic::Status::unimplemented("Not yet implemented"))
-                    }
-                }
-            }
-            (false, false, false) => {
-                quote! {
-                    #method_doc
-                    async fn #name(#self_param, request: tonic::Request<#req_message>)
-                        -> std::result::Result<tonic::Response<#res_message>, tonic::Status>;
-                }
-            }
-            (true, false, true) => {
-                quote! {
-                    #method_doc
-                    async fn #name(#self_param, request: tonic::Request<tonic::Streaming<#req_message>>)
-                        -> std::result::Result<tonic::Response<#res_message>, tonic::Status> {
-                        Err(tonic::Status::unimplemented("Not yet implemented"))
-                    }
-                }
-            }
-            (true, false, false) => {
-                quote! {
-                    #method_doc
-                    async fn #name(#self_param, request: tonic::Request<tonic::Streaming<#req_message>>)
-                        -> std::result::Result<tonic::Response<#res_message>, tonic::Status>;
-                }
-            }
-            (false, true, true) => {
-                quote! {
-                    #method_doc
-                    async fn #name(#self_param, request: tonic::Request<#req_message>)
-                        -> std::result::Result<tonic::Response<BoxStream<#res_message>>, tonic::Status> {
-                        Err(tonic::Status::unimplemented("Not yet implemented"))
-                    }
+        let req_param_type = if method.client_streaming() {
+            quote!(tonic::Request<tonic::Streaming<#req_message>>)
+        } else {
+            quote!(tonic::Request<#req_message>)
+        };
+
+        let partial_sig = quote! {
+            #method_doc
+            async fn #name(#self_param, request: #req_param_type)
+        };
+
+        let body_or_semicolon = if generate_default_stubs {
+            quote! {
+                {
+                    Err(tonic::Status::unimplemented("Not yet implemented"))
                 }
             }
-            (false, true, false) => {
-                let stream = quote::format_ident!("{}Stream", method.identifier());
-                let stream_doc = generate_doc_comment(format!(
-                    " Server streaming response type for the {} method.",
-                    method.identifier()
-                ));
-
-                quote! {
-                    #stream_doc
-                    type #stream: tonic::codegen::tokio_stream::Stream<Item = std::result::Result<#res_message, tonic::Status>> + std::marker::Send + 'static;
-
-                    #method_doc
-                    async fn #name(#self_param, request: tonic::Request<#req_message>)
-                        -> std::result::Result<tonic::Response<Self::#stream>, tonic::Status>;
-                }
+        } else {
+            quote!(;)
+        };
+
+        let result = |ok| quote!(std::result::Result<#ok, tonic::Status>);
+        let response_result = |message| result(quote!(tonic::Response<#message>));
+
+        let method = if !method.server_streaming() {
+            let return_ty = response_result(res_message);
+            quote! {
+                #partial_sig -> #return_ty #body_or_semicolon
             }
-            (true, true, true) => {
-                quote! {
-                    #method_doc
-                    async fn #name(#self_param, request: tonic::Request<tonic::Streaming<#req_message>>)
-                        -> std::result::Result<tonic::Response<BoxStream<#res_message>>, tonic::Status> {
-                        Err(tonic::Status::unimplemented("Not yet implemented"))
-                    }
-                }
+        } else if generate_default_stubs {
+            let return_ty = response_result(quote!(BoxStream<#res_message>));
+            quote! {
+                #partial_sig -> #return_ty #body_or_semicolon
             }
-            (true, true, false) => {
-                let stream = quote::format_ident!("{}Stream", method.identifier());
-                let stream_doc = generate_doc_comment(format!(
-                    " Server streaming response type for the {} method.",
-                    method.identifier()
-                ));
-
-                quote! {
-                    #stream_doc
-                    type #stream: tonic::codegen::tokio_stream::Stream<Item = std::result::Result<#res_message, tonic::Status>> + std::marker::Send + 'static;
-
-                    #method_doc
-                    async fn #name(#self_param, request: tonic::Request<tonic::Streaming<#req_message>>)
-                        -> std::result::Result<tonic::Response<Self::#stream>, tonic::Status>;
-                }
+        } else {
+            let stream = quote::format_ident!("{}Stream", method.identifier());
+            let stream_doc = generate_doc_comment(format!(
+                " Server streaming response type for the {} method.",
+                method.identifier()
+            ));
+            let stream_item_ty = result(res_message);
+            let stream_ty = quote!(tonic::codegen::tokio_stream::Stream<Item = #stream_item_ty> + std::marker::Send + 'static);
+            let return_ty = response_result(quote!(Self::#stream));
+            quote! {
+                #stream_doc
+                type #stream: #stream_ty;
+
+                #partial_sig -> #return_ty #body_or_semicolon
             }
         };