Codegen: option to use BTreeMap as map representation

Cargo.toml

@@ -16,6 +16,7 @@ members = [
     "test-crates/perftest/bytes",
     "test-crates/perftest/misc",
     "test-crates/perftest/vs-cxx",
+    "test-crates/protobuf-codegen-btreemap-test",
     "test-crates/protobuf-codegen-identical-test",
     "test-crates/protobuf-codegen-protoc-test",
     "test-crates/protobuf-codegen-pure-test",

protobuf-codegen/src/customize/mod.rs

@@ -109,6 +109,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) btreemaps: Option<bool>,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -171,6 +173,14 @@ impl Customize {
         self
     }
 
+    /// Use btreemaps for maps representation
+    pub fn btreemaps(self, use_btreemaps: bool) -> Self {
+        Self {
+            btreemaps: 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 {
@@ -197,6 +207,9 @@ impl Customize {
         if let Some(v) = that.inside_protobuf {
             self.inside_protobuf = Some(v);
         }
+        if let Some(v) = that.btreemaps {
+            self.btreemaps = Some(v);
+        }
     }
 
     /// Update unset fields of self with fields from other customize

protobuf-codegen/src/customize/rustproto_proto.rs

@@ -15,6 +15,7 @@ pub(crate) fn customize_from_rustproto_for_message(source: &MessageOptions) -> C
     let lite_runtime = None;
     let gen_mod_rs = None;
     let inside_protobuf = None;
+    let btreemaps = None;
     Customize {
         before,
         generate_accessors,
@@ -24,6 +25,7 @@ pub(crate) fn customize_from_rustproto_for_message(source: &MessageOptions) -> C
         lite_runtime,
         gen_mod_rs,
         inside_protobuf,
+        btreemaps,
     }
 }
 
@@ -40,6 +42,7 @@ pub(crate) fn customize_from_rustproto_for_field(source: &FieldOptions) -> Custo
     let lite_runtime = None;
     let gen_mod_rs = None;
     let inside_protobuf = None;
+    let btreemaps = None;
     Customize {
         before,
         generate_accessors,
@@ -49,6 +52,7 @@ pub(crate) fn customize_from_rustproto_for_field(source: &FieldOptions) -> Custo
         lite_runtime,
         gen_mod_rs,
         inside_protobuf,
+        btreemaps,
     }
 }
 
@@ -61,6 +65,7 @@ pub(crate) fn customize_from_rustproto_for_file(source: &FileOptions) -> Customi
     let lite_runtime = rustproto::exts::lite_runtime_all.get(source);
     let gen_mod_rs = None;
     let inside_protobuf = None;
+    let btreemaps = None;
     Customize {
         before,
         generate_accessors,
@@ -70,5 +75,6 @@ pub(crate) fn customize_from_rustproto_for_file(source: &FileOptions) -> Customi
         lite_runtime,
         inside_protobuf,
         gen_mod_rs,
+        btreemaps,
     }
 }

protobuf-codegen/src/gen/field/accessor.rs

@@ -66,7 +66,7 @@ impl FieldGen<'_> {
         let MapField { .. } = map_field;
         AccessorFn {
             name: "make_map_simpler_accessor".to_owned(),
-            type_params: vec![format!("_"), format!("_")],
+            type_params: vec![format!("_")],
             callback_params: self.make_accessor_fns_lambda(),
         }
     }

protobuf-codegen/src/gen/field/mod.rs

@@ -276,7 +276,7 @@ impl<'a> FieldGen<'a> {
             FieldKind::Repeated(ref repeated) => repeated.rust_type(reference),
             FieldKind::Map(MapField {
                 ref key, ref value, ..
-            }) => RustType::HashMap(
+            }) => RustType::Map(
                 Box::new(key.rust_storage_elem_type(reference)),
                 Box::new(value.rust_storage_elem_type(reference)),
             ),

protobuf-codegen/src/gen/rust_types_values.rs

