From 38d02ff7c69d757deb0290786bdf0de38351534b Mon Sep 17 00:00:00 2001
From: Nazar Mokrynskyi <nazar@mokrynskyi.com>
Date: Thu, 27 Jul 2023 13:05:34 +0300
Subject: [PATCH] Reduce number of PoT checkpoints 2x

---
 crates/subspace-proof-of-time/benches/pot.rs |  2 +-
 crates/subspace-proof-of-time/src/pot_aes.rs | 22 +++++++++++++-------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/crates/subspace-proof-of-time/benches/pot.rs b/crates/subspace-proof-of-time/benches/pot.rs
index adbb6ed68f..4a27cbc0dd 100644
--- a/crates/subspace-proof-of-time/benches/pot.rs
+++ b/crates/subspace-proof-of-time/benches/pot.rs
@@ -11,7 +11,7 @@ fn criterion_benchmark(c: &mut Criterion) {
     let slot_number = 1;
     let mut injected_block_hash = BlockHash::default();
     thread_rng().fill(injected_block_hash.as_mut());
-    let checkpoints = 16;
+    let checkpoints = 8;
     // About 1s on 5.5 GHz Raptor Lake CPU
     let iterations = 166_000_000;
     let proof_of_time_sequential = ProofOfTime::new(1, iterations);
diff --git a/crates/subspace-proof-of-time/src/pot_aes.rs b/crates/subspace-proof-of-time/src/pot_aes.rs
index efb257383e..79bb6a8fd8 100644
--- a/crates/subspace-proof-of-time/src/pot_aes.rs
+++ b/crates/subspace-proof-of-time/src/pot_aes.rs
@@ -3,7 +3,7 @@
 extern crate alloc;
 
 use aes::cipher::generic_array::GenericArray;
-use aes::cipher::{BlockEncrypt, KeyInit};
+use aes::cipher::{BlockDecrypt, BlockEncrypt, KeyInit};
 use aes::Aes128;
 use alloc::vec::Vec;
 use subspace_core_primitives::{PotBytes, PotCheckpoint, PotKey, PotSeed};
@@ -31,12 +31,16 @@ pub(crate) fn create(
 }
 
 /// Verifies the AES based proof sequentially.
+///
+/// Panics if `checkpoint_iterations` is not a multiple of `2`.
 pub(crate) fn verify_sequential(
     seed: &PotSeed,
     key: &PotKey,
     checkpoints: &[PotCheckpoint],
     checkpoint_iterations: u32,
 ) -> bool {
+    assert_eq!(checkpoint_iterations % 2, 0);
+
     let key = GenericArray::from(PotBytes::from(*key));
     let cipher = Aes128::new(&key);
 
@@ -45,15 +49,17 @@ pub(crate) fn verify_sequential(
     for checkpoint in checkpoints.iter().rev().skip(1).rev() {
         inputs.push(GenericArray::from(PotBytes::from(*checkpoint)));
     }
+    let mut outputs = checkpoints
+        .iter()
+        .map(|checkpoint| GenericArray::from(PotBytes::from(*checkpoint)))
+        .collect::<Vec<_>>();
 
-    for _ in 0..checkpoint_iterations {
+    for _ in 0..checkpoint_iterations / 2 {
         cipher.encrypt_blocks(&mut inputs);
+        cipher.decrypt_blocks(&mut outputs);
     }
 
-    inputs
-        .iter()
-        .zip(checkpoints)
-        .all(|(a, b)| a.as_slice() == b.as_ref())
+    inputs == outputs
 }
 
 #[cfg(test)]
@@ -111,13 +117,13 @@ mod tests {
             &seed,
             &key,
             &checkpoints,
-            checkpoint_iterations + 1
+            checkpoint_iterations + 2
         ));
         assert!(!verify_sequential(
             &seed,
             &key,
             &checkpoints,
-            checkpoint_iterations - 1
+            checkpoint_iterations - 2
         ));
 
         // Decryption with wrong seed fails.