From 38df06f7abb4e1d028e5b1c6d76f1f0697e8a2f0 Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Tue, 8 Jul 2025 18:50:04 +0200
Subject: [PATCH] Fixed lut generation

---
 core/src/blind_rotation/lut.rs            |  7 ++++++-
 core/src/blind_rotation/test_fft64/lut.rs | 12 ++++++------
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/core/src/blind_rotation/lut.rs b/core/src/blind_rotation/lut.rs
index 96c1422..7446e9a 100644
--- a/core/src/blind_rotation/lut.rs
+++ b/core/src/blind_rotation/lut.rs
@@ -40,8 +40,13 @@ impl LookUpTable {
         // Get the number minimum limb to store the message modulus
         let limbs: usize = k.div_ceil(1 << basek);
 
+        #[cfg(debug_assertions)]
+        {
+            assert!(limbs <= self.data[0].size());
+        }
+
         // Scaling factor
-        let scale: i64 = (1 << (basek * limbs - 1)).div_round(k) as i64;
+        let scale: i64 = 1 << (k % basek) as i64;
 
         // #elements in lookup table
         let f_len: usize = f.len();
diff --git a/core/src/blind_rotation/test_fft64/lut.rs b/core/src/blind_rotation/test_fft64/lut.rs
index 3738b62..02f710d 100644
--- a/core/src/blind_rotation/test_fft64/lut.rs
+++ b/core/src/blind_rotation/test_fft64/lut.rs
@@ -12,7 +12,7 @@ fn standard() {
     let message_modulus: usize = 16;
     let extension_factor: usize = 1;
 
-    let scale: usize = (1 << (basek - 1)) / message_modulus;
+    let log_scale: usize = basek + 1;
 
     let mut f: Vec<i64> = vec![0i64; message_modulus];
     f.iter_mut()
@@ -20,7 +20,7 @@ fn standard() {
         .for_each(|(i, x)| *x = (i as i64) - 8);
 
     let mut lut: LookUpTable = LookUpTable::alloc(&module, basek, k_lut, extension_factor);
-    lut.set(&module, &f, message_modulus);
+    lut.set(&module, &f, log_scale);
 
     let half_step: i64 = lut.domain_size().div_round(message_modulus << 1) as i64;
     lut.rotate(half_step);
@@ -31,7 +31,7 @@ fn standard() {
         (0..step).for_each(|_| {
             assert_eq!(
                 f[i / step] % message_modulus as i64,
-                lut.data[0].raw()[0] / scale as i64
+                lut.data[0].raw()[0] / (1 << (log_scale % basek)) as i64
             );
             lut.rotate(-1);
         });
@@ -46,7 +46,7 @@ fn extended() {
     let message_modulus: usize = 16;
     let extension_factor: usize = 4;
 
-    let scale: usize = (1 << (basek - 1)) / message_modulus;
+    let log_scale: usize = basek + 1;
 
     let mut f: Vec<i64> = vec![0i64; message_modulus];
     f.iter_mut()
@@ -54,7 +54,7 @@ fn extended() {
         .for_each(|(i, x)| *x = (i as i64) - 8);
 
     let mut lut: LookUpTable = LookUpTable::alloc(&module, basek, k_lut, extension_factor);
-    lut.set(&module, &f, message_modulus);
+    lut.set(&module, &f, log_scale);
 
     let half_step: i64 = lut.domain_size().div_round(message_modulus << 1) as i64;
     lut.rotate(half_step);
@@ -65,7 +65,7 @@ fn extended() {
         (0..step).for_each(|_| {
             assert_eq!(
                 f[i / step] % message_modulus as i64,
-                lut.data[0].raw()[0] / scale as i64
+                lut.data[0].raw()[0] / (1 << (log_scale % basek)) as i64
             );
             lut.rotate(-1);
         });