feat(versionable): impl Versionize for HashSet/HashMap

zama-ai · Jul 31, 2024 · dc39d69 · dc39d69
1 parent 351fc47
commit dc39d69
Showing 1 changed file with 61 additions and 0 deletions.
diff --git a/utils/tfhe-versionable/src/lib.rs b/utils/tfhe-versionable/src/lib.rs
@@ -11,6 +11,7 @@ pub mod upgrade;
 
 use aligned_vec::{ABox, AVec};
 use num_complex::Complex;
+use std::collections::{HashMap, HashSet};
 use std::convert::Infallible;
 use std::error::Error;
 use std::fmt::Display;
@@ -568,3 +569,63 @@ impl<T: Unversionize, U: Unversionize, V: Unversionize> Unversionize for (T, U,
 }
 
 impl<T: NotVersioned, U: NotVersioned, V: NotVersioned> NotVersioned for (T, U, V) {}
+
+// converts to `Vec<T::Versioned>` for the versioned type, so we don't have to derive
+// Eq/Hash on it.
+impl<T: Versionize> Versionize for HashSet<T> {
+    type Versioned<'vers> = Vec<T::Versioned<'vers>>
+    where
+        T: 'vers;
+
+    fn versionize(&self) -> Self::Versioned<'_> {
+        self.iter().map(|val| val.versionize()).collect()
+    }
+}
+
+impl<T: VersionizeOwned> VersionizeOwned for HashSet<T> {
+    type VersionedOwned = Vec<T::VersionedOwned>;
+
+    fn versionize_owned(self) -> Self::VersionedOwned {
+        self.into_iter().map(|val| val.versionize_owned()).collect()
+    }
+}
+
+impl<T: Unversionize + std::hash::Hash + Eq> Unversionize for HashSet<T> {
+    fn unversionize(versioned: Self::VersionedOwned) -> Result<Self, UnversionizeError> {
+        versioned
+            .into_iter()
+            .map(|val| T::unversionize(val))
+            .collect()
+    }
+}
+
+// converts to `Vec<(K::Versioned, V::Versioned)>` for the versioned type, so we don't have to
+// derive Eq/Hash on it.
+impl<K: Versionize, V: Versionize> Versionize for HashMap<K, V> {
+    type Versioned<'vers> = Vec<(K::Versioned<'vers>, V::Versioned<'vers>)> where K: 'vers, V: 'vers;
+
+    fn versionize(&self) -> Self::Versioned<'_> {
+        self.iter()
+            .map(|(key, val)| (key.versionize(), val.versionize()))
+            .collect()
+    }
+}
+
+impl<K: VersionizeOwned, V: VersionizeOwned> VersionizeOwned for HashMap<K, V> {
+    type VersionedOwned = Vec<(K::VersionedOwned, V::VersionedOwned)>;
+
+    fn versionize_owned(self) -> Self::VersionedOwned {
+        self.into_iter()
+            .map(|(key, val)| (key.versionize_owned(), val.versionize_owned()))
+            .collect()
+    }
+}
+
+impl<K: Unversionize + std::hash::Hash + Eq, V: Unversionize> Unversionize for HashMap<K, V> {
+    fn unversionize(versioned: Self::VersionedOwned) -> Result<Self, UnversionizeError> {
+        versioned
+            .into_iter()
+            .map(|(key, val)| Ok((K::unversionize(key)?, V::unversionize(val)?)))
+            .collect()
+    }
+}