@@ -33,7 +33,7 @@ pub(crate) enum RustType {
     Float(u32),
     Bool,
     Vec(Box<RustType>),
-    HashMap(Box<RustType>, Box<RustType>),
+    Map(Box<RustType>, Box<RustType>),
     String,
     // [T], not &[T]
     Slice(Box<RustType>),
@@ -70,11 +70,19 @@ impl RustType {
             RustType::Float(bits) => format!("f{}", bits),
             RustType::Bool => format!("bool"),
             RustType::Vec(ref param) => format!("::std::vec::Vec<{}>", param.to_code(customize)),
-            RustType::HashMap(ref key, ref value) => format!(
-                "::std::collections::HashMap<{}, {}>",
-                key.to_code(customize),
-                value.to_code(customize)
-            ),
+            RustType::Map(ref key, ref value) => {
+                let ty = if let Some(true) = customize.btreemaps {
+                    "::std::collections::BTreeMap"
+                } else {
+                    "::std::collections::HashMap"
+                };
+                format!(
+                    "{}<{}, {}>",
+                    ty,
+                    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"),
@@ -220,7 +228,14 @@ impl RustType {
             RustType::Float(..) => "0.".to_string(),
             RustType::Bool => "false".to_string(),
             RustType::Vec(..) => EXPR_VEC_NEW.to_string(),
-            RustType::HashMap(..) => "::std::collections::HashMap::new()".to_string(),
+            RustType::Map(..) => {
+                let ty = if let Some(true) = customize.btreemaps {
+                    "::std::collections::BTreeMap"
+                } else {
+                    "::std::collections::HashMap"
+                };
+                format!("{}::new()", ty)
+            }
             RustType::String => "::std::string::String::new()".to_string(),
             RustType::Bytes => "::bytes::Bytes::new()".to_string(),
             RustType::Chars => format!("{}::Chars::new()", protobuf_crate_path(customize)),
@@ -266,7 +281,7 @@ impl RustType {
             | RustType::Chars
             | RustType::String
             | RustType::MessageField(..)
-            | RustType::HashMap(..) => format!("{}.clear()", v),
+            | RustType::Map(..) => format!("{}.clear()", v),
             RustType::Bool
             | RustType::Float(..)
             | RustType::Int(..)

protobuf/src/reflect/acc/v2/map.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 use std::fmt;
 use std::hash::Hash;
 
@@ -8,7 +8,7 @@
 use crate::reflect::acc::FieldAccessor;
 use crate::reflect::map::ReflectMapMut;
 use crate::reflect::map::ReflectMapRef;
-use crate::reflect::runtime_types::RuntimeTypeHashable;
+use crate::reflect::runtime_types::RuntimeTypeMapKey;
 use crate::reflect::runtime_types::RuntimeTypeTrait;
 use crate::reflect::ProtobufValue;
 use crate::reflect::RuntimeType;
@@ -29,21 +29,46 @@
     }
 }
 
-struct MapFieldAccessorImpl<M, K, V>
+struct MapFieldAccessorImpl<M, T>
 where
     M: MessageFull,
-    K: ProtobufValue,
-    V: ProtobufValue,
 {
-    get_field: fn(&M) -> &HashMap<K, V>,
-    mut_field: fn(&mut M) -> &mut HashMap<K, V>,
+    get_field: fn(&M) -> &T,
+    mut_field: fn(&mut M) -> &mut T,
 }
 
-impl<M, K, V> MapFieldAccessor for MapFieldAccessorImpl<M, K, V>
+impl<M, K, V> MapFieldAccessor for MapFieldAccessorImpl<M, HashMap<K, V>>
 where
     M: MessageFull,
     K: ProtobufValue + Eq + Hash,
-    K::RuntimeType: RuntimeTypeHashable,
+    K::RuntimeType: RuntimeTypeMapKey,
+    V: ProtobufValue,
+{
+    fn get_reflect<'a>(&self, m: &'a dyn MessageDyn) -> ReflectMapRef<'a> {
+        let m = m.downcast_ref().unwrap();
+        let map = (self.get_field)(m);
+        ReflectMapRef::new(map)
+    }
+
+    fn mut_reflect<'a>(&self, m: &'a mut dyn MessageDyn) -> ReflectMapMut<'a> {
+        let m = m.downcast_mut().unwrap();
+        let map = (self.mut_field)(m);
+        ReflectMapMut::new(map)
+    }
+
+    fn element_type(&self) -> (RuntimeType, RuntimeType) {
+        (
+            K::RuntimeType::runtime_type_box(),
+            V::RuntimeType::runtime_type_box(),
+        )
+    }
+}
+
+impl<M, K, V> MapFieldAccessor for MapFieldAccessorImpl<M, BTreeMap<K, V>>
+where
+    M: MessageFull,
+    K: ProtobufValue + Ord,
+    K::RuntimeType: RuntimeTypeMapKey,
     V: ProtobufValue,
 {
     fn get_reflect<'a>(&self, m: &'a dyn MessageDyn) -> ReflectMapRef<'a> {
@@ -67,21 +92,19 @@
 }
 
 /// Make accessor for map field
-pub fn make_map_simpler_accessor<M, K, V>(
+pub fn make_map_simpler_accessor<M, T>(
     name: &'static str,
-    get_field: for<'a> fn(&'a M) -> &'a HashMap<K, V>,
-    mut_field: for<'a> fn(&'a mut M) -> &'a mut HashMap<K, V>,
+    get_field: for<'a> fn(&'a M) -> &'a T,
+    mut_field: for<'a> fn(&'a mut M) -> &'a mut T,
 ) -> FieldAccessor
 where
