From 836df871fe4ddca838df3e2e371f30253de27c5f Mon Sep 17 00:00:00 2001
From: Pro7ech <jeanphilippe.bossuat@gmail.com>
Date: Fri, 7 Nov 2025 16:30:47 +0100
Subject: [PATCH] Add normalize in cmux & uint_prepared to uint

---
 .../bdd_arithmetic/ciphertexts/fhe_uint.rs    | 29 ++++++++++++++++++-
 .../src/tfhe/bdd_arithmetic/eval.rs           | 15 ++++++----
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs
index 5ae7145..6bd8be1 100644
--- a/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs
+++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs
@@ -13,7 +13,7 @@ use poulpy_hal::{
 };
 use std::{collections::HashMap, marker::PhantomData};
 
-use crate::tfhe::bdd_arithmetic::{FromBits, ToBits, UnsignedInteger};
+use crate::tfhe::bdd_arithmetic::{Cmux, FheUintPrepared, FromBits, GetGGSWBit, ToBits, UnsignedInteger};
 
 /// An FHE ciphertext encrypting the bits of an [UnsignedInteger].
 pub struct FheUint<D: Data, T: UnsignedInteger> {
@@ -362,6 +362,33 @@ impl<D: DataRef, T: UnsignedInteger> GLWEToRef for FheUint<D, T> {
 }
 
 impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
+    pub fn from_fhe_uint_prepared<M, DR, H, K, BE: Backend>(
+        &mut self,
+        module: &M,
+        other: &FheUintPrepared<DR, T, BE>,
+        keys: &H,
+        scratch: &mut Scratch<BE>,
+    ) where
+        DR: DataRef,
+        M: Cmux<BE> + ModuleLogN + GLWEPacking<BE> + GLWECopy,
+        Scratch<BE>: ScratchTakeCore<BE>,
+        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>,
+    {
+        let zero: GLWE<Vec<u8>> = GLWE::alloc_from_infos(self);
+        let mut one: GLWE<Vec<u8>> = GLWE::alloc_from_infos(self);
+        one.data_mut()
+            .encode_coeff_i64(self.base2k().into(), 0, 2, 0, 1);
+
+        let (mut out_bits, scratch_1) = scratch.take_glwe_slice(T::BITS as usize, self);
+
+        for i in 0..T::BITS as usize {
+            module.cmux(&mut out_bits[i], &one, &zero, &other.get_bit(i), scratch_1);
+        }
+
+        self.pack(module, out_bits, keys, scratch_1);
+    }
+
     pub fn zero_byte<M, K, H, BE: Backend>(&mut self, module: &M, byte: usize, keys: &H, scratch: &mut Scratch<BE>)
     where
         H: GLWEAutomorphismKeyHelper<K, BE>,
diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/eval.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/eval.rs
index 5369903..dbd1862 100644
--- a/poulpy-schemes/src/tfhe/bdd_arithmetic/eval.rs
+++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/eval.rs
@@ -2,8 +2,7 @@ use core::panic;
 
 use itertools::Itertools;
 use poulpy_core::{
-    GLWEAdd, GLWECopy, GLWEExternalProduct, GLWESub, ScratchTakeCore,
-    layouts::{GGSWInfos, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos, prepared::GGSWPreparedToRef},
+    GLWEAdd, GLWECopy, GLWEExternalProduct, GLWENormalize, GLWESub, ScratchTakeCore, layouts::{GGSWInfos, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos, prepared::GGSWPreparedToRef}
 };
 use poulpy_hal::layouts::{Backend, DataMut, Module, Scratch, ZnxZero};
 
@@ -71,8 +70,8 @@ where
     {
         #[cfg(debug_assertions)]
         {
-            assert!(inputs.bit_size() >= circuit.input_size());
-            assert!(out.len() >= circuit.output_size());
+            assert!(inputs.bit_size() >= circuit.input_size(), "inputs.bit_size(): {} < circuit.input_size():{}", inputs.bit_size(), circuit.input_size());
+            assert!(out.len() >= circuit.output_size(), "out.len(): {} < circuit.output_size(): {}", out.len(), circuit.output_size());
         }
 
         for (i, out_i) in out.iter_mut().enumerate().take(circuit.output_size()) {
@@ -165,7 +164,7 @@ pub enum Node {
 
 pub trait Cmux<BE: Backend>
 where
-    Self: GLWEExternalProduct<BE> + GLWESub + GLWEAdd,
+    Self: GLWEExternalProduct<BE> + GLWESub + GLWEAdd + GLWENormalize<BE>,
 {
     fn cmux_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
     where
@@ -185,8 +184,10 @@ where
         Scratch<BE>: ScratchTakeCore<BE>,
     {
         self.glwe_sub(res, t, f);
+        self.glwe_normalize_inplace(res, scratch);
         self.glwe_external_product_inplace(res, s, scratch);
         self.glwe_add_inplace(res, f);
+        self.glwe_normalize_inplace(res, scratch);
     }
 
     fn cmux_inplace<R, A, S>(&self, res: &mut R, a: &A, s: &S, scratch: &mut Scratch<BE>)
@@ -197,14 +198,16 @@ where
         Scratch<BE>: ScratchTakeCore<BE>,
     {
         self.glwe_sub_inplace(res, a);
+        self.glwe_normalize_inplace(res, scratch);
         self.glwe_external_product_inplace(res, s, scratch);
         self.glwe_add_inplace(res, a);
+        self.glwe_normalize_inplace(res, scratch);
     }
 }
 
 impl<BE: Backend> Cmux<BE> for Module<BE>
 where
-    Self: GLWEExternalProduct<BE> + GLWESub + GLWEAdd,
+    Self: GLWEExternalProduct<BE> + GLWESub + GLWEAdd + GLWENormalize<BE>,
     Scratch<BE>: ScratchTakeCore<BE>,
 {
 }