From 3b7d88997197eae0b16b1c5aa08f42240a0d0aff Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Thu, 16 Jan 2025 17:59:01 +0100
Subject: [PATCH] fixed overflow

---
 math/src/automorphism.rs               | 10 +++++++++-
 math/src/ring.rs                       |  5 ++---
 math/src/ring/impl_u64/automorphism.rs |  7 +++++++
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/math/src/automorphism.rs b/math/src/automorphism.rs
index 4c9a2be..97ccca1 100644
--- a/math/src/automorphism.rs
+++ b/math/src/automorphism.rs
@@ -18,7 +18,12 @@ impl AutoPermMap {
         gal_el
     }
 
-    pub fn gen<O: Unsigned, const NTT: bool>(&mut self, ring: &Ring<O>, gen_1: usize, gen_2: bool) -> usize {
+    pub fn gen<O: Unsigned, const NTT: bool>(
+        &mut self,
+        ring: &Ring<O>,
+        gen_1: usize,
+        gen_2: bool,
+    ) -> usize {
         self.insert(AutoPerm::new::<O, NTT>(ring, gen_1, gen_2))
     }
 
@@ -27,7 +32,9 @@ impl AutoPermMap {
     }
 }
 
+#[derive(Debug)]
 pub struct AutoPerm {
+    pub ntt: bool,
     pub gal_el: usize,
     pub permutation: Vec<usize>,
 }
@@ -64,6 +71,7 @@ impl AutoPerm {
         }
 
         Self {
+            ntt: NTT,
             gal_el: gal_el,
             permutation: permutation,
         }
diff --git a/math/src/ring.rs b/math/src/ring.rs
index 1387c9a..df88c1a 100644
--- a/math/src/ring.rs
+++ b/math/src/ring.rs
@@ -45,11 +45,10 @@ impl<O: Unsigned> Ring<O> {
             e >>= 1;
         }
 
-        let nth_root = 1 << self.cyclotomic_order;
-        gal_el &= nth_root - 1;
+        gal_el &= self.cyclotomic_order - 1;
 
         if gen_2 {
-            return nth_root - gal_el;
+            return self.cyclotomic_order - gal_el;
         }
         gal_el
     }
diff --git a/math/src/ring/impl_u64/automorphism.rs b/math/src/ring/impl_u64/automorphism.rs
index 25579dd..f4fd4a6 100644
--- a/math/src/ring/impl_u64/automorphism.rs
+++ b/math/src/ring/impl_u64/automorphism.rs
@@ -158,6 +158,13 @@ impl Ring<u64> {
             b.n()
         );
 
+        assert!(
+            NTT == auto_perm.ntt,
+            "missmatch between AutoPerm NTT flag={} and method NTT flag={}",
+            auto_perm.ntt,
+            NTT
+        );
+
         let b_vec: &mut Vec<u64> = &mut b.0;
         let a_vec: &Vec<u64> = &a.0;