From f58d06ddf508b246ba1dcd0d9adffde14024d8a3 Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Wed, 28 May 2025 15:19:52 +0200
Subject: [PATCH] Added bytes_of to structs that can be created through the
 trait extension of Scratch

---
 core/src/automorphism.rs            |  4 ++++
 core/src/gglwe_ciphertext.rs        |  4 ++++
 core/src/ggsw_ciphertext.rs         |  4 ++++
 core/src/glwe_ciphertext.rs         |  4 ++++
 core/src/glwe_ciphertext_fourier.rs |  4 ++++
 core/src/keys.rs                    | 12 ++++++++++++
 core/src/keyswitch_key.rs           |  4 ++++
 core/src/tensor_key.rs              |  5 +++++
 8 files changed, 41 insertions(+)

diff --git a/core/src/automorphism.rs b/core/src/automorphism.rs
index 9be5a8a..3e309f9 100644
--- a/core/src/automorphism.rs
+++ b/core/src/automorphism.rs
@@ -26,6 +26,10 @@ impl AutomorphismKey<Vec<u8>, FFT64> {
             p: 0,
         }
     }
+
+    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> usize {
+        GLWESwitchingKey::<Vec<u8>, FFT64>::bytes_of(module, basek, k, rows, rank, rank)
+    }
 }
 
 impl<T, B: Backend> Infos for AutomorphismKey<T, B> {
diff --git a/core/src/gglwe_ciphertext.rs b/core/src/gglwe_ciphertext.rs
index 23d74e0..0417707 100644
--- a/core/src/gglwe_ciphertext.rs
+++ b/core/src/gglwe_ciphertext.rs
@@ -27,6 +27,10 @@ impl<B: Backend> GGLWECiphertext<Vec<u8>, B> {
             k,
         }
     }
+
+    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank_in: usize, rank_out: usize) -> usize {
+        module.bytes_of_mat_znx_dft(rows, rank_in, rank_out + 1, derive_size(basek, k))
+    }
 }
 
 impl<T, B: Backend> Infos for GGLWECiphertext<T, B> {
diff --git a/core/src/ggsw_ciphertext.rs b/core/src/ggsw_ciphertext.rs
index 1d4d3e4..4144cf4 100644
--- a/core/src/ggsw_ciphertext.rs
+++ b/core/src/ggsw_ciphertext.rs
@@ -31,6 +31,10 @@ impl<B: Backend> GGSWCiphertext<Vec<u8>, B> {
             k: k,
         }
     }
+
+    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> usize {
+        module.bytes_of_mat_znx_dft(rows, rank + 1, rank + 1, derive_size(basek, k))
+    }
 }
 
 impl<T, B: Backend> Infos for GGSWCiphertext<T, B> {
diff --git a/core/src/glwe_ciphertext.rs b/core/src/glwe_ciphertext.rs
index 8870af9..a01c41c 100644
--- a/core/src/glwe_ciphertext.rs
+++ b/core/src/glwe_ciphertext.rs
@@ -32,6 +32,10 @@ impl GLWECiphertext<Vec<u8>> {
             k,
         }
     }
+
+    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
+        module.bytes_of_vec_znx(rank + 1, derive_size(basek, k))
+    }
 }
 
 impl<T> Infos for GLWECiphertext<T> {
diff --git a/core/src/glwe_ciphertext_fourier.rs b/core/src/glwe_ciphertext_fourier.rs
index 811eab1..cb2ec21 100644
--- a/core/src/glwe_ciphertext_fourier.rs
+++ b/core/src/glwe_ciphertext_fourier.rs
@@ -23,6 +23,10 @@ impl<B: Backend> GLWECiphertextFourier<Vec<u8>, B> {
             k: k,
         }
     }
+
+    pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize {
+        module.bytes_of_vec_znx_dft(rank + 1, derive_size(basek, k))
+    }
 }
 
 impl<T, B: Backend> Infos for GLWECiphertextFourier<T, B> {
diff --git a/core/src/keys.rs b/core/src/keys.rs
index cd69d28..f48e7bb 100644
--- a/core/src/keys.rs
+++ b/core/src/keys.rs
@@ -26,6 +26,10 @@ impl SecretKey<Vec<u8>> {
             dist: SecretDistribution::NONE,
         }
     }
+
+    pub fn bytes_of(module: &Module<FFT64>, rank: usize) -> usize {
+        module.bytes_of_scalar_znx(rank + 1)
+    }
 }
 
 impl<DataSelf> SecretKey<DataSelf> {
@@ -89,6 +93,10 @@ impl<B: Backend> SecretKeyFourier<Vec<u8>, B> {
             dist: SecretDistribution::NONE,
         }
     }
+
+    pub fn bytes_of(module: &Module<B>, rank: usize) -> usize {
+        module.bytes_of_scalar_znx_dft(rank + 1)
+    }
 }
 
 impl<D: AsRef<[u8]> + AsMut<[u8]>> SecretKeyFourier<D, FFT64> {
@@ -124,6 +132,10 @@ impl<B: Backend> GLWEPublicKey<Vec<u8>, B> {
             dist: SecretDistribution::NONE,
         }
     }
+
+    pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize {
+        GLWECiphertextFourier::<Vec<u8>, B>::bytes_of(module, basek, k, rank)
+    }
 }
 
 impl<T, B: Backend> Infos for GLWEPublicKey<T, B> {
diff --git a/core/src/keyswitch_key.rs b/core/src/keyswitch_key.rs
index a68a39e..645b24b 100644
--- a/core/src/keyswitch_key.rs
+++ b/core/src/keyswitch_key.rs
@@ -17,6 +17,10 @@ impl GLWESwitchingKey<Vec<u8>, FFT64> {
             module, basek, k, rows, rank_in, rank_out,
         ))
     }
+
+    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank_in: usize, rank_out: usize) -> usize {
+        GGLWECiphertext::<Vec<u8>, FFT64>::bytes_of(module, basek, k, rows, rank_in, rank_out)
+    }
 }
 
 impl<T, B: Backend> Infos for GLWESwitchingKey<T, B> {
diff --git a/core/src/tensor_key.rs b/core/src/tensor_key.rs
index d8cc459..8267a25 100644
--- a/core/src/tensor_key.rs
+++ b/core/src/tensor_key.rs
@@ -20,6 +20,11 @@ impl TensorKey<Vec<u8>, FFT64> {
         });
         Self { keys: keys }
     }
+
+    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> usize {
+        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
+        pairs * GLWESwitchingKey::<Vec<u8>, FFT64>::bytes_of(module, basek, k, rows, 1, rank)
+    }
 }
 
 impl<T, B: Backend> Infos for TensorKey<T, B> {