Codegen: option to use BTreeMap as map representation

As of now, protobuf maps are represented with [std::collections::HashMap](https://doc.rust-lang.org/stable/std/collections/struct.HashMap.html), which serves as a robust default. In rarer instances, opting for a different implementation, such as BTreeMap, might be desirable, e.g. for deterministic serialization. This change * adds `btreemaps` codegen option to generate [std::collections::BTreeMap](https://doc.rust-lang.org/std/collections/struct.BTreeMap.html) for maps representation.
stepancheg · Jun 26, 2024 · 263ee76 · 263ee76
1 parent eb5a051
commit 263ee76
Show file tree

Hide file tree

Showing 12 changed files with 238 additions and 100 deletions.
diff --git a/proto/rustproto.proto b/proto/rustproto.proto
@@ -21,6 +21,8 @@ extend google.protobuf.FileOptions {
     optional bool tokio_bytes_for_string_all = 17012;
     // When false, `#[non_exhaustive]` is not generated for `oneof` fields.
     optional bool oneofs_non_exhaustive_all = 17013;
+    // When true, generate `BTreeMap` instead of `HashMap` for map fields.
+    optional bool btreemap_all = 17014;
 
     // When true, will only generate codes that works with lite runtime.
     optional bool lite_runtime_all = 17035;
@@ -37,6 +39,8 @@ extend google.protobuf.MessageOptions {
     optional bool tokio_bytes_for_string = 17012;
     // When false, `#[non_exhaustive]` is not generated for `oneof` fields.
     optional bool oneofs_non_exhaustive = 17013;
+    // When true, generate `BTreeMap` instead of `HashMap` for map fields.
+    optional bool btreemap = 17014;
 }
 
 extend google.protobuf.FieldOptions {
@@ -50,4 +54,6 @@ extend google.protobuf.FieldOptions {
     optional bool tokio_bytes_for_string_field = 17012;
     // When false, `#[non_exhaustive]` is not generated for `oneof` fields.
     optional bool oneofs_non_exhaustive_field = 17013;
+    // When true, generate `BTreeMap` instead of `HashMap` for map fields.
+    optional bool btreemap_field = 17014;
 }
diff --git a/protobuf-codegen/src/customize/mod.rs b/protobuf-codegen/src/customize/mod.rs
@@ -111,6 +111,8 @@ pub struct Customize {
     /// Used internally to generate protos bundled in protobuf crate
     /// like `descriptor.proto`
     pub(crate) inside_protobuf: Option<bool>,
+    /// When true, protobuf maps are represented with `std::collections::BTreeMap`
+    pub(crate) btreemap: Option<bool>,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -178,6 +180,14 @@ impl Customize {
         self
     }
 
+    /// Use btreemaps for maps representation
+    pub fn btreemaps(self, use_btreemaps: bool) -> Self {
+        Self {
+            btreemap: Some(use_btreemaps),
+            ..self
+        }
+    }
+
     /// Update fields of self with fields defined in other customize
     pub fn update_with(&mut self, that: &Customize) {
         if let Some(v) = &that.before {
@@ -204,6 +214,9 @@ impl Customize {
         if let Some(v) = that.inside_protobuf {
             self.inside_protobuf = Some(v);
         }
+        if let Some(v) = that.btreemap {
+            self.btreemap = Some(v);
+        }
     }
 
     /// Update unset fields of self with fields from other customize
@@ -243,6 +256,8 @@ impl Customize {
                 r.lite_runtime = Some(parse_bool(v)?);
             } else if n == "gen_mod_rs" {
                 r.gen_mod_rs = Some(parse_bool(v)?);
+            } else if n == "btreemap" {
+                r.btreemap = Some(parse_bool(v)?);
             } else if n == "inside_protobuf" {
                 r.inside_protobuf = Some(parse_bool(v)?);
             } else if n == "lite" {

diff --git a/protobuf-codegen/src/customize/rustproto_proto.rs b/protobuf-codegen/src/customize/rustproto_proto.rs
@@ -13,6 +13,7 @@ pub(crate) fn customize_from_rustproto_for_message(source: &MessageOptions) -> C
     let tokio_bytes = rustproto::exts::tokio_bytes.get(source);
     let tokio_bytes_for_string = rustproto::exts::tokio_bytes_for_string.get(source);
     let oneofs_non_exhaustive = rustproto::exts::oneofs_non_exhaustive.get(source);
+    let btreemap = rustproto::exts::btreemap.get(source);
     let lite_runtime = None;
     let gen_mod_rs = None;
     let inside_protobuf = None;
@@ -26,6 +27,7 @@ pub(crate) fn customize_from_rustproto_for_message(source: &MessageOptions) -> C
         lite_runtime,
         gen_mod_rs,
         inside_protobuf,
+        btreemap,
     }
 }
 
@@ -40,6 +42,7 @@ pub(crate) fn customize_from_rustproto_for_field(source: &FieldOptions) -> Custo
     let tokio_bytes = rustproto::exts::tokio_bytes_field.get(source);
     let tokio_bytes_for_string = rustproto::exts::tokio_bytes_for_string_field.get(source);
     let oneofs_non_exhaustive = rustproto::exts::oneofs_non_exhaustive_field.get(source);
+    let btreemap = rustproto::exts::btreemap_field.get(source);
     let lite_runtime = None;
     let gen_mod_rs = None;
     let inside_protobuf = None;
@@ -53,6 +56,7 @@ pub(crate) fn customize_from_rustproto_for_field(source: &FieldOptions) -> Custo
         lite_runtime,
         gen_mod_rs,
         inside_protobuf,
+        btreemap,
     }
 }
 
@@ -64,6 +68,7 @@ pub(crate) fn customize_from_rustproto_for_file(source: &FileOptions) -> Customi
     let tokio_bytes_for_string = rustproto::exts::tokio_bytes_for_string_all.get(source);
     let oneofs_non_exhaustive = rustproto::exts::oneofs_non_exhaustive_all.get(source);
     let lite_runtime = rustproto::exts::lite_runtime_all.get(source);
+    let btreemap = rustproto::exts::btreemap_all.get(source);
     let gen_mod_rs = None;
     let inside_protobuf = None;
     Customize {
@@ -76,5 +81,6 @@ pub(crate) fn customize_from_rustproto_for_file(source: &FileOptions) -> Customi
         lite_runtime,
         inside_protobuf,
         gen_mod_rs,
+        btreemap,
     }
 }
diff --git a/protobuf-codegen/src/gen/field/mod.rs b/protobuf-codegen/src/gen/field/mod.rs
@@ -302,6 +302,12 @@ impl<'a> FieldGen<'a> {
     pub(crate) fn full_storage_type(&self, reference: &FileAndMod) -> RustType {
         match self.kind {
             FieldKind::Repeated(ref repeated) => repeated.rust_type(reference),
+            FieldKind::Map(MapField {
+                ref key, ref value, ..
+            }) if self.customize.btreemap == Some(true) => RustType::BTreeMap(
+                Box::new(key.rust_storage_elem_type(reference)),
+                Box::new(value.rust_storage_elem_type(reference)),
+            ),
             FieldKind::Map(MapField {
                 ref key, ref value, ..
             }) => RustType::HashMap(

diff --git a/protobuf-codegen/src/gen/rust_types_values.rs b/protobuf-codegen/src/gen/rust_types_values.rs
@@ -34,6 +34,7 @@ pub(crate) enum RustType {
     Bool,
     Vec(Box<RustType>),
     HashMap(Box<RustType>, Box<RustType>),
+    BTreeMap(Box<RustType>, Box<RustType>),
     String,
     // [T], not &[T]
     Slice(Box<RustType>),
@@ -75,6 +76,11 @@ impl RustType {
                 key.to_code(customize),
                 value.to_code(customize)
             ),
+            RustType::BTreeMap(ref key, ref value) => format!(
+                "::std::collections::BTreeMap<{}, {}>",
+                key.to_code(customize),
+                value.to_code(customize)
+            ),
             RustType::String => format!("::std::string::String"),
             RustType::Slice(ref param) => format!("[{}]", param.to_code(customize)),
             RustType::Str => format!("str"),
@@ -221,6 +227,7 @@ impl RustType {
             RustType::Bool => "false".to_owned(),
             RustType::Vec(..) => EXPR_VEC_NEW.to_owned(),
             RustType::HashMap(..) => "::std::collections::HashMap::new()".to_owned(),
+            RustType::BTreeMap(..) => "::std::collections::BTreeMap::new()".to_owned(),
             RustType::String => "::std::string::String::new()".to_owned(),
             RustType::Bytes => "::bytes::Bytes::new()".to_owned(),
             RustType::Chars => format!("{}::Chars::new()", protobuf_crate_path(customize)),
@@ -266,7 +273,8 @@ impl RustType {
             | RustType::Chars
             | RustType::String
             | RustType::MessageField(..)
-            | RustType::HashMap(..) => format!("{}.clear()", v),
+            | RustType::HashMap(..)
+            | RustType::BTreeMap(..) => format!("{}.clear()", v),
             RustType::Bool
             | RustType::Float(..)
             | RustType::Int(..)

diff --git a/protobuf-parse/src/proto/rustproto.proto b/protobuf-parse/src/proto/rustproto.proto
@@ -21,6 +21,8 @@ extend google.protobuf.FileOptions {
     optional bool tokio_bytes_for_string_all = 17012;
     // When false, `#[non_exhaustive]` is not generated for `oneof` fields.
     optional bool oneofs_non_exhaustive_all = 17013;
+    // When true, generate `BTreeMap` instead of `HashMap` for map fields.
+    optional bool btreemap_all = 17014;
 
     // When true, will only generate codes that works with lite runtime.
     optional bool lite_runtime_all = 17035;
@@ -37,6 +39,8 @@ extend google.protobuf.MessageOptions {
     optional bool tokio_bytes_for_string = 17012;
     // When false, `#[non_exhaustive]` is not generated for `oneof` fields.
     optional bool oneofs_non_exhaustive = 17013;
+    // When true, generate `BTreeMap` instead of `HashMap` for map fields.
+    optional bool btreemap = 17014;
 }
 
 extend google.protobuf.FieldOptions {
@@ -50,4 +54,6 @@ extend google.protobuf.FieldOptions {
     optional bool tokio_bytes_for_string_field = 17012;
     // When false, `#[non_exhaustive]` is not generated for `oneof` fields.
     optional bool oneofs_non_exhaustive_field = 17013;
+    // When true, generate `BTreeMap` instead of `HashMap` for map fields.
+    optional bool btreemap_field = 17014;
 }