feat: OpenEnum, Config::typed_enum_fields

Introduce OpenEnum to represent fields of Protobuf enumerated types
with the possibility of unknown values.

A new enum_type annotation is supported in the prost attribute
inside derives, which allows to specify the type-checked representation
of enum types in message fields and oneof variants.
The accepted values are "open" or "closed".

Use OpenEnum in generated code for enum-typed fields of messages
instead of i32, if the `enum_type` annotation is set to "open".

OpenEnum features get and set methods, which are meant to be
the primary ways to get and set open enumeration fields,
instead of the getters and setters generated by Message derive.

Add typed_enum_fields method to prost-build configuration, which
allows type-checked representation of enumerations in fields of
message structs and variants of oneof enums. The argument and the
invocation order works like with the boxed method.

Depending on the syntax (and preparing for the future support of
editions), the type-checked representation can be closed (for proto2)
or open (for proto3). The former is represented by the generated
enum type itself, while the latter is represented by OpenEnum
wrapping the enum type.

Author: mzabaluev

benchmarks/benches/dataset.rs

@@ -37,7 +37,7 @@ fn benchmark_dataset<M>(criterion: &mut Criterion, name: &str, dataset: &'static
 where
     M: prost::Message + Default + 'static,
 {
-    let mut group = criterion.benchmark_group(&format!("dataset/{}", name));
+    let mut group = criterion.benchmark_group(format!("dataset/{}", name));
 
     group.bench_function("merge", move |b| {
         let dataset = load_dataset(dataset).unwrap();

prost-build/src/code_generator.rs

@@ -388,7 +388,7 @@ impl<'b> CodeGenerator<'_, 'b> {
 
     fn append_field(&mut self, fq_message_name: &str, field: &Field) {
         let type_ = field.descriptor.r#type();
-        let repeated = field.descriptor.label() == Label::Repeated;
+        let repeated = field.descriptor.label == Some(Label::Repeated as i32);
         let deprecated = self.deprecated(&field.descriptor);
         let optional = self.optional(&field.descriptor);
         let boxed = self
@@ -415,12 +415,16 @@ impl<'b> CodeGenerator<'_, 'b> {
         let type_tag = self.field_type_tag(&field.descriptor);
         self.buf.push_str(&type_tag);
 
-        if type_ == Type::Bytes {
-            let bytes_type = self
-                .context
-                .bytes_type(fq_message_name, field.descriptor.name());
-            self.buf
-                .push_str(&format!("={:?}", bytes_type.annotation()));
+        match type_ {
+            Type::Bytes => {
+                let bytes_type = self
+                    .context
+                    .bytes_type(fq_message_name, field.descriptor.name());
+                self.buf
+                    .push_str(&format!("={:?}", bytes_type.annotation()));
+            }
+            Type::Enum => self.push_enum_type_annotation(fq_message_name, field.descriptor.name()),
+            _ => {}
         }
 
         match field.descriptor.label() {
@@ -537,12 +541,16 @@ impl<'b> CodeGenerator<'_, 'b> {
         let value_tag = self.map_value_type_tag(value);
 
         self.buf.push_str(&format!(
-            "#[prost({}=\"{}, {}\", tag=\"{}\")]\n",
+            "#[prost({}=\"{}, {}\"",
             map_type.annotation(),
             key_tag,
             value_tag,
-            field.descriptor.number()
         ));
+        if value.r#type() == Type::Enum {
+            self.push_enum_type_annotation(fq_message_name, field.descriptor.name());
+        }
+        self.buf
+            .push_str(&format!(", tag=\"{}\")]\n", field.descriptor.number()));
         self.append_field_attributes(fq_message_name, field.descriptor.name());
         self.push_indent();
         self.buf.push_str(&format!(
@@ -630,11 +638,12 @@ impl<'b> CodeGenerator<'_, 'b> {
 
             self.push_indent();
             let ty_tag = self.field_type_tag(&field.descriptor);
-            self.buf.push_str(&format!(
-                "#[prost({}, tag=\"{}\")]\n",
-                ty_tag,
-                field.descriptor.number()
-            ));
+            self.buf.push_str(&format!("#[prost({}", ty_tag,));
+            if field.descriptor.r#type() == Type::Enum {
+                self.push_enum_type_annotation(&oneof_name, field.descriptor.name());
+            }
+            self.buf
+                .push_str(&format!(", tag=\"{}\")]\n", field.descriptor.number()));
             self.append_field_attributes(&oneof_name, field.descriptor.name());
 
             self.push_indent();
@@ -930,13 +939,21 @@ impl<'b> CodeGenerator<'_, 'b> {
         self.buf.push_str("}\n");
     }
 
+    fn push_enum_type_annotation(&mut self, fq_message_name: &str, field_name: &str) {
+        match self.enum_field_repr(fq_message_name, field_name) {
+            EnumRepr::Int => {}
+            EnumRepr::Open => self.buf.push_str(", enum_type=\"open\""),
+            EnumRepr::Closed => self.buf.push_str(", enum_type=\"closed\""),
+        }
+    }
+
     fn resolve_type(&self, field: &FieldDescriptorProto, fq_message_name: &str) -> String {
         match field.r#type() {
             Type::Float => String::from("f32"),
             Type::Double => String::from("f64"),
             Type::Uint32 | Type::Fixed32 => String::from("u32"),
             Type::Uint64 | Type::Fixed64 => String::from("u64"),
-            Type::Int32 | Type::Sfixed32 | Type::Sint32 | Type::Enum => String::from("i32"),
+            Type::Int32 | Type::Sfixed32 | Type::Sint32 => String::from("i32"),
             Type::Int64 | Type::Sfixed64 | Type::Sint64 => String::from("i64"),
             Type::Bool => String::from("bool"),
             Type::String => format!("{}::alloc::string::String", self.context.prost_path()),
@@ -946,6 +963,15 @@ impl<'b> CodeGenerator<'_, 'b> {
                 .rust_type()
                 .to_owned(),
             Type::Group | Type::Message => self.resolve_ident(field.type_name()),
+            Type::Enum => match self.enum_field_repr(fq_message_name, field.name()) {
+                EnumRepr::Int => String::from("i32"),
+                EnumRepr::Open => format!(
+                    "{}::OpenEnum<{}>",
+                    self.context.prost_path(),
+                    self.resolve_ident(field.type_name())
+                ),
+                EnumRepr::Closed => self.resolve_ident(field.type_name()),
+            },
         }
     }
 
@@ -987,6 +1013,24 @@ impl<'b> CodeGenerator<'_, 'b> {
             .join("::")
     }
 
+    fn enum_field_repr(&self, fq_message_name: &str, field_name: &str) -> EnumRepr {
+        if self
+            .context
+            .is_typed_enum_field(fq_message_name, field_name)
+        {
+            // FIXME: store information for the code generator to know when
+            // proto3 enums are used in proto2, where they should be open
+            // accordingly to the spec:
+            // https://protobuf.dev/programming-guides/enum/#spec
+            match self.syntax {
+                Syntax::Proto2 => EnumRepr::Closed,
+                Syntax::Proto3 => EnumRepr::Open,
+            }
+        } else {
+            EnumRepr::Int
+        }
+    }
+
     fn field_type_tag(&self, field: &FieldDescriptorProto) -> Cow<'static, str> {
         match field.r#type() {
             Type::Float => Cow::Borrowed("float"),
@@ -1077,6 +1121,12 @@ fn can_pack(field: &FieldDescriptorProto) -> bool {
     )
 }
 
+enum EnumRepr {
+    Int,
+    Closed,
+    Open,
+}
+
 struct EnumVariantMapping<'a> {
     path_idx: usize,
     proto_name: &'a str,

prost-build/src/config.rs

@@ -37,6 +37,7 @@ pub struct Config {
     pub(crate) enum_attributes: PathMap<String>,
     pub(crate) field_attributes: PathMap<String>,
     pub(crate) boxed: PathMap<()>,
+    pub(crate) typed_enum_fields: PathMap<()>,
     pub(crate) prost_types: bool,
     pub(crate) strip_enum_prefix: bool,
     pub(crate) out_dir: Option<PathBuf>,
@@ -374,6 +375,30 @@ impl Config {
         self
     }
 
+    /// Represent Protobuf enum types encountered in matched fields with types
+    /// bound to their corresponding Rust enum types, rather than the default `i32`.
+    ///
+    /// Depending on the proto file syntax, the representation type can be:
+    /// * For closed enums (in proto2), the corresponding Rust enum type.
+    /// * For open enums (in proto3), the Rust enum type wrapped in [`OpenEnum`](prost::OpenEnum).
+    ///
+    /// # Arguments
+    ///
+    /// **`path`** - a path matching any number of fields. These fields will get the type-checked
+    /// enum representation.
+    /// For details about matching fields see [`btree_map`](#method.btree_map).
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # let mut config = prost_build::Config::new();
+    /// config.typed_enum_fields(".my_messages");
+    /// ```
+    pub fn typed_enum_fields(&mut self, path: impl AsRef<str>) -> &mut Self {
+        self.typed_enum_fields.insert(path.as_ref().to_owned(), ());
+        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);
@@ -1179,6 +1204,7 @@ impl default::Default for Config {
             enum_attributes: PathMap::default(),
             field_attributes: PathMap::default(),
             boxed: PathMap::default(),
+            typed_enum_fields: PathMap::default(),
             prost_types: true,
             strip_enum_prefix: true,
             out_dir: None,

prost-build/src/context.rs

@@ -168,6 +168,15 @@ impl<'a> Context<'a> {
         false
     }
 
+    /// Returns `true` if the named message field should be generated using the
+    /// type-safe `OpenEnum` representation.
+    pub fn is_typed_enum_field(&self, fq_message_name: &str, field_name: &str) -> bool {
+        self.config
+            .typed_enum_fields
+            .get_first_field(fq_message_name, field_name)
+            .is_some()
+    }
+
     /// Returns `true` if this message can automatically derive Copy trait.
     pub fn can_message_derive_copy(&self, fq_message_name: &str) -> bool {
         assert_eq!(".", &fq_message_name[..1]);