From 5b530152fe73a93d4a4add90e56da8092156ba13 Mon Sep 17 00:00:00 2001
From: Nicolas Sarlin <nicolas.sarlin@zama.ai>
Date: Mon, 10 Jun 2024 14:34:52 +0200
Subject: [PATCH] feat(versionable): Add versionize support for aligned-vec
 types

---
 utils/tfhe-versionable/Cargo.toml |  1 +
 utils/tfhe-versionable/src/lib.rs | 48 +++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/utils/tfhe-versionable/Cargo.toml b/utils/tfhe-versionable/Cargo.toml
index 0eaba1b2f..644219d2f 100644
--- a/utils/tfhe-versionable/Cargo.toml
+++ b/utils/tfhe-versionable/Cargo.toml
@@ -27,3 +27,4 @@ toml = "0.8"
 serde = { version = "1.0", features = ["derive"] }
 tfhe-versionable-derive = { version = "0.1.0", path = "../tfhe-versionable-derive" }
 num-complex = { version = "0.4", features = ["serde"] }
+aligned-vec = { version = "0.5", features = ["serde"] }
diff --git a/utils/tfhe-versionable/src/lib.rs b/utils/tfhe-versionable/src/lib.rs
index 4fbb9c487..5b8c9b9db 100644
--- a/utils/tfhe-versionable/src/lib.rs
+++ b/utils/tfhe-versionable/src/lib.rs
@@ -9,6 +9,7 @@
 pub mod derived_traits;
 pub mod upgrade;
 
+use aligned_vec::{ABox, AVec};
 use num_complex::Complex;
 use std::convert::Infallible;
 use std::fmt::Display;
@@ -369,3 +370,50 @@ impl<T: Unversionize> Unversionize for Complex<T> {
 }
 
 impl<T: NotVersioned> NotVersioned for Complex<T> {}
+
+impl<T: Versionize> Versionize for ABox<T> {
+    type Versioned<'vers> = T::Versioned<'vers> where T: 'vers;
+
+    fn versionize(&self) -> Self::Versioned<'_> {
+        self.as_ref().versionize()
+    }
+
+    // Alignment doesn't matter for versioned types
+    type VersionedOwned = Box<T::VersionedOwned>;
+
+    fn versionize_owned(&self) -> Self::VersionedOwned {
+        Box::new(T::versionize_owned(self))
+    }
+}
+
+impl<T: Unversionize> Unversionize for ABox<T>
+where
+    T::VersionedOwned: Clone,
+{
+    fn unversionize(versioned: Self::VersionedOwned) -> Result<Self, UnversionizeError> {
+        Ok(ABox::new(0, T::unversionize((*versioned).to_owned())?))
+    }
+}
+
+impl<T: VersionizeVec> Versionize for AVec<T> {
+    type Versioned<'vers> = T::VersionedSlice<'vers> where T: 'vers;
+
+    fn versionize(&self) -> Self::Versioned<'_> {
+        T::versionize_slice(self)
+    }
+
+    // Alignment doesn't matter for versioned types
+    type VersionedOwned = T::VersionedVec;
+
+    fn versionize_owned(&self) -> Self::VersionedOwned {
+        T::versionize_vec(self)
+    }
+}
+
+impl<T: UnversionizeVec> Unversionize for AVec<T> {
+    fn unversionize(versioned: Self::VersionedOwned) -> Result<Self, UnversionizeError> {
+        T::unversionize_vec(versioned).map(|unver| AVec::from_iter(0, unver))
+    }
+}
+
+impl<T: NotVersioned + Clone + Serialize + DeserializeOwned> NotVersioned for AVec<T> {}