Update reference submission to handle all batch sizes

.gitignore

@@ -4,7 +4,7 @@ io/
 third_party/
 measurements/
 harness/__pycache__/
-submission/target
+submission/target/
 
 # Prerequisites
 *.d

submission/Cargo.toml

@@ -13,6 +13,7 @@ aligned-vec = { version = "0.5.0" }
 serde = { version = "1", features = ["derive"] }
 aes = "0.8.4"
 itertools = "0.14.0"
+rayon = "1.10"
 
 
 

submission/src/bin/client_decrypt_decode.rs

@@ -3,7 +3,7 @@ use std::fs;
 
 use submission::help_fun::decrypt_decode_lwe_list;
 use submission::help_fun::get_size_string;
-use tfhe::core_crypto::prelude::{Container, ContiguousEntityContainer, LweCiphertext, LweCiphertextList, LweSecretKey, UnsignedInteger, decrypt_lwe_ciphertext};
+use tfhe::core_crypto::prelude::{LweCiphertextList, LweSecretKey, };
 
 pub fn main() -> Result<(), Box<dyn std::error::Error>> {
     let args: Vec<String> = env::args().collect();

submission/src/bin/client_encode_encrypt.rs

@@ -1,11 +1,21 @@
 use std::{env, fs};
 
-use auto_base_conv::{AES_TIGHT, AesParam};
-use submission::{aes_manager::Aes128Manager, data_struct::{AllRdKeys, get_0_round_key, get_8_to_1_round_key, get_10_9_round_key}, help_fun::get_size_string};
-use tfhe::core_crypto::{prelude::{ActivatedRandomGenerator, EncryptionRandomGenerator, GlweSecretKey, SecretRandomGenerator}, seeders::new_seeder};
-
-
-
+use auto_base_conv::{AES_SET_2, AesParam, generate_vec_keyed_lut_accumulator};
+use submission::{
+    aes_manager::Aes128Manager,
+    aes_ref::*,
+    data_struct::{AllRdKeys, AllRdKeys2, get_0_round_key, get_8_to_1_round_key, get_10_9_round_key},
+    help_fun::get_size_string,
+};
+use tfhe::core_crypto::{
+    algorithms::encrypt_lwe_ciphertext_list,
+    commons::parameters::{LweCiphertextCount, LweSize},
+    entities::{LweCiphertextList, LweCiphertextListOwned, PlaintextList},
+    prelude::{
+        ActivatedRandomGenerator, EncryptionRandomGenerator, GlweSecretKey,
+    },
+    seeders::new_seeder,
+};
 
 pub fn gen_transciphering_keys(
     param: &AesParam<u64>,
@@ -21,6 +31,75 @@ pub fn gen_transciphering_keys(
     }
 }
 
+pub fn gen_transciphering_keys_2(
+    param: &AesParam<u64>,
+    glwe_sk: &GlweSecretKey<Vec<u64>>,
+    aes_key: &[u8; 16],
+    encryption_generator: &mut EncryptionRandomGenerator<ActivatedRandomGenerator>,
+) -> AllRdKeys2 {
+    let aes = Aes128Ref::new(&aes_key);
+    let large_lwe_size = LweSize(param.glwe_dimension().0 * param.polynomial_size().0 + 1);
+    let round_keys = aes.get_round_keys();
+    let mut he_round_keys = Vec::<LweCiphertextListOwned<u64>>::with_capacity(NUM_ROUNDS + 1);
+    for r in 0..=NUM_ROUNDS {
+        let mut lwe_list_rk = LweCiphertextList::new(
+            0u64,
+            large_lwe_size,
+            LweCiphertextCount(BLOCKSIZE_IN_BIT),
+            param.ciphertext_modulus(),
+        );
+
+        let rk = PlaintextList::from_container(
+            (0..BLOCKSIZE_IN_BIT)
+                .map(|i| {
+                    let byte_idx = i / BYTESIZE;
+                    let bit_idx = i % BYTESIZE;
+                    let round_key_byte = round_keys[r][byte_idx];
+                    let round_key_bit = (round_key_byte & (1 << bit_idx)) >> bit_idx;
+                    (round_key_bit as u64) << 63
+                })
+                .collect::<Vec<u64>>(),
+        );
+        encrypt_lwe_ciphertext_list(
+            &glwe_sk.clone().into_lwe_secret_key(),
+            &mut lwe_list_rk,
+            &rk,
+            param.glwe_modular_std_dev(),
+            encryption_generator,
+        );
+
+        he_round_keys.push(lwe_list_rk);
+    }
+    let vec_keyed_sbox_round_1 = generate_vec_keyed_lut_accumulator(
+        aes.get_keyed_sbox(0),
+        u64::BITS as usize - 1,
+        &glwe_sk,
+        param.glwe_modular_std_dev(),
+        param.ciphertext_modulus(),
+        encryption_generator,
+    );
+    let vec_keyed_sbox_round_1_mult_by_2 = generate_vec_keyed_lut_accumulator(
+        aes.get_keyed_sbox_mult_by_2(0),
+        u64::BITS as usize - 1,
+        &glwe_sk,
+        param.glwe_modular_std_dev(),
+        param.ciphertext_modulus(),
+        encryption_generator,
+    );
+    let vec_keyed_sbox_round_1_mult_by_3 = generate_vec_keyed_lut_accumulator(
+        aes.get_keyed_sbox_mult_by_3(0),
+        u64::BITS as usize - 1,
+        &glwe_sk,
+        param.glwe_modular_std_dev(),
+        param.ciphertext_modulus(),
+        encryption_generator,
+    );
+    AllRdKeys2 {
+        _0_round_key: vec![vec_keyed_sbox_round_1, vec_keyed_sbox_round_1_mult_by_2, vec_keyed_sbox_round_1_mult_by_3],
+        other_round_keys: he_round_keys,
+    }
+}
+
 pub fn main() -> Result<(), Box<dyn std::error::Error>> {
     let args: Vec<String> = env::args().collect();
     if args.len() < 2 {
@@ -36,31 +115,50 @@ pub fn main() -> Result<(), Box<dyn std::error::Error>> {
     let hex_string = fs::read_to_string(&aes_key_path)?.trim().to_string();
 
     let mut aes_key: [u8; 16] = [0u8; 16];
-    for (i, byte) in aes_key.iter_mut().enumerate() {
-        let hex_pair = &hex_string[i * 2..i * 2 + 2];
-        *byte = u8::from_str_radix(hex_pair, 16)?;
-    }
 
     let secret_keys_dir = format!("{}/secret_keys", io_dir);
     let glwe_sk_path = format!("{}/glwe_sk.bin", secret_keys_dir);
     let glwe_sk_bytes = fs::read(&glwe_sk_path)?;
     let glwe_sk: GlweSecretKey<Vec<u64>> = bincode::deserialize(&glwe_sk_bytes)?;
 
-    let param = &*AES_TIGHT;
+    let param = AES_SET_2.clone(); //AES_TIGHT
     let mut boxed_seeder = new_seeder();
     let seeder = boxed_seeder.as_mut();
     let mut encryption_generator =
         EncryptionRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed(), seeder);
-
-    let trans_key = gen_transciphering_keys(param, &glwe_sk, &aes_key, &mut encryption_generator);
-
-    let ciphertext_upload_dir = format!("{}/ciphertexts_upload", io_dir);
-    fs::create_dir_all(&ciphertext_upload_dir)?;
-
-    let trans_key_path = format!("{}/trans_key.bin", ciphertext_upload_dir);
-    fs::write(&trans_key_path, bincode::serialize(&trans_key)?)?;
-
-    println!("Transciphering keys saved to {}", ciphertext_upload_dir);
+
+    if size == "0" {
+        for (i, byte) in aes_key.iter_mut().enumerate() {
+            let hex_pair = &hex_string[i * 2..i * 2 + 2];
+            *byte = u8::from_str_radix(hex_pair, 16)?;
+        }
+        let trans_key =
+            gen_transciphering_keys(&param, &glwe_sk, &aes_key, &mut encryption_generator);
+
+        let ciphertext_upload_dir = format!("{}/ciphertexts_upload", io_dir);
+        fs::create_dir_all(&ciphertext_upload_dir)?;
+
+        let trans_key_path = format!("{}/trans_key.bin", ciphertext_upload_dir);
+        fs::write(&trans_key_path, bincode::serialize(&trans_key)?)?;
+
+        println!("Transciphering keys saved to {}", ciphertext_upload_dir);
+    } else if size == "1" || size == "2" {
+        for (i, byte) in aes_key.iter_mut().enumerate() {
+            let hex_pair = &hex_string[i * 2..i * 2 + 2];
+            *byte = u8::from_str_radix(hex_pair, 16)?;
+        }
+        let trans_keys_2 =
+            gen_transciphering_keys_2(&param, &glwe_sk, &aes_key, &mut encryption_generator);
+        let ciphertext_upload_dir = format!("{}/ciphertexts_upload", io_dir);
+        fs::create_dir_all(&ciphertext_upload_dir)?;
+
+        let trans_key_path = format!("{}/trans_key.bin", ciphertext_upload_dir);
+        fs::write(&trans_key_path, bincode::serialize(&trans_keys_2)?)?;
+
+        println!("Transciphering keys saved to {}", ciphertext_upload_dir);
+    } else {
+        return Err(Box::from(format!("Unexpected size {}", size)));
+    }
 
     Ok(())
 }

submission/src/bin/client_key_generation.rs

@@ -1,11 +1,19 @@
-use submission::help_fun::get_size_string;
-use std::{collections::HashMap, env};
 use std::fs;
+use std::{collections::HashMap, env};
+use submission::help_fun::get_size_string;
 
 use aligned_vec::ABox;
-use auto_base_conv::{AES_TIGHT, AesParam, AutomorphKey, AutomorphKeySerializable, GlweKeyswitchKeyOwned, gen_all_auto_keys, generate_scheme_switching_key, keygen_pbs_with_glwe_ks};
-use tfhe::core_crypto::{prelude::{ActivatedRandomGenerator, EncryptionRandomGenerator, GgswCiphertextList, GlweSecretKeyOwned, LweBootstrapKeyOwned, LweSecretKeyOwned, SecretRandomGenerator}, seeders::new_seeder};
+use auto_base_conv::{
+    AES_SET_2, AesParam, AutomorphKey, AutomorphKeySerializable, GlweKeyswitchKeyOwned, gen_all_auto_keys, generate_scheme_switching_key, keygen_pbs_with_glwe_ks
+};
 use tfhe::core_crypto::fft_impl::fft64::c64;
+use tfhe::core_crypto::{
+    prelude::{
+        ActivatedRandomGenerator, EncryptionRandomGenerator, GgswCiphertextList,
+        GlweSecretKeyOwned, LweBootstrapKeyOwned, LweSecretKeyOwned, SecretRandomGenerator,
+    },
+    seeders::new_seeder,
+};
 
 pub fn generate_fhe_keys(
     param: &AesParam<u64>,
@@ -54,7 +62,7 @@ pub fn generate_fhe_keys(
         secret_generator,
         encryption_generator,
     );
-    
+
     let ss_key = generate_scheme_switching_key(
         &glwe_sk,
         ss_base_log,
@@ -77,57 +85,53 @@ pub fn generate_fhe_keys(
     (lwe_sk, glwe_sk, bsk, ksk, auto_keys, ss_key)
 }
 
-
 pub fn main() -> Result<(), Box<dyn std::error::Error>> {
     let args: Vec<String> = env::args().collect();
     if args.len() < 2 {
         eprintln!("Usage: {} <size>", args[0]);
-        std::process::exit(1); 
+        std::process::exit(1);
     }
     let size = args[1].clone();
     let io_dir = "io/".to_owned() + get_size_string(size.parse::<usize>()?);
-   
-    let param = &*AES_TIGHT;
+
+    let param = AES_SET_2.clone(); //AES_TIGHT
     let mut boxed_seeder = new_seeder();
     let seeder = boxed_seeder.as_mut();
     let mut secret_generator =
         SecretRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed());
     let mut encryption_generator =
         EncryptionRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed(), seeder);
 
-    let fhe_keys =
-        generate_fhe_keys(param, &mut secret_generator, &mut encryption_generator);
-
+    let fhe_keys = generate_fhe_keys(&param, &mut secret_generator, &mut encryption_generator);
+
     let (lwe_sk, glwe_sk, bsk, ksk, auto_keys, ss_key) = fhe_keys;
     let serialize_auto_keys = auto_keys
         .into_iter()
         .map(|(k, v)| (k, v.to_serializable()))
         .collect::<HashMap<usize, AutomorphKeySerializable>>();
 
-
     // create secret keys directory
     let secret_keys_dir = format!("{}/secret_keys", io_dir);
     let public_keys_dir = format!("{}/public_keys", io_dir);
     fs::create_dir_all(&secret_keys_dir)?;
     fs::create_dir_all(&public_keys_dir)?;
-    
+
     // save secret keys
     let lwe_sk_path = format!("{}/lwe_sk.bin", secret_keys_dir);
     let glwe_sk_path = format!("{}/glwe_sk.bin", secret_keys_dir);
     fs::write(&lwe_sk_path, bincode::serialize(&lwe_sk)?)?;
     fs::write(&glwe_sk_path, bincode::serialize(&glwe_sk)?)?;
-    
+
     // save public/evaluation keys
     let bsk_path = format!("{}/bsk.bin", public_keys_dir);
     let ksk_path = format!("{}/ksk.bin", public_keys_dir);
     let auto_keys_path = format!("{}/auto_keys.bin", public_keys_dir);
     let ss_key_path = format!("{}/ss_key.bin", public_keys_dir);
-    
+
     fs::write(&bsk_path, bincode::serialize(&bsk)?)?;
     fs::write(&ksk_path, bincode::serialize(&ksk)?)?;
     fs::write(&auto_keys_path, bincode::serialize(&serialize_auto_keys)?)?;
     fs::write(&ss_key_path, bincode::serialize(&ss_key)?)?;
-
 
     Ok(())
-}
+}