From 13e26c815282780630426be3aff8c0065814813e Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Mon, 19 May 2025 13:17:56 +0200
Subject: [PATCH] Added test for automorphism inplace

---
 core/src/test_fft64/glwe.rs | 110 ++++++++++++++++++++++++++++++++----
 1 file changed, 99 insertions(+), 11 deletions(-)

diff --git a/core/src/test_fft64/glwe.rs b/core/src/test_fft64/glwe.rs
index 525de22..53c06fe 100644
--- a/core/src/test_fft64/glwe.rs
+++ b/core/src/test_fft64/glwe.rs
@@ -448,18 +448,9 @@ fn test_automorphism(
     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);
-
     pt_want
-        .to_mut()
-        .at_mut(0, 1)
-        .iter_mut()
-        .enumerate()
-        .for_each(|(i, x)| {
-            *x = i as i64;
-        });
+        .data
+        .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
 
     let mut scratch: ScratchOwned = ScratchOwned::new(
         GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, autokey.size())
@@ -524,6 +515,103 @@ fn test_automorphism(
     );
 }
 
+#[test]
+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);
+    });
+}
+
+fn test_automorphism_inplace(
+    log_n: usize,
+    basek: usize,
+    p: i64,
+    k_autokey: usize,
+    k_ct: usize,
+    rank: usize,
+    sigma: f64,
+) {
+    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
+    let rows: usize = (k_ct + basek - 1) / basek;
+
+    let mut autokey: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, p, k_autokey, rows, rank);
+    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, basek, k_ct, rank);
+    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ct);
+    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ct);
+
+    let mut source_xs: Source = Source::new([0u8; 32]);
+    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);
+
+    let mut scratch: ScratchOwned = ScratchOwned::new(
+        GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, autokey.size())
+            | GLWECiphertext::decrypt_scratch_space(&module, ct.size())
+            | GLWECiphertext::encrypt_sk_scratch_space(&module, ct.size())
+            | GLWECiphertext::automorphism_inplace_scratch_space(&module, ct.size(), rank, autokey.size()),
+    );
+
+    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
+    sk.fill_ternary_prob(0.5, &mut source_xs);
+
+    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
+    sk_dft.dft(&module, &sk);
+
+    autokey.encrypt_sk(
+        &module,
+        p,
+        &sk,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+
+    ct.encrypt_sk(
+        &module,
+        &pt_want,
+        &sk_dft,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+
+    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);
+
+    let noise_have: f64 = pt_have.data.std(0, basek).log2();
+    let noise_want: f64 = noise_gglwe_product(
+        module.n() as f64,
+        basek,
+        0.5,
+        0.5,
+        0f64,
+        sigma * sigma,
+        0f64,
+        rank as f64,
+        k_ct,
+        k_autokey,
+    );
+
+    assert!(
+        (noise_have - noise_want).abs() <= 0.1,
+        "{} {}",
+        noise_have,
+        noise_want
+    );
+}
+
 fn test_external_product(log_n: usize, basek: usize, k_ggsw: usize, k_ct_in: usize, k_ct_out: usize, rank: usize, sigma: f64) {
     let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);