Merge pull request #808 from s3rius/feature/custom-hashers

Added custom hashers support for HashSet and HashMap.
scylladb · Sep 20, 2023 · ceba60a · ceba60a
2 parents e9af9a1 + 30c6e12
commit ceba60a
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 15 deletions.
diff --git a/scylla-cql/src/frame/response/cql_to_rust.rs b/scylla-cql/src/frame/response/cql_to_rust.rs
@@ -4,7 +4,7 @@ use bigdecimal::BigDecimal;
 use chrono::{DateTime, Duration, NaiveDate, TimeZone, Utc};
 use num_bigint::BigInt;
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
-use std::hash::Hash;
+use std::hash::{BuildHasher, Hash};
 use std::net::IpAddr;
 use thiserror::Error;
 use uuid::Uuid;
@@ -200,27 +200,29 @@ impl<T: FromCqlVal<CqlValue>> FromCqlVal<CqlValue> for Vec<T> {
     }
 }
 
-impl<T1: FromCqlVal<CqlValue> + Eq + Hash, T2: FromCqlVal<CqlValue>> FromCqlVal<CqlValue>
-    for HashMap<T1, T2>
+impl<T1: FromCqlVal<CqlValue> + Eq + Hash, T2: FromCqlVal<CqlValue>, T3: BuildHasher + Default>
+    FromCqlVal<CqlValue> for HashMap<T1, T2, T3>
 {
     fn from_cql(cql_val: CqlValue) -> Result<Self, FromCqlValError> {
         let vec = cql_val.into_pair_vec().ok_or(FromCqlValError::BadCqlType)?;
-        let mut res = HashMap::with_capacity(vec.len());
+        let mut res = HashMap::with_capacity_and_hasher(vec.len(), T3::default());
         for (key, value) in vec {
             res.insert(T1::from_cql(key)?, T2::from_cql(value)?);
         }
         Ok(res)
     }
 }
 
-impl<T: FromCqlVal<CqlValue> + Eq + Hash> FromCqlVal<CqlValue> for HashSet<T> {
+impl<T: FromCqlVal<CqlValue> + Eq + Hash, S: BuildHasher + Default> FromCqlVal<CqlValue>
+    for HashSet<T, S>
+{
     fn from_cql(cql_val: CqlValue) -> Result<Self, FromCqlValError> {
         cql_val
             .into_vec()
             .ok_or(FromCqlValError::BadCqlType)?
             .into_iter()
             .map(T::from_cql)
-            .collect::<Result<HashSet<T>, FromCqlValError>>()
+            .collect::<Result<HashSet<T, S>, FromCqlValError>>()
     }
 }
 

diff --git a/scylla-cql/src/frame/value.rs b/scylla-cql/src/frame/value.rs
@@ -8,6 +8,7 @@ use num_bigint::BigInt;
 use std::borrow::Cow;
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use std::convert::TryInto;
+use std::hash::BuildHasher;
 use std::net::IpAddr;
 use thiserror::Error;
 use uuid::Uuid;
@@ -662,13 +663,13 @@ fn serialize_list_or_set<'a, V: 'a + Value>(
     Ok(())
 }
 
-impl<V: Value> Value for HashSet<V> {
+impl<V: Value, S: BuildHasher + Default> Value for HashSet<V, S> {
     fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), ValueTooBig> {
         serialize_list_or_set(self.iter(), self.len(), buf)
     }
 }
 
-impl<K: Value, V: Value> Value for HashMap<K, V> {
+impl<K: Value, V: Value, S: BuildHasher> Value for HashMap<K, V, S> {
     fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), ValueTooBig> {
         serialize_map(self.iter(), self.len(), buf)
     }
@@ -851,9 +852,9 @@ impl<T: Value> ValueList for Vec<T> {
 }
 
 // Implement ValueList for maps, which serializes named values
-macro_rules! impl_value_list_for_map {
-    ($map_type:ident, $key_type:ty) => {
-        impl<T: Value> ValueList for $map_type<$key_type, T> {
+macro_rules! impl_value_list_for_btree_map {
+    ($key_type:ty) => {
+        impl<T: Value> ValueList for BTreeMap<$key_type, T> {
             fn serialized(&self) -> SerializedResult<'_> {
                 let mut result = SerializedValues::with_capacity(self.len());
                 for (key, val) in self {
@@ -866,10 +867,26 @@ macro_rules! impl_value_list_for_map {
     };
 }
 
-impl_value_list_for_map!(HashMap, String);
-impl_value_list_for_map!(HashMap, &str);
-impl_value_list_for_map!(BTreeMap, String);
-impl_value_list_for_map!(BTreeMap, &str);
+// Implement ValueList for maps, which serializes named values
+macro_rules! impl_value_list_for_hash_map {
+    ($key_type:ty) => {
+        impl<T: Value, S: BuildHasher> ValueList for HashMap<$key_type, T, S> {
+            fn serialized(&self) -> SerializedResult<'_> {
+                let mut result = SerializedValues::with_capacity(self.len());
+                for (key, val) in self {
+                    result.add_named_value(key, val)?;
+                }
+
+                Ok(Cow::Owned(result))
+            }
+        }
+    };
+}
+
+impl_value_list_for_hash_map!(String);
+impl_value_list_for_hash_map!(&str);
+impl_value_list_for_btree_map!(String);
+impl_value_list_for_btree_map!(&str);
 
 // Implement ValueList for tuples of Values of size up to 16