From e7bf8e9307fe7d13e76c6fda3331c8ab7335f0a4 Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Sat, 8 Nov 2025 18:11:11 +0100
Subject: [PATCH] fix sext

---
 .../bdd_arithmetic/ciphertexts/fhe_uint.rs    |  4 +-
 .../tests/test_suite/fheuint.rs               | 55 +++++++++++++------
 2 files changed, 41 insertions(+), 18 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 819af68..19629f6 100644
--- a/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs
+++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs
@@ -426,7 +426,7 @@ impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
         assert!(byte < (1 << T::LOG_BYTES));
 
         let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
-        let rot: i64 = (T::bit_index(byte << 3) << log_gap) as i64;
+        let rot: i64 = (T::bit_index((byte << 3) + 7) << log_gap) as i64;
 
         let (mut sext, scratch_1) = scratch.take_glwe(self);
 
@@ -443,7 +443,7 @@ impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
 
         // Splice sext
         let (mut tmp, scratch_2) = scratch_1.take_glwe(self);
-        for i in byte..(1 << T::LOG_BYTES) as usize {
+        for i in (byte + 1)..(1 << T::LOG_BYTES) as usize {
             FheUint::<&mut [u8], T>::from_glwe_to_mut(&mut tmp).splice_u8(module, i, 0, &self.bits, &sext, keys, scratch_2);
             module.glwe_copy(&mut self.bits, &tmp);
         }
diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/fheuint.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/fheuint.rs
index fc1af3f..7d19305 100644
--- a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/fheuint.rs
+++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/fheuint.rs
@@ -37,29 +37,52 @@ where
     let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 
     for j in 0..3 {
-        for i in 0..32 {
-            let a: u32 = 0xFFFFFFFF >> i;
+        let a: u32 = 0x8483_8281;
+        a_enc.encrypt_sk(
+            module,
+            a,
+            sk,
+            &mut source_xa,
+            &mut source_xe,
+            scratch.borrow(),
+        );
 
-            a_enc.encrypt_sk(
-                module,
-                a,
-                sk,
-                &mut source_xa,
-                &mut source_xe,
-                scratch.borrow(),
-            );
+        a_enc.sext(module, j, keys, scratch.borrow());
 
-            a_enc.sext(module, j, keys, scratch.borrow());
+        // println!("{:08x} -> {:08x} {:08x}", a, sext(a, j), a_enc.decrypt(module, sk, scratch.borrow()));
 
-            // println!("{:08x} -> {:08x} {:08x}", a, sext(a, j), a_enc.decrypt(module, sk, scratch.borrow()));
+        assert_eq!(
+            sext(a, ((1 + j as u32) << 3) - 1),
+            a_enc.decrypt(module, sk, scratch.borrow())
+        );
+    }
 
-            assert_eq!(sext(a, j), a_enc.decrypt(module, sk, scratch.borrow()));
-        }
+    for j in 0..3 {
+        let a: u32 = 0x4443_4241;
+        a_enc.encrypt_sk(
+            module,
+            a,
+            sk,
+            &mut source_xa,
+            &mut source_xe,
+            scratch.borrow(),
+        );
+
+        a_enc.sext(module, j, keys, scratch.borrow());
+
+        // println!("{:08x} -> {:08x} {:08x}", a, sext(a, j), a_enc.decrypt(module, sk, scratch.borrow()));
+
+        assert_eq!(
+            sext(a, ((1 + j as u32) << 3) - 1),
+            a_enc.decrypt(module, sk, scratch.borrow())
+        );
     }
 }
 
-pub fn sext(x: u32, byte: usize) -> u32 {
-    x | (((x >> (byte << 3)) & 1) * (0xFFFF_FFFF << (byte << 3)))
+pub(crate) fn sext(x: u32, bits: u32) -> u32 {
+    let lo: u32 = x << (u32::BITS - bits) >> (u32::BITS - bits);
+    let hi: u32 = ((x >> bits) & 1) * (0xFFFF_FFFF << bits);
+    hi | lo
 }
 
 pub fn test_fhe_uint_splice_u8<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)