-    M: MessageFull + 'static,
-    K: ProtobufValue + Hash + Eq,
-    K::RuntimeType: RuntimeTypeHashable,
-    V: ProtobufValue,
+    M: MessageFull,
+    MapFieldAccessorImpl<M, T>: MapFieldAccessor,
 {
     FieldAccessor::new(
         name,
         AccessorV2::Map(MapFieldAccessorHolder {
-            accessor: Box::new(MapFieldAccessorImpl::<M, K, V> {
+            accessor: Box::new(MapFieldAccessorImpl::<M, T> {
                 get_field,
                 mut_field,
             }),

protobuf/src/reflect/map/generated.rs

@@ -1,11 +1,11 @@
-use std::collections::hash_map;
-use std::collections::HashMap;
+use std::collections::{btree_map, hash_map};
+use std::collections::{BTreeMap, HashMap};
 use std::hash::Hash;
 
 use crate::reflect::map::ReflectMap;
 use crate::reflect::map::ReflectMapIter;
 use crate::reflect::map::ReflectMapIterTrait;
-use crate::reflect::runtime_types::RuntimeTypeHashable;
+use crate::reflect::runtime_types::RuntimeTypeMapKey;
 use crate::reflect::runtime_types::RuntimeTypeTrait;
 use crate::reflect::ProtobufValue;
 use crate::reflect::ReflectValueBox;
@@ -16,10 +16,12 @@ impl<K, V> ReflectMap for HashMap<K, V>
 where
     K: ProtobufValue + Eq + Hash,
     V: ProtobufValue,
-    K::RuntimeType: RuntimeTypeHashable,
+    K::RuntimeType: RuntimeTypeMapKey,
 {
     fn reflect_iter<'a>(&'a self) -> ReflectMapIter<'a> {
-        ReflectMapIter::new(GeneratedMapIterImpl::<'a, K, V> { iter: self.iter() })
+        ReflectMapIter::new(GeneratedMapIterImpl::<'a, K, V, hash_map::Iter<'a, K, V>> {
+            iter: self.iter(),
+        })
     }
 
     fn len(&self) -> usize {
@@ -31,7 +33,7 @@ where
     }
 
     fn get<'a>(&'a self, key: ReflectValueRef) -> Option<ReflectValueRef<'a>> {
-        <K::RuntimeType as RuntimeTypeHashable>::hash_map_get(self, key).map(V::RuntimeType::as_ref)
+        <K::RuntimeType as RuntimeTypeMapKey>::hash_map_get(self, key).map(V::RuntimeType::as_ref)
     }
 
     fn insert(&mut self, key: ReflectValueBox, value: ReflectValueBox) {
@@ -53,12 +55,55 @@ where
     }
 }
 
-struct GeneratedMapIterImpl<'a, K: Eq + Hash + 'static, V: 'static> {
-    iter: hash_map::Iter<'a, K, V>,
+impl<K, V> ReflectMap for BTreeMap<K, V>
+where
+    K: ProtobufValue + Ord,
+    V: ProtobufValue,
+    K::RuntimeType: RuntimeTypeMapKey,
+{
+    fn reflect_iter<'a>(&'a self) -> ReflectMapIter<'a> {
+        ReflectMapIter::new(
+            GeneratedMapIterImpl::<'a, K, V, btree_map::Iter<'a, K, V>> { iter: self.iter() },
+        )
+    }
+
+    fn len(&self) -> usize {
+        BTreeMap::len(self)
+    }
+
+    fn is_empty(&self) -> bool {
+        self.is_empty()
+    }
+
+    fn get<'a>(&'a self, key: ReflectValueRef) -> Option<ReflectValueRef<'a>> {
+        <K::RuntimeType as RuntimeTypeMapKey>::btree_map_get(self, key).map(V::RuntimeType::as_ref)
+    }
+
+    fn insert(&mut self, key: ReflectValueBox, value: ReflectValueBox) {
+        let key: K = key.downcast().expect("wrong key type");
+        let value: V = value.downcast().expect("wrong value type");
+        self.insert(key, value);
+    }
+
+    fn clear(&mut self) {
+        self.clear();
+    }
+
+    fn key_type(&self) -> RuntimeType {
+        K::RuntimeType::runtime_type_box()
+    }
+
+    fn value_type(&self) -> RuntimeType {
+        V::RuntimeType::runtime_type_box()
+    }
+}
+
+struct GeneratedMapIterImpl<'a, K: 'static, V: 'static, I: Iterator<Item = (&'a K, &'a V)>> {
+    iter: I,
 }
 
-impl<'a, K: ProtobufValue + Eq + Hash, V: ProtobufValue> ReflectMapIterTrait<'a>
-    for GeneratedMapIterImpl<'a, K, V>
+impl<'a, K: ProtobufValue, V: ProtobufValue, I: Iterator<Item = (&'a K, &'a V)>>
+    ReflectMapIterTrait<'a> for GeneratedMapIterImpl<'a, K, V, I>
 {
     fn next(&mut self) -> Option<(ReflectValueRef<'a>, ReflectValueRef<'a>)> {
         match self.iter.next() {

protobuf/src/reflect/map/mod.rs

@@ -12,7 +12,7 @@ use crate::reflect::RuntimeType;
 mod empty;
 mod generated;
 
-/// Implemented for `HashMap` with appropriate keys and values
+/// Implemented for `HashMap`, `BTreeMap` with appropriate keys and values
 pub(crate) trait ReflectMap: Debug + Send + Sync + 'static {
     fn reflect_iter(&self) -> ReflectMapIter;