From c5fe07188fbfb0cd2e317e745213095cdbc879db Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Mon, 19 May 2025 14:22:05 +0200
Subject: [PATCH] fixed tests of automorphism over glwe

---
 core/src/automorphism.rs                | 18 ++++++++++++++++++
 core/src/glwe_ciphertext.rs             |  6 +++---
 core/src/test_fft64/automorphism_key.rs |  2 +-
 core/src/test_fft64/glwe.rs             | 18 ++++++++----------
 4 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/core/src/automorphism.rs b/core/src/automorphism.rs
index 6913fab..64b3cf7 100644
--- a/core/src/automorphism.rs
+++ b/core/src/automorphism.rs
@@ -142,6 +142,10 @@ impl AutomorphismKey<Vec<u8>, FFT64> {
         tmp_dft + tmp_idft + idft + keyswitch
     }
 
+    pub fn automorphism_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, ksk_size: usize, rank: usize) -> usize {
+        AutomorphismKey::automorphism_scratch_space(module, out_size, out_size, ksk_size, rank)
+    }
+
     pub fn external_product_scratch_space(
         module: &Module<FFT64>,
         out_size: usize,
@@ -318,6 +322,20 @@ where
         self.p = (lhs.p * rhs.p) % (module.cyclotomic_order() as i64);
     }
 
+    pub fn automorphism_inplace<DataRhs>(
+        &mut self,
+        module: &Module<FFT64>,
+        rhs: &AutomorphismKey<DataRhs, FFT64>,
+        scratch: &mut Scratch,
+    ) where
+        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
+    {
+        unsafe {
+            let self_ptr: *mut AutomorphismKey<DataSelf, FFT64> = self as *mut AutomorphismKey<DataSelf, FFT64>;
+            self.automorphism(&module, &*self_ptr, rhs, scratch);
+        }
+    }
+
     pub fn keyswitch<DataLhs, DataRhs>(
         &mut self,
         module: &Module<FFT64>,
diff --git a/core/src/glwe_ciphertext.rs b/core/src/glwe_ciphertext.rs
index 0fd6242..ca94db1 100644
--- a/core/src/glwe_ciphertext.rs
+++ b/core/src/glwe_ciphertext.rs
@@ -294,9 +294,9 @@ where
         MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
     {
         self.keyswitch(module, lhs, &rhs.key, scratch);
-        //(0..self.rank() + 1).for_each(|i| {
-        //    module.vec_znx_automorphism_inplace(rhs.p(), self, i);
-        //})
+        (0..self.rank() + 1).for_each(|i| {
+            module.vec_znx_automorphism_inplace(rhs.p(), self, i);
+        })
     }
 
     pub fn automorphism_inplace<DataRhs>(
diff --git a/core/src/test_fft64/automorphism_key.rs b/core/src/test_fft64/automorphism_key.rs
index 0e81578..9705a3f 100644
--- a/core/src/test_fft64/automorphism_key.rs
+++ b/core/src/test_fft64/automorphism_key.rs
@@ -1,4 +1,4 @@
-use base2k::{FFT64, Module, ScalarZnxOps, ScalarZnxToRef, ScratchOwned, Stats, VecZnxOps, ZnxView};
+use base2k::{FFT64, Module, ScalarZnxOps, ScratchOwned, Stats, VecZnxOps};
 use sampling::source::Source;
 
 use crate::{
diff --git a/core/src/test_fft64/glwe.rs b/core/src/test_fft64/glwe.rs
index 687d49a..54f389c 100644
--- a/core/src/test_fft64/glwe.rs
+++ b/core/src/test_fft64/glwe.rs
@@ -1,6 +1,6 @@
 use base2k::{
     Decoding, Encoding, FFT64, FillUniform, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned, Stats, VecZnxOps, VecZnxToMut,
-    ZnxViewMut, ZnxZero,
+    ZnxView, ZnxViewMut, ZnxZero,
 };
 use itertools::izip;
 use sampling::source::Source;
@@ -420,7 +420,7 @@ fn test_keyswitch_inplace(log_n: usize, basek: usize, k_ksk: usize, k_ct: usize,
 fn automorphism() {
     (1..4).for_each(|rank| {
         println!("test automorphism rank: {}", rank);
-        test_automorphism(12, 12, 1, 60, 45, 60, rank, 3.2);
+        test_automorphism(12, 12, -5, 60, 45, 60, rank, 3.2);
     });
 }
 
@@ -447,7 +447,6 @@ fn test_automorphism(
     let mut source_xe: Source = Source::new([0u8; 32]);
     let mut source_xa: Source = Source::new([0u8; 32]);
 
-    // Random input plaintext
     pt_want
         .data
         .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
@@ -486,14 +485,15 @@ fn test_automorphism(
     );
 
     ct_out.automorphism(&module, &ct_in, &autokey, scratch.borrow());
-
     ct_out.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
-
     module.vec_znx_automorphism_inplace(p, &mut pt_want, 0);
-
     module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
+    module.vec_znx_normalize_inplace(basek, &mut pt_have, 0, scratch.borrow());
 
     let noise_have: f64 = pt_have.data.std(0, basek).log2();
+
+    println!("{}", noise_have);
+
     let noise_want: f64 = noise_gglwe_product(
         module.n() as f64,
         basek,
@@ -519,7 +519,7 @@ fn test_automorphism(
 fn automorphism_inplace() {
     (1..4).for_each(|rank| {
         println!("test automorphism_inplace rank: {}", rank);
-        test_automorphism_inplace(12, 12, 1, 60, 60, rank, 3.2);
+        test_automorphism_inplace(12, 12, -5, 60, 60, rank, 3.2);
     });
 }
 
@@ -575,12 +575,10 @@ fn test_automorphism_inplace(log_n: usize, basek: usize, p: i64, k_autokey: usiz
     );
 
     ct.automorphism_inplace(&module, &autokey, scratch.borrow());
-
     ct.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
-
     module.vec_znx_automorphism_inplace(p, &mut pt_want, 0);
-
     module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
+    module.vec_znx_normalize_inplace(basek, &mut pt_have, 0, scratch.borrow());
 
     let noise_have: f64 = pt_have.data.std(0, basek).log2();
     let noise_want: f64 = noise_gglwe_product(