prost-build: support boxing fields

This allows the user to request boxing of specific fields. This is useful for
enum types that might have variants of very different size.

Author: krallin

prost-build/src/code_generator.rs

@@ -310,10 +310,15 @@ impl<'a> CodeGenerator<'a> {
         let ty = self.resolve_type(&field, fq_message_name);
 
         let boxed = !repeated
-            && (type_ == Type::Message || type_ == Type::Group)
-            && self
-                .message_graph
-                .is_nested(field.type_name(), fq_message_name);
+            && ((type_ == Type::Message || type_ == Type::Group)
+                && self
+                    .message_graph
+                    .is_nested(field.type_name(), fq_message_name))
+            || (self
+                .config
+                .boxed
+                .get_first_field(&fq_message_name, field.name())
+                .is_some());
 
         debug!(
             "    field: {:?}, type: {:?}, boxed: {}",
@@ -557,10 +562,15 @@ impl<'a> CodeGenerator<'a> {
             self.push_indent();
             let ty = self.resolve_type(&field, fq_message_name);
 
-            let boxed = (type_ == Type::Message || type_ == Type::Group)
+            let boxed = ((type_ == Type::Message || type_ == Type::Group)
                 && self
                     .message_graph
-                    .is_nested(field.type_name(), fq_message_name);
+                    .is_nested(field.type_name(), fq_message_name))
+                || (self
+                    .config
+                    .boxed
+                    .get_first_field(&oneof_name, field.name())
+                    .is_some());
 
             debug!(
                 "    oneof: {:?}, type: {:?}, boxed: {}",

prost-build/src/fixtures/field_attributes/_expected_field_attributes.rs

@@ -0,0 +1,32 @@
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct Container {
+    #[prost(oneof = "container::Data", tags = "1, 2")]
+    pub data: ::core::option::Option<container::Data>,
+}
+/// Nested message and enum types in `Container`.
+pub mod container {
+    #[allow(clippy::derive_partial_eq_without_eq)]
+    #[derive(Clone, PartialEq, ::prost::Oneof)]
+    pub enum Data {
+        #[prost(message, tag = "1")]
+        Foo(::prost::alloc::boxed::Box<super::Foo>),
+        #[prost(message, tag = "2")]
+        Bar(super::Bar),
+    }
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct Foo {
+    #[prost(string, tag = "1")]
+    pub foo: ::prost::alloc::string::String,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct Bar {
+    #[prost(message, optional, boxed, tag = "1")]
+    pub qux: ::core::option::Option<::prost::alloc::boxed::Box<Qux>>,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct Qux {}

prost-build/src/fixtures/field_attributes/field_attributes.proto

@@ -0,0 +1,21 @@
+syntax = "proto3";
+
+package field_attributes;
+
+message Container {
+  oneof data {
+    Foo foo = 1;
+    Bar bar = 2;
+  }
+}
+
+message Foo {
+  string foo = 1;
+}
+
+message Bar {
+  Qux qux = 1;
+}
+
+message Qux {
+}

prost-build/src/lib.rs

@@ -247,6 +247,7 @@ pub struct Config {
     message_attributes: PathMap<String>,
     enum_attributes: PathMap<String>,
     field_attributes: PathMap<String>,
+    boxed: PathMap<()>,
     prost_types: bool,
     strip_enum_prefix: bool,
     out_dir: Option<PathBuf>,
@@ -558,6 +559,27 @@ impl Config {
         self
     }
 
+    /// Wrap matched fields in a `Box`.
+    ///
+    /// # Arguments
+    ///
+    /// **`path`** - a path matching any number of fields. These fields get the attribute.
+    /// For details about matching fields see [`btree_map`](#method.btree_map).
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # let mut config = prost_build::Config::new();
+    /// config.boxed(".my_messages.MyMessageType.my_field");
+    /// ```
+    pub fn boxed<P>(&mut self, path: P) -> &mut Self
+    where
+        P: AsRef<str>,
+    {
+        self.boxed.insert(path.as_ref().to_string(), ());
+        self
+    }
+
     /// Configures the code generator to use the provided service generator.
     pub fn service_generator(&mut self, service_generator: Box<dyn ServiceGenerator>) -> &mut Self {
         self.service_generator = Some(service_generator);
@@ -1215,6 +1237,7 @@ impl default::Default for Config {
             message_attributes: PathMap::default(),
             enum_attributes: PathMap::default(),
             field_attributes: PathMap::default(),
+            boxed: PathMap::default(),
             prost_types: true,
             strip_enum_prefix: true,
             out_dir: None,
@@ -1651,6 +1674,39 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_generate_field_attributes() {
+        let _ = env_logger::try_init();
+
+        let out_dir = std::env::temp_dir();
+
+        Config::new()
+            .out_dir(out_dir.clone())
+            .boxed("Container.data.foo")
+            .boxed("Bar.qux")
+            .compile_protos(
+                &["src/fixtures/field_attributes/field_attributes.proto"],
+                &["src/fixtures/field_attributes"],
+            )
+            .unwrap();
+
+        let out_file = out_dir
+            .join("field_attributes.rs")
+            .as_path()
+            .display()
+            .to_string();
+
+        let content = read_all_content(&out_file).replace("\r\n", "\n");
+        let expected_content =
+            read_all_content("src/fixtures/field_attributes/_expected_field_attributes.rs")
+                .replace("\r\n", "\n");
+        assert_eq!(
+            expected_content, content,
+            "Unexpected content: \n{}",
+            content
+        );
+    }
+
     #[test]
     fn deterministic_include_file() {
         let _ = env_logger::try_init();