Added custom hashers support for HashSet and HashMap.

Signed-off-by: Pavel Kirilin <[email protected]>
scylladb · Sep 9, 2023 · cea41c0 · cea41c0
1 parent cabb6f1
commit cea41c0
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 10 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;
@@ -213,14 +213,16 @@ impl<T1: FromCqlVal<CqlValue> + Eq + Hash, T2: FromCqlVal<CqlValue>> FromCqlVal<
     }
 }
 
-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,7 +852,7 @@ impl<T: Value> ValueList for Vec<T> {
 }
 
 // Implement ValueList for maps, which serializes named values
-macro_rules! impl_value_list_for_map {
+macro_rules! impl_value_list_for_btree_map {
     ($map_type:ident, $key_type:ty) => {
         impl<T: Value> ValueList for $map_type<$key_type, T> {
             fn serialized(&self) -> SerializedResult<'_> {
@@ -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 {
+    ($map_type:ident, $key_type:ty) => {
+        impl<T: Value, S: BuildHasher> ValueList for $map_type<$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!(HashMap, String);
+impl_value_list_for_hash_map!(HashMap, &str);
+impl_value_list_for_btree_map!(BTreeMap, String);
+impl_value_list_for_btree_map!(BTreeMap, &str);
 
 // Implement ValueList for tuples of Values of size up to 16