From 2602c9e1b3768065910fa6284d738265b159dccc Mon Sep 17 00:00:00 2001
From: Arthur Meyre <arthur.meyre@zama.ai>
Date: Mon, 29 Sep 2025 17:35:23 +0200
Subject: [PATCH] fix(hlapi): clear rerand metadata once rerand is done

---
 tfhe/src/high_level_api/booleans/base.rs              |  2 ++
 tfhe/src/high_level_api/integers/signed/base.rs       |  2 ++
 tfhe/src/high_level_api/integers/unsigned/base.rs     |  2 ++
 tfhe/src/high_level_api/re_randomization.rs           |  4 ++++
 tfhe/src/high_level_api/tag.rs                        | 10 ++++++++++
 tfhe/src/high_level_api/tests/cpk_re_randomization.rs |  6 ++++++
 6 files changed, 26 insertions(+)

diff --git a/tfhe/src/high_level_api/booleans/base.rs b/tfhe/src/high_level_api/booleans/base.rs
index 1e153ae41..60312d7bd 100644
--- a/tfhe/src/high_level_api/booleans/base.rs
+++ b/tfhe/src/high_level_api/booleans/base.rs
@@ -2071,6 +2071,8 @@ impl ReRandomize for FheBool {
                     seed,
                 )?;
 
+                self.re_randomization_metadata_mut().clear();
+
                 Ok(())
             }
             #[cfg(feature = "gpu")]
diff --git a/tfhe/src/high_level_api/integers/signed/base.rs b/tfhe/src/high_level_api/integers/signed/base.rs
index 393f46161..ea48644a2 100644
--- a/tfhe/src/high_level_api/integers/signed/base.rs
+++ b/tfhe/src/high_level_api/integers/signed/base.rs
@@ -1162,6 +1162,8 @@ where
                     seed,
                 )?;
 
+                self.re_randomization_metadata_mut().clear();
+
                 Ok(())
             }
             #[cfg(feature = "gpu")]
diff --git a/tfhe/src/high_level_api/integers/unsigned/base.rs b/tfhe/src/high_level_api/integers/unsigned/base.rs
index dc318a1f9..61612b182 100644
--- a/tfhe/src/high_level_api/integers/unsigned/base.rs
+++ b/tfhe/src/high_level_api/integers/unsigned/base.rs
@@ -1672,6 +1672,8 @@ where
                     seed,
                 )?;
 
+                self.re_randomization_metadata_mut().clear();
+
                 Ok(())
             }
             #[cfg(feature = "gpu")]
diff --git a/tfhe/src/high_level_api/re_randomization.rs b/tfhe/src/high_level_api/re_randomization.rs
index 3a8e8f65b..61b20a99f 100644
--- a/tfhe/src/high_level_api/re_randomization.rs
+++ b/tfhe/src/high_level_api/re_randomization.rs
@@ -231,4 +231,8 @@ impl ReRandomizationMetadata {
     pub fn set_data(&mut self, data: &[u8]) {
         self.inner.set_data(data);
     }
+
+    pub fn clear(&mut self) {
+        self.inner.clear();
+    }
 }
diff --git a/tfhe/src/high_level_api/tag.rs b/tfhe/src/high_level_api/tag.rs
index cc4a72f8a..fe89e3a2a 100644
--- a/tfhe/src/high_level_api/tag.rs
+++ b/tfhe/src/high_level_api/tag.rs
@@ -174,6 +174,16 @@ impl SmallVec {
         self.set_data(le_bytes.as_slice());
     }
 
+    /// Clears the vector, removing all values.
+    ///
+    /// Note that this method has no effect on the allocated capacity of the vector.
+    pub fn clear(&mut self) {
+        match self {
+            Self::Stack { bytes: _, len } => *len = 0,
+            Self::Heap(items) => items.clear(),
+        }
+    }
+
     // Creates a SmallVec from the vec, but, only re-uses the vec
     // if its len would not fit on the stack part.
     //
diff --git a/tfhe/src/high_level_api/tests/cpk_re_randomization.rs b/tfhe/src/high_level_api/tests/cpk_re_randomization.rs
index 5f9bf35bf..086a58bb5 100644
--- a/tfhe/src/high_level_api/tests/cpk_re_randomization.rs
+++ b/tfhe/src/high_level_api/tests/cpk_re_randomization.rs
@@ -78,8 +78,10 @@ fn test_re_rand() {
         let mut seed_gen = re_rand_context.finalize();
 
         a.re_randomize(&cpk, seed_gen.next_seed().unwrap()).unwrap();
+        assert!(a.re_randomization_metadata().data().is_empty());
 
         b.re_randomize(&cpk, seed_gen.next_seed().unwrap()).unwrap();
+        assert!(b.re_randomization_metadata().data().is_empty());
 
         let c = a + b;
         let dec: u64 = c.decrypt(&cks);
@@ -130,8 +132,10 @@ fn test_re_rand() {
         let mut seed_gen = re_rand_context.finalize();
 
         a.re_randomize(&cpk, seed_gen.next_seed().unwrap()).unwrap();
+        assert!(a.re_randomization_metadata().data().is_empty());
 
         b.re_randomize(&cpk, seed_gen.next_seed().unwrap()).unwrap();
+        assert!(b.re_randomization_metadata().data().is_empty());
 
         let c = a + b;
         let dec: i8 = c.decrypt(&cks);
@@ -181,8 +185,10 @@ fn test_re_rand() {
                 let mut seed_gen = re_rand_context.finalize();
 
                 a.re_randomize(&cpk, seed_gen.next_seed().unwrap()).unwrap();
+                assert!(a.re_randomization_metadata().data().is_empty());
 
                 b.re_randomize(&cpk, seed_gen.next_seed().unwrap()).unwrap();
+                assert!(b.re_randomization_metadata().data().is_empty());
 
                 let c = a & b;
                 let dec: bool = c.decrypt(&cks);