Faster PoT verification without allocation or std

Author: nazar-pc

Cargo.lock

@@ -41,6 +41,7 @@ cc = "1.1.23"
 chacha20 = { version = "0.9.1", default-features = false }
 clap = "4.5.18"
 core_affinity = "0.8.1"
+cpufeatures = "0.2.17"
 criterion = { version = "0.5.1", default-features = false }
 cross-domain-message-gossip = { version = "0.1.0", path = "domains/client/cross-domain-message-gossip" }
 derive_more = { version = "1.0.0", default-features = false }

Cargo.toml

@@ -272,7 +272,7 @@ impl PotVerifier {
             drop(cache);
 
             let verified_successfully =
-                subspace_proof_of_time::verify(seed, slot_iterations, checkpoints.as_slice())
+                subspace_proof_of_time::verify(seed, slot_iterations, checkpoints)
                     .unwrap_or_default();
 
             if !verified_successfully {

crates/sc-proof-of-time/src/verifier.rs

@@ -2,9 +2,9 @@
 
 use crate::hashes::{blake3_hash, blake3_hash_list, Blake3Hash};
 use crate::Randomness;
-use core::fmt;
 use core::num::NonZeroU8;
 use core::str::FromStr;
+use core::{fmt, mem};
 use derive_more::{AsMut, AsRef, Deref, DerefMut, From};
 use parity_scale_codec::{Decode, Encode, MaxEncodedLen};
 use scale_info::TypeInfo;
@@ -213,6 +213,7 @@ impl PotSeed {
     TypeInfo,
     MaxEncodedLen,
 )]
+#[repr(C)]
 pub struct PotOutput([u8; Self::SIZE]);
 
 impl fmt::Debug for PotOutput {
@@ -291,6 +292,20 @@ impl PotOutput {
         seed.copy_from_slice(&hash[..Self::SIZE]);
         seed
     }
+
+    /// Convenient conversion from slice of underlying representation for efficiency purposes
+    #[inline(always)]
+    pub const fn slice_from_repr(value: &[[u8; Self::SIZE]]) -> &[Self] {
+        // SAFETY: `PotOutput` is `#[repr(C)]` and guaranteed to have the same memory layout
+        unsafe { mem::transmute(value) }
+    }
+
+    /// Convenient conversion to slice of underlying representation for efficiency purposes
+    #[inline(always)]
+    pub const fn repr_from_slice(value: &[Self]) -> &[[u8; Self::SIZE]] {
+        // SAFETY: `PotOutput` is `#[repr(C)]` and guaranteed to have the same memory layout
+        unsafe { mem::transmute(value) }
+    }
 }
 
 /// Proof of time checkpoints, result of proving

crates/subspace-core-primitives/src/pot.rs

@@ -3,7 +3,7 @@ name = "subspace-proof-of-time"
 description = "Subspace proof of time implementation"
 license = "0BSD"
 version = "0.1.0"
-authors = ["Rahul Subramaniyam <rahulksnv@gmail.com>"]
+authors = ["Nazar Mokrynskyi <nazar@mokrynskyi.com>"]
 edition = "2021"
 include = [
     "/src",
@@ -19,13 +19,14 @@ aes.workspace = true
 subspace-core-primitives.workspace = true
 thiserror.workspace = true
 
-# This is required to for benchmark dependency features to work correctly
-rand = { workspace = true, optional = true }
+[target.'cfg(target_arch = "x86_64")'.dependencies]
+cpufeatures = { workspace = true }
 
 [dev-dependencies]
 core_affinity.workspace = true
 criterion.workspace = true
-rand.workspace = true
+rand_core = { workspace = true }
+rand_chacha = { workspace = true }
 
 [[bench]]
 name = "pot"
@@ -34,12 +35,3 @@ harness = false
 [[bench]]
 name = "pot-compare-cpu-cores"
 harness = false
-
-[features]
-default = ["std"]
-std = [
-    "subspace-core-primitives/std",
-    "thiserror/std",
-    "rand?/std",
-    "rand?/std_rng",
-]

crates/subspace-proof-of-time/Cargo.toml

@@ -1,12 +1,15 @@
 use core::num::NonZeroU32;
-use criterion::{black_box, criterion_group, criterion_main, Criterion};
-use rand::{thread_rng, Rng};
+use criterion::{criterion_group, criterion_main, Criterion};
+use rand_chacha::ChaCha8Rng;
+use rand_core::{RngCore, SeedableRng};
+use std::hint::black_box;
 use subspace_core_primitives::pot::PotSeed;
 use subspace_proof_of_time::prove;
 
 fn criterion_benchmark(c: &mut Criterion) {
+    let mut rng = ChaCha8Rng::from_seed(Default::default());
     let mut seed = PotSeed::default();
-    thread_rng().fill(seed.as_mut());
+    rng.fill_bytes(seed.as_mut());
     // About 1s on 6.0 GHz Raptor Lake CPU (14900K)
     let pot_iterations = NonZeroU32::new(200_032_000).expect("Not zero; qed");
 

crates/subspace-proof-of-time/benches/pot-compare-cpu-cores.rs

@@ -1,12 +1,15 @@
 use core::num::NonZeroU32;
-use criterion::{black_box, criterion_group, criterion_main, Criterion};
-use rand::{thread_rng, Rng};
+use criterion::{criterion_group, criterion_main, Criterion};
+use rand_chacha::ChaCha8Rng;
+use rand_core::{RngCore, SeedableRng};
+use std::hint::black_box;
 use subspace_core_primitives::pot::PotSeed;
 use subspace_proof_of_time::{prove, verify};
 
 fn criterion_benchmark(c: &mut Criterion) {
+    let mut rng = ChaCha8Rng::from_seed(Default::default());
     let mut seed = PotSeed::default();
-    thread_rng().fill(seed.as_mut());
+    rng.fill_bytes(seed.as_mut());
     // About 1s on 6.0 GHz Raptor Lake CPU (14900K)
     let pot_iterations = NonZeroU32::new(200_032_000).expect("Not zero; qed");
 
@@ -23,7 +26,7 @@ fn criterion_benchmark(c: &mut Criterion) {
             black_box(verify(
                 black_box(seed),
                 black_box(pot_iterations),
-                black_box(&*checkpoints),
+                black_box(&checkpoints),
             ))
             .unwrap();
         })

crates/subspace-proof-of-time/benches/pot.rs

@@ -1,24 +1,22 @@
 //! AES related functionality.
 
-#[cfg(all(feature = "std", target_arch = "x86_64"))]
+#[cfg(target_arch = "x86_64")]
 mod x86_64;
 
-#[cfg(not(feature = "std"))]
-extern crate alloc;
-
 use aes::cipher::array::Array;
 use aes::cipher::{BlockCipherDecrypt, BlockCipherEncrypt, KeyInit};
 use aes::Aes128;
-#[cfg(not(feature = "std"))]
-use alloc::vec::Vec;
 use subspace_core_primitives::pot::{PotCheckpoints, PotKey, PotOutput, PotSeed};
 
 /// Creates the AES based proof.
 #[inline(always)]
 pub(crate) fn create(seed: PotSeed, key: PotKey, checkpoint_iterations: u32) -> PotCheckpoints {
-    #[cfg(all(feature = "std", target_arch = "x86_64"))]
-    if std::is_x86_feature_detected!("aes") {
-        return unsafe { x86_64::create(seed.as_ref(), key.as_ref(), checkpoint_iterations) };
+    #[cfg(target_arch = "x86_64")]
+    {
+        cpufeatures::new!(has_aes, "aes");
+        if has_aes::get() {
+            return unsafe { x86_64::create(seed.as_ref(), key.as_ref(), checkpoint_iterations) };
+        }
     }
 
     create_generic(seed, key, checkpoint_iterations)
@@ -48,27 +46,26 @@ fn create_generic(seed: PotSeed, key: PotKey, checkpoint_iterations: u32) -> Pot
 pub(crate) fn verify_sequential(
     seed: PotSeed,
     key: PotKey,
-    checkpoints: &[PotOutput],
+    checkpoints: &PotCheckpoints,
     checkpoint_iterations: u32,
 ) -> bool {
     assert_eq!(checkpoint_iterations % 2, 0);
 
     let key = Array::from(*key);
     let cipher = Aes128::new(&key);
 
-    let mut inputs = Vec::with_capacity(checkpoints.len());
-    inputs.push(Array::from(*seed));
-    for &checkpoint in checkpoints.iter().rev().skip(1).rev() {
-        inputs.push(Array::from(*checkpoint));
-    }
-    let mut outputs = checkpoints
-        .iter()
-        .map(|&checkpoint| Array::from(*checkpoint))
-        .collect::<Vec<_>>();
+    let mut inputs = [[0u8; 16]; PotCheckpoints::NUM_CHECKPOINTS.get() as usize];
+    inputs[0] = *seed;
+    inputs[1..].copy_from_slice(PotOutput::repr_from_slice(
+        &checkpoints[..PotCheckpoints::NUM_CHECKPOINTS.get() as usize - 1],
+    ));
+
+    let mut outputs = [[0u8; 16]; PotCheckpoints::NUM_CHECKPOINTS.get() as usize];
+    outputs.copy_from_slice(PotOutput::repr_from_slice(checkpoints.as_slice()));
 
     for _ in 0..checkpoint_iterations / 2 {
-        cipher.encrypt_blocks(&mut inputs);
-        cipher.decrypt_blocks(&mut outputs);
+        cipher.encrypt_blocks(Array::cast_slice_from_core_mut(&mut inputs));
+        cipher.decrypt_blocks(Array::cast_slice_from_core_mut(&mut outputs));
     }
 
     inputs == outputs
@@ -77,6 +74,7 @@ pub(crate) fn verify_sequential(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use subspace_core_primitives::pot::PotOutput;
 
     const SEED: [u8; 16] = [
         0xd6, 0x66, 0xcc, 0xd8, 0xd5, 0x93, 0xc2, 0x3d, 0xa8, 0xdb, 0x6b, 0x5b, 0x14, 0x13, 0xb1,
@@ -100,7 +98,7 @@ mod tests {
     fn test_create_verify() {
         let seed = PotSeed::from(SEED);
         let key = PotKey::from(KEY);
-        let checkpoint_iterations = 100;
+        let checkpoint_iterations = 20;
 
         // Can encrypt/decrypt.
         let checkpoints = create(seed, key, checkpoint_iterations);
@@ -112,7 +110,7 @@ mod tests {
         assert!(verify_sequential(
             seed,
             key,
-            &*checkpoints,
+            &checkpoints,
             checkpoint_iterations,
         ));
 
@@ -122,37 +120,37 @@ mod tests {
         assert!(!verify_sequential(
             seed,
             key,
-            &*checkpoints_1,
+            &checkpoints_1,
             checkpoint_iterations,
         ));
 
         // Decryption with wrong number of iterations fails.
         assert!(!verify_sequential(
             seed,
             key,
-            &*checkpoints,
+            &checkpoints,
             checkpoint_iterations + 2,
         ));
         assert!(!verify_sequential(
             seed,
             key,
-            &*checkpoints,
+            &checkpoints,
             checkpoint_iterations - 2,
         ));
 
         // Decryption with wrong seed fails.
         assert!(!verify_sequential(
             PotSeed::from(SEED_1),
             key,
-            &*checkpoints,
+            &checkpoints,
             checkpoint_iterations,
         ));
 
         // Decryption with wrong key fails.
         assert!(!verify_sequential(
             seed,
             PotKey::from(KEY_1),
-            &*checkpoints,
+            &checkpoints,
             checkpoint_iterations,
         ));
     }

crates/subspace-proof-of-time/src/aes.rs

@@ -1,10 +1,11 @@
 //! Proof of time implementation.
 
-#![cfg_attr(not(feature = "std"), no_std)]
+#![no_std]
+
 mod aes;
 
 use core::num::NonZeroU32;
-use subspace_core_primitives::pot::{PotCheckpoints, PotOutput, PotSeed};
+use subspace_core_primitives::pot::{PotCheckpoints, PotSeed};
 
 /// Proof of time error
 #[derive(Debug, thiserror::Error)]
@@ -46,7 +47,7 @@ pub fn prove(seed: PotSeed, iterations: NonZeroU32) -> Result<PotCheckpoints, Po
 pub fn verify(
     seed: PotSeed,
     iterations: NonZeroU32,
-    checkpoints: &[PotOutput],
+    checkpoints: &PotCheckpoints,
 ) -> Result<bool, PotError> {
     let num_checkpoints = checkpoints.len() as u32;
     if iterations.get() % (num_checkpoints * 2) != 0 {