fix(integer): fix unsigned_overflowing_sub on trivials

unsigned_overflowing_sub does an independant subtraction
on each blocks with a correcting term being added to avoid
trashing the padding bit (lhs - rhs + correction).

The correction depended on rhs's degree.
e.g. if rhs's degree was in range 1..(msg_mod-1) -> correction =
     msg_mod

However if rhs's degree was zero (so rhs is a trivial 0), the correction
was also 0, however the borrow propagation rely on that correction to
always be added.

Author: tmontaigu

tfhe/src/integer/server_key/radix/sub.rs

@@ -416,7 +416,17 @@ impl ServerKey {
         let mut borrow = self.key.create_trivial(0);
         let mut new_blocks = Vec::with_capacity(lhs.blocks.len());
         for (lhs_b, rhs_b) in lhs.blocks.iter().zip(rhs.blocks.iter()) {
-            let mut result_block = self.key.unchecked_sub(lhs_b, rhs_b);
+            let (mut result_block, correction) =
+                self.key.unchecked_sub_with_correcting_term(lhs_b, rhs_b);
+            if correction == 0 {
+                // When rhs_block is a trivial zero, the correcting term added is 0
+                // However we rely on that correcting term to be added regardless
+                assert_eq!(rhs_b.degree.0, 0);
+                self.key.unchecked_scalar_add_assign(
+                    &mut result_block,
+                    self.key.message_modulus.0 as u8,
+                );
+            }
             // Here unchecked_sub_assign does not give correct result, we don't want
             // the correcting term to be used
             // -> This is ok as the value returned by unchecked_sub is in range 1..(message_mod * 2)

tfhe/src/integer/server_key/radix_parallel/sub.rs

@@ -361,7 +361,21 @@ impl ServerKey {
                 .blocks
                 .iter()
                 .zip(rhs.blocks.iter())
-                .map(|(lhs_block, rhs_block)| self.key.unchecked_sub(lhs_block, rhs_block))
+                .map(|(lhs_block, rhs_block)| {
+                    let (mut result_block, correction) = self
+                        .key
+                        .unchecked_sub_with_correcting_term(lhs_block, rhs_block);
+                    if correction == 0 {
+                        // When rhs_block is a trivial zero, the correcting term added is 0
+                        // However we rely on that correcting term to be added regardless
+                        assert_eq!(rhs_block.degree.0, 0);
+                        self.key.unchecked_scalar_add_assign(
+                            &mut result_block,
+                            self.key.message_modulus.0 as u8,
+                        );
+                    }
+                    result_block
+                })
                 .collect::<Vec<_>>();
             let mut ct = RadixCiphertext::from(ct);
 

tfhe/src/integer/server_key/radix_parallel/tests_cases_unsigned.rs

@@ -1822,6 +1822,40 @@ where
             );
         }
     }
+
+    // Test with trivial inputs, as it was bugged at some point
+    for _ in 0..4 {
+        // Reduce maximum value of random number such that at least the last block is a trivial 0
+        // (This is how the reproducing case was found)
+        let clear_0 = rng.gen::<u64>() % (modulus / sks.key.message_modulus.0 as u64);
+        let clear_1 = rng.gen::<u64>() % (modulus / sks.key.message_modulus.0 as u64);
+
+        let a: RadixCiphertext = sks.create_trivial_radix(clear_0, NB_CTXT);
+        let b: RadixCiphertext = sks.create_trivial_radix(clear_1, NB_CTXT);
+
+        assert_eq!(a.blocks[NB_CTXT - 1].degree.0, 0);
+        assert_eq!(b.blocks[NB_CTXT - 1].degree.0, 0);
+
+        let (encrypted_result, encrypted_overflow) =
+            sks.unchecked_unsigned_overflowing_sub_parallelized(&a, &b);
+
+        let (expected_result, expected_overflowed) =
+            overflowing_sub_under_modulus(clear_0, clear_1, modulus);
+
+        let decrypted_result: u64 = cks.decrypt(&encrypted_result);
+        let decrypted_overflowed = cks.decrypt_one_block(&encrypted_overflow) == 1;
+        assert_eq!(
+            decrypted_result, expected_result,
+            "Invalid result for sub, for ({clear_0} - {clear_1}) % {modulus} \
+                expected {expected_result}, got {decrypted_result}"
+        );
+        assert_eq!(
+            decrypted_overflowed,
+            expected_overflowed,
+            "Invalid overflow flag result for overflowing_sub, for ({clear_0} - {clear_1}) % {modulus} \
+                expected overflow flag {expected_overflowed}, got {decrypted_overflowed}"
+        );
+    }
 }
 
 pub(crate) fn default_sub_test<P, T>(param: P, mut executor: T)