From c845a52b4ecb5e8e8c66050d4585ed30a110648d Mon Sep 17 00:00:00 2001
From: Job Doesburg <job.doesburg@gmail.com>
Date: Thu, 7 Nov 2024 17:27:53 +0100
Subject: [PATCH] Batch

---
 src/distributed.rs      |  8 ++++++++
 src/high_level/ops.rs   | 23 ++++++++++++++++++++++-
 src/tests/high_level.rs | 26 +++++++++++---------------
 3 files changed, 41 insertions(+), 16 deletions(-)

diff --git a/src/distributed.rs b/src/distributed.rs
index 1fb3055..bc39852 100644
--- a/src/distributed.rs
+++ b/src/distributed.rs
@@ -195,6 +195,14 @@ impl PEPSystem {
         pseudonymize(p, pseudonymization_info)
     }
 
+    pub fn rekey_batch<R: RngCore + CryptoRng>(&self, encrypted: &[EncryptedDataPoint], rekey_info: &RekeyInfo, rng: &mut R) -> Box<[EncryptedDataPoint]> {
+        rekey_batch(encrypted, rekey_info, rng)
+    }
+
+    pub fn pseudonymize_batch<R: RngCore + CryptoRng>(&self, encrypted: &[EncryptedPseudonym], pseudonymization_info: &PseudonymizationInfo, rng: &mut R) -> Box<[EncryptedPseudonym]> {
+        pseudonymize_batch(encrypted, pseudonymization_info, rng)
+    }
+
     pub fn transcrypt<E: Encrypted>(
         &self,
         encrypted: &E,
diff --git a/src/high_level/ops.rs b/src/high_level/ops.rs
index cf3af95..809fdca 100644
--- a/src/high_level/ops.rs
+++ b/src/high_level/ops.rs
@@ -97,6 +97,25 @@ pub fn rekey(p: &EncryptedDataPoint, rekey_info: &RekeyInfo) -> EncryptedDataPoi
     EncryptedDataPoint::from(rekey2(&p.value, &rekey_info.from.0, &rekey_info.to.0))
 }
 
+pub fn pseudonymize_batch<R: RngCore + CryptoRng>(
+    encrypted: &[EncryptedPseudonym],
+    pseudonymization_info: &PseudonymizationInfo,
+    rng: &mut R,
+) -> Box<[EncryptedPseudonym]> {
+    let mut encrypted_copy = encrypted.to_vec();
+    encrypted_copy.shuffle(rng);
+    encrypted_copy.iter()
+        .map(|x| pseudonymize(x, pseudonymization_info))
+        .collect()
+}
+pub fn rekey_batch<R: RngCore + CryptoRng>(encrypted: &[EncryptedDataPoint], rekey_info: &RekeyInfo, rng: &mut R) -> Box<[EncryptedDataPoint]> {
+    let mut encrypted_copy = encrypted.to_vec();
+    encrypted_copy.shuffle(rng);
+    encrypted_copy.iter()
+        .map(|x| rekey(x, rekey_info))
+        .collect()
+}
+
 pub fn transcrypt<E: Encrypted>(encrypted: &E, transcryption_info: &TranscryptionInfo) -> E {
     if E::IS_PSEUDONYM {
         E::from_value(rsk2(
@@ -109,4 +128,6 @@ pub fn transcrypt<E: Encrypted>(encrypted: &E, transcryption_info: &Transcryptio
     } else {
         E::from_value(rekey2(&encrypted.value(), &transcryption_info.k.from.0, &transcryption_info.k.to.0))
     }
-}
\ No newline at end of file
+}
+
+// TODO: batch transcrypt
\ No newline at end of file
diff --git a/src/tests/high_level.rs b/src/tests/high_level.rs
index 93c11e9..f0eba1b 100644
--- a/src/tests/high_level.rs
+++ b/src/tests/high_level.rs
@@ -1,4 +1,3 @@
-use std::any::Any;
 use crate::arithmetic::GroupElement;
 use crate::high_level::contexts::*;
 use crate::high_level::keys::*;
@@ -185,20 +184,11 @@ fn test_batch() {
     let (_session2_public, session2_secret) =
         make_session_keys(&global_secret, &enc_context2, &enc_secret);
 
-
-    let mut ciphertexts = Vec::new();
-
+    let mut data = vec![];
+    let mut pseudonyms = vec![];
     for _ in 0..10 {
-        let mut entity_data: [dyn Encrypted<UnencryptedType=dyn Any>] = *[];
-        let pseudonym = Pseudonym::random(rng);
-        let enc_pseudonym = encrypt(&pseudonym, &session1_public, rng);
-        entity_data.push(enc_pseudonym);
-        for _ in 0..5 {
-            let data = DataPoint::random(rng);
-            let enc_data = encrypt(&data, &session1_public, rng);
-            entity_data.push(enc_data);
-        }
-        ciphertexts.push(entity_data);
+        data.push(encrypt(&DataPoint::random(rng), &session1_public, rng));
+        pseudonyms.push(encrypt(&Pseudonym::random(rng), &session1_public, rng));
     }
 
     let transcryption_info = TranscryptionInfo::new(
@@ -209,5 +199,11 @@ fn test_batch() {
         &pseudo_secret,
         &enc_secret,
     );
-    let mut new_ciphertexts = transcrypt_batch(&mut ciphertexts, &transcryption_info, rng);
+
+    let rekey_info = RekeyInfo::from(transcryption_info);
+
+    let rekeyed = rekey_batch(&data, &rekey_info, rng);
+    let pseudonymized = pseudonymize_batch(&pseudonyms, &transcryption_info, rng);
+
+    // TODO check that the batch is indeed shuffled
 }
\ No newline at end of file