From 142100af5c3febf36c0d4c946c05c351185ab9bf Mon Sep 17 00:00:00 2001
From: sinu <65924192+sinui0@users.noreply.github.com>
Date: Fri, 6 Oct 2023 10:32:58 -0700
Subject: [PATCH] change encrypt_many_blocks to operate in-place

---
 mpz-core/benches/aes.rs | 4 ++--
 mpz-core/src/aes.rs     | 5 ++---
 mpz-core/src/prg.rs     | 5 +++--
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/mpz-core/benches/aes.rs b/mpz-core/benches/aes.rs
index 6946e281..be2acd4b 100644
--- a/mpz-core/benches/aes.rs
+++ b/mpz-core/benches/aes.rs
@@ -18,10 +18,10 @@ fn criterion_benchmark(c: &mut Criterion) {
     c.bench_function("aes::encrypt_many_blocks::<8>", move |bench| {
         let key = rand::random::<Block>();
         let aes = AesEncryptor::new(key);
-        let blks = rand::random::<[Block; 8]>();
+        let mut blks = rand::random::<[Block; 8]>();
 
         bench.iter(|| {
-            let z = aes.encrypt_many_blocks(black_box(blks));
+            let z = aes.encrypt_many_blocks(black_box(&mut blks));
             black_box(z);
         });
     });
diff --git a/mpz-core/src/aes.rs b/mpz-core/src/aes.rs
index 7623f6c7..5c5188c5 100644
--- a/mpz-core/src/aes.rs
+++ b/mpz-core/src/aes.rs
@@ -152,12 +152,11 @@ impl AesEncryptor {
         blk
     }
 
-    /// Encrypt many blocks.
+    /// Encrypt many blocks in-place.
     #[inline(always)]
-    pub fn encrypt_many_blocks<const N: usize>(&self, mut blks: [Block; N]) -> [Block; N] {
+    pub fn encrypt_many_blocks<const N: usize>(&self, blks: &mut [Block; N]) {
         self.0
             .encrypt_blocks(Block::as_generic_array_mut_slice(blks.as_mut_slice()));
-        blks
     }
 
     /// Encrypt slice of blocks in-place.
diff --git a/mpz-core/src/prg.rs b/mpz-core/src/prg.rs
index e748d0e0..0d909f45 100644
--- a/mpz-core/src/prg.rs
+++ b/mpz-core/src/prg.rs
@@ -21,7 +21,7 @@ impl BlockRngCore for PrgCore {
     // Compute [AES(state)..AES(state+8)]
     #[inline(always)]
     fn generate(&mut self, results: &mut Self::Results) {
-        let states = [0; AesEncryptor::AES_BLOCK_COUNT].map(
+        let mut states = [0; AesEncryptor::AES_BLOCK_COUNT].map(
             #[inline(always)]
             |_| {
                 let x = self.state;
@@ -29,7 +29,8 @@ impl BlockRngCore for PrgCore {
                 Block::from(bytemuck::cast::<_, [u8; 16]>([x, 0u64]))
             },
         );
-        *results = bytemuck::cast(self.aes.encrypt_many_blocks(states))
+        self.aes.encrypt_many_blocks(&mut states);
+        *results = bytemuck::cast(states);
     }
 }