OpenEnum: an enum to represent protobuf's enumeration field values

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

@@ -427,15 +427,19 @@ impl CodeGenerator<'_> {
         let type_tag = self.field_type_tag(&field.descriptor);
         self.buf.push_str(&type_tag);
 
-        if type_ == Type::Bytes {
-            let bytes_type = self
-                .config
-                .bytes_type
-                .get_first_field(fq_message_name, field.descriptor.name())
-                .copied()
-                .unwrap_or_default();
-            self.buf
-                .push_str(&format!("={:?}", bytes_type.annotation()));
+        match type_ {
+            Type::Bytes => {
+                let bytes_type = self
+                    .config
+                    .bytes_type
+                    .get_first_field(fq_message_name, field.descriptor.name())
+                    .copied()
+                    .unwrap_or_default();
+                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() {
@@ -555,12 +559,16 @@ impl CodeGenerator<'_> {
         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!(
@@ -639,11 +647,12 @@ impl CodeGenerator<'_> {
 
             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();
@@ -947,13 +956,21 @@ impl CodeGenerator<'_> {
         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", prost_path(self.config)),
@@ -966,6 +983,15 @@ impl CodeGenerator<'_> {
                 .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<{}>",
+                    prost_path(self.config),
+                    self.resolve_ident(field.type_name())
+                ),
+                EnumRepr::Closed => self.resolve_ident(field.type_name()),
+            },
         }
     }
 
@@ -1007,6 +1033,26 @@ impl CodeGenerator<'_> {
             .join("::")
     }
 
+    fn enum_field_repr(&self, fq_message_name: &str, field_name: &str) -> EnumRepr {
+        if self
+            .config
+            .typed_enum_fields
+            .get_first_field(fq_message_name, field_name)
+            .is_some()
+        {
+            // 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"),
@@ -1143,6 +1189,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

@@ -36,6 +36,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>,
@@ -372,6 +373,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);
@@ -1158,6 +1183,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,