From d75844dea54c923c5392a903a4b3deed0e241c64 Mon Sep 17 00:00:00 2001
From: Arthur Meyre <arthur.meyre@zama.ai>
Date: Wed, 12 Nov 2025 11:35:57 +0100
Subject: [PATCH] fix(core): fix decomposition algorithm not matching the
 theory

- problem arose from a shift being done on an unsigned value which did not
keep the signed characteristics of the represented signed value
- introduce an arithmetic_shift on the UnsignedInteger trait with a blanket
implementation
- add the edge case which revelead the issue
- the asm has been verified to only change for the shift operation being
applied, meaning no performance regression will occurr
---
 .../core_crypto/commons/math/decomposition/iter.rs |  4 ++--
 .../commons/math/decomposition/tests.rs            | 14 ++++++++++++++
 tfhe/src/core_crypto/commons/numeric/unsigned.rs   |  6 ++++++
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/tfhe/src/core_crypto/commons/math/decomposition/iter.rs b/tfhe/src/core_crypto/commons/math/decomposition/iter.rs
index 40bedb67e..03a0a9976 100644
--- a/tfhe/src/core_crypto/commons/math/decomposition/iter.rs
+++ b/tfhe/src/core_crypto/commons/math/decomposition/iter.rs
@@ -143,9 +143,9 @@ pub(crate) fn decompose_one_level<S: UnsignedInteger>(
     mod_b_mask: S,
 ) -> S {
     let res = *state & mod_b_mask;
-    *state >>= base_log;
+    *state = (*state).arithmetic_shr(base_log);
     let carry = decomposition_bit_trick(res, *state, base_log);
-    *state += carry;
+    *state = (*state).wrapping_add(carry);
     res.wrapping_sub(carry << base_log)
 }
 
diff --git a/tfhe/src/core_crypto/commons/math/decomposition/tests.rs b/tfhe/src/core_crypto/commons/math/decomposition/tests.rs
index 9ba0a177b..0f88a6a9a 100644
--- a/tfhe/src/core_crypto/commons/math/decomposition/tests.rs
+++ b/tfhe/src/core_crypto/commons/math/decomposition/tests.rs
@@ -224,3 +224,17 @@ fn test_single_level_decompose_balanced() {
     // We expect an average value of 0 so the sum is also 0
     assert_eq!(sum, 0);
 }
+
+#[test]
+fn test_decomposition_edge_case_sign_handling() {
+    let decomposer = SignedDecomposer::new(DecompositionBaseLog(17), DecompositionLevelCount(3));
+    let val: u64 = 0x8000_00e3_55b0_c827;
+
+    let decomp = decomposer.decompose(val);
+
+    let expected = [44422i64, 909, -65536];
+
+    for (term, expect) in decomp.zip(expected) {
+        assert_eq!(term.value() as i64, expect, "Problem with term {term:?}");
+    }
+}
diff --git a/tfhe/src/core_crypto/commons/numeric/unsigned.rs b/tfhe/src/core_crypto/commons/numeric/unsigned.rs
index 33c928f68..e255e9b31 100644
--- a/tfhe/src/core_crypto/commons/numeric/unsigned.rs
+++ b/tfhe/src/core_crypto/commons/numeric/unsigned.rs
@@ -102,6 +102,12 @@ pub trait UnsignedInteger:
     /// Integer division rounding up.
     #[must_use]
     fn div_ceil(self, divisor: Self) -> Self;
+    /// Interpret the value as signed and apply an arithmetic right shift, sign extending like on a
+    /// signed value.
+    #[must_use]
+    fn arithmetic_shr(self, rhs: usize) -> Self {
+        self.into_signed().shr(rhs).into_unsigned()
+    }
     /// Return the casting of the current value to the signed type of the same size.
     fn into_signed(self) -> Self::Signed;
     /// Return a bit representation of the integer, where blocks of length `block_length` are