compiler: support protobuf well-known types

Cargo.toml

@@ -25,7 +25,7 @@ ttrpc-codegen = { version = "0.6.0", path = "./ttrpc-codegen" }
 ttrpc-compiler = { version = "0.8.0", path = "./compiler" }
 protobuf = "3.7.2"
 protobuf-codegen = "3.7.2"
-protobuf-support = "3.7.2"
+protobuf-parse = "3.7.2"
 
 [package]
 name = "ttrpc"

README.md

@@ -51,6 +51,9 @@ fn main() {
 }
 ```
 
+Canonical Google well-known type imports, such as `google/protobuf/timestamp.proto`, are available
+automatically when generating code programmatically and do not require an additional include path.
+
 # async/.await
 ttrpc-rust supports async/.await. By using async/.await you can reduce the overhead and resource consumption caused by threads.
 

compiler/README.md

@@ -7,6 +7,12 @@ generate rust version ttrpc code from proto files.
 - [Manual Generation](https://github.com/containerd/ttrpc-rust#1-generate-with-protoc-command) uses ttrpc-compiler as a protoc plugin
 - [Programmatic Generation](https://github.com/containerd/ttrpc-rust#2-generate-programmatically) uses ttrpc-compiler as a rust crate
 
+## Well-known types
+
+RPC inputs and outputs from canonical Google well-known proto dependencies reference the
+corresponding types provided by the `protobuf` runtime. Well-known proto files explicitly selected
+for generation continue to use their locally generated modules.
+
 ## Versions
 | ttrpc-compiler version | ttrpc version |
 | ------------- | ------------- |

compiler/src/codegen.rs

@@ -43,7 +43,10 @@ use std::{
 };
 
 use crate::{
-    util::proto_path_to_rust_mod, util::scope::RootScope, util::writer::CodeWriter, Customize,
+    util::proto_path_to_rust_mod,
+    util::scope::{RootScope, RustType},
+    util::writer::CodeWriter,
+    Customize,
 };
 use protobuf::{
     descriptor::*,
@@ -88,21 +91,28 @@ impl<'a> MethodGen<'a> {
     }
 
     fn input(&self) -> String {
-        format!(
-            "super::{}",
-            self.root_scope
-                .find_message(self.proto.input_type())
-                .rust_fq_name()
-        )
+        self.root_scope
+            .rust_type(self.proto.input_type())
+            .to_string()
     }
 
     fn output(&self) -> String {
-        format!(
-            "super::{}",
-            self.root_scope
-                .find_message(self.proto.output_type())
-                .rust_fq_name()
-        )
+        self.root_scope
+            .rust_type(self.proto.output_type())
+            .to_string()
+    }
+
+    fn request_handler_call(&self, macro_name: &str, request: &str) -> String {
+        match self.root_scope.rust_type(self.proto.input_type()) {
+            RustType::Generated { module, name } => format!(
+                "::ttrpc::{macro_name}!(self, ctx, {request}, {module}, {name}, {});",
+                self.name()
+            ),
+            rust_type @ RustType::ProtobufRuntime { .. } => format!(
+                "::ttrpc::{macro_name}!(self, ctx, {request}, {rust_type}, {});",
+                self.name()
+            ),
+        }
     }
 
     fn method_type(&self) -> (MethodType, String) {
@@ -164,10 +174,7 @@ impl<'a> MethodGen<'a> {
         |w| {
             w.block("fn handler(&self, ctx: ::ttrpc::TtrpcContext, req: ::ttrpc::Request) -> ::ttrpc::Result<()> {", "}",
             |w| {
-                w.write_line(format!("::ttrpc::request_handler!(self, ctx, req, {}, {}, {});",
-                                        proto_path_to_rust_mod(self.root_scope.find_message(self.proto.input_type()).fd.name()),
-                                        self.root_scope.find_message(self.proto.input_type()).rust_name(),
-                                        self.name()));
+                w.write_line(self.request_handler_call("request_handler", "req"));
                 w.write_line("Ok(())");
             });
         });
@@ -181,10 +188,7 @@ impl<'a> MethodGen<'a> {
                 |w| {
                     w.block("async fn handler(&self, ctx: ::ttrpc::r#async::TtrpcContext, req: ::ttrpc::Request) -> ::ttrpc::Result<::ttrpc::Response> {", "}",
                         |w| {
-                            w.write_line(format!("::ttrpc::async_request_handler!(self, ctx, req, {}, {}, {});",
-                                        proto_path_to_rust_mod(self.root_scope.find_message(self.proto.input_type()).fd.name()),
-                                        self.root_scope.find_message(self.proto.input_type()).rust_name(),
-                                        self.name()));
+                            w.write_line(self.request_handler_call("async_request_handler", "req"));
                     });
             });
             }
@@ -205,10 +209,7 @@ impl<'a> MethodGen<'a> {
                 |w| {
                     w.block("async fn handler(&self, ctx: ::ttrpc::r#async::TtrpcContext, mut inner: ::ttrpc::r#async::StreamInner) -> ::ttrpc::Result<Option<::ttrpc::Response>> {", "}",
                         |w| {
-                            w.write_line(format!("::ttrpc::async_server_streamimg_handler!(self, ctx, inner, {}, {}, {});",
-                                        proto_path_to_rust_mod(self.root_scope.find_message(self.proto.input_type()).fd.name()),
-                                        self.root_scope.find_message(self.proto.input_type()).rust_name(),
-                                        self.name()));
+                            w.write_line(self.request_handler_call("async_server_streamimg_handler", "inner"));
                     });
             });
             }
@@ -712,7 +713,7 @@ pub fn gen(
     let files_map: HashMap<&str, &FileDescriptorProto> =
         file_descriptors.iter().map(|f| (f.name(), f)).collect();
 
-    let root_scope = RootScope { file_descriptors };
+    let root_scope = RootScope::new(file_descriptors, files_to_generate);
 
     let mut results = CodeGeneratorResponse::new();
     results.set_supported_features(CodeGeneratorResponse_Feature::FEATURE_PROTO3_OPTIONAL as _);

compiler/src/util/scope.rs

@@ -2,9 +2,28 @@
 //! the protobuf / protobuf-codegen crates, but were then removed.
 //! The missing functionalities have been reimplemented in this module.
 
+use std::fmt;
+
 use protobuf::descriptor::{DescriptorProto, FileDescriptorProto};
 
-// vendered from https://github.com/stepancheg/rust-protobuf/blob/v3.7.2/protobuf-codegen/src/gen/rust/keywords.rs
+use super::to_snake_case;
+
+const DESCRIPTOR_PROTO_FILE: &str = "google/protobuf/descriptor.proto";
+
+const WELL_KNOWN_TYPE_PROTO_FILES: &[&str] = &[
+    "google/protobuf/any.proto",
+    "google/protobuf/api.proto",
+    "google/protobuf/duration.proto",
+    "google/protobuf/empty.proto",
+    "google/protobuf/field_mask.proto",
+    "google/protobuf/source_context.proto",
+    "google/protobuf/struct.proto",
+    "google/protobuf/timestamp.proto",
+    "google/protobuf/type.proto",
+    "google/protobuf/wrappers.proto",
+];
+
+// vendored from https://github.com/stepancheg/rust-protobuf/blob/v3.7.2/protobuf-codegen/src/gen/rust/keywords.rs
 fn is_rust_keyword(ident: &str) -> bool {
     #[rustfmt::skip]
     static RUST_KEYWORDS: &[&str] = &[
@@ -71,7 +90,8 @@ fn is_rust_keyword(ident: &str) -> bool {
 // reimplementation based on https://github.com/stepancheg/rust-protobuf/blob/v3.7.2/protobuf-codegen/src/gen/scope.rs#L26
 // it only implements the `find_message` method with not extra dependencies
 pub struct RootScope<'a> {
-    pub file_descriptors: &'a [FileDescriptorProto],
+    file_descriptors: &'a [FileDescriptorProto],
+    files_to_generate: &'a [String],
 }
 
 // re-implementation of https://github.com/stepancheg/rust-protobuf/blob/v3.7.2/protobuf-codegen/src/gen/scope.rs#L340
@@ -82,6 +102,21 @@ pub struct ScopedMessage<'a> {
     pub msg: &'a DescriptorProto,
 }
 
+#[derive(Debug, Eq, PartialEq)]
+pub enum RustType {
+    Generated { module: String, name: String },
+    ProtobufRuntime { path: String },
+}
+
+impl fmt::Display for RustType {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Generated { module, name } => write!(f, "super::{module}::{name}"),
+            Self::ProtobufRuntime { path } => f.write_str(path),
+        }
+    }
+}
+
 impl ScopedMessage<'_> {
     pub fn prefix(&self) -> String {
         let mut prefix = String::new();
@@ -108,17 +143,57 @@ impl ScopedMessage<'_> {
         r
     }
 
-    // fully-qualified name of this type
-    pub fn rust_fq_name(&self) -> String {
-        format!(
-            "{}::{}",
-            super::proto_path_to_rust_mod(self.fd.name()),
-            self.rust_name()
-        )
+    fn protobuf_runtime_path(&self) -> String {
+        let mut path = String::new();
+        for message in &self.path {
+            path.push_str(&to_snake_case(message.name()));
+            path.push_str("::");
+        }
+        path.push_str(self.msg.name());
+        path
     }
 }
 
 impl<'a> RootScope<'a> {
+    pub fn new(
+        file_descriptors: &'a [FileDescriptorProto],
+        files_to_generate: &'a [String],
+    ) -> Self {
+        Self {
+            file_descriptors,
+            files_to_generate,
+        }
+    }
+
+    pub fn rust_type(&'a self, fqn: impl AsRef<str>) -> RustType {
+        let message = self.find_message(fqn);
+        let file_name = message.fd.name();
+        let module = super::proto_path_to_rust_mod(file_name);
+        let name = message.rust_name();
+
+        if !self.files_to_generate.iter().any(|file| file == file_name) {
+            if file_name == DESCRIPTOR_PROTO_FILE {
+                return RustType::ProtobufRuntime {
+                    path: format!(
+                        "::protobuf::descriptor::{}",
+                        message.protobuf_runtime_path()
+                    ),
+                };
+            }
+
+            if WELL_KNOWN_TYPE_PROTO_FILES.contains(&file_name) {
+                return RustType::ProtobufRuntime {
+                    path: format!(
+                        "::protobuf::well_known_types::{module}::{}",
+                        message.protobuf_runtime_path()
+                    ),
+                };
+            }
+        }
+
+        RustType::Generated { module, name }
+    }
+
     pub fn find_message(&'a self, fqn: impl AsRef<str>) -> ScopedMessage<'a> {
         let Some(fqn1) = fqn.as_ref().strip_prefix(".") else {
             panic!("name must start with dot: {}", fqn.as_ref())
@@ -147,6 +222,124 @@ impl<'a> RootScope<'a> {
                 }
             }
         }
-        panic!("enum not found by name: {}", fqn.as_ref())
+        panic!("message not found by name: {}", fqn.as_ref())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn file_descriptor(name: &str, package: &str, message: &str) -> FileDescriptorProto {
+        let mut descriptor = FileDescriptorProto::new();
+        descriptor.set_name(name.to_owned());
+        descriptor.set_package(package.to_owned());
+
+        let mut message_descriptor = DescriptorProto::new();
+        message_descriptor.set_name(message.to_owned());
+        descriptor.message_type.push(message_descriptor);
+        descriptor
+    }
+
+    #[test]
+    fn well_known_dependencies_use_protobuf_runtime() {
+        let cases = [
+            ("any.proto", "Any", "any"),
+            ("api.proto", "Api", "api"),
+            ("duration.proto", "Duration", "duration"),
+            ("empty.proto", "Empty", "empty"),
+            ("field_mask.proto", "FieldMask", "field_mask"),
+            ("source_context.proto", "SourceContext", "source_context"),
+            ("struct.proto", "Struct", "struct_"),
+            ("timestamp.proto", "Timestamp", "timestamp"),
+            ("type.proto", "Type", "type_"),
+            ("wrappers.proto", "StringValue", "wrappers"),
+        ];
+        let files_to_generate = ["service.proto".to_owned()];
+
+        for (proto, message, module) in cases {
+            let descriptors = [file_descriptor(
+                &format!("google/protobuf/{proto}"),
+                "google.protobuf",
+                message,
+            )];
+            let scope = RootScope::new(&descriptors, &files_to_generate);
+
+            assert_eq!(
+                format!("::protobuf::well_known_types::{module}::{message}"),
+                scope
+                    .rust_type(format!(".google.protobuf.{message}"))
+                    .to_string()
+            );
+        }
+    }
+
+    #[test]
+    fn explicitly_generated_well_known_type_uses_local_module() {
+        let descriptors = [file_descriptor(
+            "google/protobuf/timestamp.proto",
+            "google.protobuf",
+            "Timestamp",
+        )];
+        let files_to_generate = ["google/protobuf/timestamp.proto".to_owned()];
+        let scope = RootScope::new(&descriptors, &files_to_generate);
+
+        assert_eq!(
+            "super::timestamp::Timestamp",
+            scope.rust_type(".google.protobuf.Timestamp").to_string()
+        );
+    }
+
+    #[test]
+    fn descriptor_dependency_uses_protobuf_runtime() {
+        let descriptors = [file_descriptor(
+            DESCRIPTOR_PROTO_FILE,
+            "google.protobuf",
+            "FileDescriptorProto",
+        )];
+        let files_to_generate = ["service.proto".to_owned()];
+        let scope = RootScope::new(&descriptors, &files_to_generate);
+
+        assert_eq!(
+            "::protobuf::descriptor::FileDescriptorProto",
+            scope
+                .rust_type(".google.protobuf.FileDescriptorProto")
+                .to_string()
+        );
+    }
+
+    #[test]
+    fn nested_descriptor_dependency_uses_protobuf_runtime() {
+        let mut descriptor =
+            file_descriptor(DESCRIPTOR_PROTO_FILE, "google.protobuf", "DescriptorProto");
+        let mut nested = DescriptorProto::new();
+        nested.set_name("ExtensionRange".to_owned());
+        descriptor.message_type[0].nested_type.push(nested);
+        let descriptors = [descriptor];
+        let files_to_generate = ["service.proto".to_owned()];
+        let scope = RootScope::new(&descriptors, &files_to_generate);
+
+        assert_eq!(
+            "::protobuf::descriptor::descriptor_proto::ExtensionRange",
+            scope
+                .rust_type(".google.protobuf.DescriptorProto.ExtensionRange")
+                .to_string()
+        );
+    }
+
+    #[test]
+    fn non_well_known_google_type_uses_local_module() {
+        let descriptors = [file_descriptor(
+            "google/protobuf/custom.proto",
+            "google.protobuf",
+            "Custom",
+        )];
+        let files_to_generate = ["service.proto".to_owned()];
+        let scope = RootScope::new(&descriptors, &files_to_generate);
+
+        assert_eq!(
+            "super::custom::Custom",
+            scope.rust_type(".google.protobuf.Custom").to_string()
+        );
     }
 }

example/build.rs

@@ -20,6 +20,7 @@ fn main() {
         "protocols/protos/health.proto",
         "protocols/protos/google/protobuf/empty.proto",
         "protocols/protos/oci.proto",
+        "protocols/protos/well_known.proto",
     ];
     let protobuf_customized = ProtobufCustomize::default().gen_mod_rs(true);
 
@@ -38,6 +39,7 @@ fn main() {
 
     // Only async support stream currently.
     protos.push("protocols/protos/streaming.proto");
+    protos.push("protocols/protos/well_known_streaming.proto");
 
     Codegen::new()
         .out_dir("protocols/asynchronous")