From b11213c454f90d1664c38bad6cb292586e602499 Mon Sep 17 00:00:00 2001
From: Pro7ech <jeanphilippe.bossuat@gmail.com>
Date: Wed, 12 Nov 2025 22:21:32 +0100
Subject: [PATCH] Add glwe_normalize_tmp_bytes & update assertion error

---
 poulpy-core/src/operations/glwe.rs                | 15 +++++++++++----
 .../ciphertexts/fhe_uint_prepared.rs              |  6 +++++-
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/poulpy-core/src/operations/glwe.rs b/poulpy-core/src/operations/glwe.rs
index 9802c14..3a2c3b4 100644
--- a/poulpy-core/src/operations/glwe.rs
+++ b/poulpy-core/src/operations/glwe.rs
@@ -2,8 +2,8 @@ use poulpy_hal::{
     api::{
         BivariateTensoring, ModuleN, ScratchTakeBasic, VecZnxAdd, VecZnxAddInplace, VecZnxBigNormalize, VecZnxCopy,
         VecZnxIdftApplyConsume, VecZnxMulXpMinusOne, VecZnxMulXpMinusOneInplace, VecZnxNegate, VecZnxNormalize,
-        VecZnxNormalizeInplace, VecZnxRotate, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub, VecZnxSubInplace,
-        VecZnxSubNegateInplace, VecZnxZero,
+        VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub,
+        VecZnxSubInplace, VecZnxSubNegateInplace, VecZnxZero,
     },
     layouts::{Backend, Module, Scratch, VecZnx, VecZnxBig, ZnxInfos},
     reference::vec_znx::vec_znx_rotate_inplace_tmp_bytes,
@@ -389,12 +389,19 @@ impl GLWE<Vec<u8>> {
     }
 }
 
-impl<BE: Backend> GLWENormalize<BE> for Module<BE> where Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE> {}
+impl<BE: Backend> GLWENormalize<BE> for Module<BE> where
+    Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE> + VecZnxNormalizeTmpBytes
+{
+}
 
 pub trait GLWENormalize<BE: Backend>
 where
-    Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE>,
+    Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE> + VecZnxNormalizeTmpBytes,
 {
+    fn glwe_normalize_tmp_bytes(&self) -> usize {
+        self.vec_znx_normalize_tmp_bytes()
+    }
+
     fn glwe_normalize<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
     where
         R: GLWEToMut,
diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint_prepared.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint_prepared.rs
index c4a232f..2f26b06 100644
--- a/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint_prepared.rs
+++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint_prepared.rs
@@ -331,7 +331,11 @@ where
 
         let scratch_thread_size = self.fhe_uint_prepare_tmp_bytes(cbt.block_size(), 1, res, bits, key);
 
-        assert!(scratch.available() >= threads * scratch_thread_size);
+        assert!(
+            scratch.available() >= threads * scratch_thread_size,
+            "scratch.available():{} < threads:{threads} * scratch_thread_size:{scratch_thread_size}",
+            scratch.available()
+        );
 
         let chunk_size: usize = bit_count.div_ceil(threads);