Skip to content

Commit

Permalink
prost-build: support boxing fields
Browse files Browse the repository at this point in the history
This allows the user to request boxing of specific fields. This is useful for
enum types that might have variants of very different size.
  • Loading branch information
krallin committed Mar 31, 2023
1 parent b72967a commit 4fea4ad
Show file tree
Hide file tree
Showing 4 changed files with 125 additions and 6 deletions.
22 changes: 16 additions & 6 deletions prost-build/src/code_generator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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: {}",
Expand Down Expand Up @@ -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: {}",
Expand Down
Original file line number Diff line number Diff line change
@@ -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 {}
21 changes: 21 additions & 0 deletions prost-build/src/fixtures/field_attributes/field_attributes.proto
Original file line number Diff line number Diff line change
@@ -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 {
}
56 changes: 56 additions & 0 deletions prost-build/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -1192,6 +1214,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,
Expand Down Expand Up @@ -1597,6 +1620,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();
Expand Down

0 comments on commit 4fea4ad

Please sign in to comment.