Merge pull request #639 from Chia-Network/bls-cache-api

update BLSCache::aggregate_verify() function signature
Chia-Network · Aug 2, 2024 · 044bc77 · 044bc77
2 parents ce4fff3 + a767815
commit 044bc77
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 34 deletions.
diff --git a/crates/chia-bls/benches/cache.rs b/crates/chia-bls/benches/cache.rs
@@ -16,7 +16,7 @@ fn cache_benchmark(c: &mut Criterion) {
 
     let mut agg_sig = Signature::default();
     for i in 0..1000 {
-        let derived = sk.derive_hardened(i as u32);
+        let derived = sk.derive_hardened(i);
         let pk = derived.public_key();
         let sig = sign(&derived, msg);
         agg_sig.aggregate(&sig);
@@ -28,43 +28,43 @@ fn cache_benchmark(c: &mut Criterion) {
     c.bench_function("bls_cache.aggregate_verify, 0% cache hits", |b| {
         let mut cache = bls_cache.clone();
         b.iter(|| {
-            assert!(cache.aggregate_verify(&pks, [&msg].iter().cycle(), &agg_sig));
+            assert!(cache.aggregate_verify(pks.iter().zip([&msg].iter().cycle()), &agg_sig));
         });
     });
 
     // populate 10% of keys
-    bls_cache.aggregate_verify(&pks[0..100], [&msg].iter().cycle(), &agg_sig);
+    bls_cache.aggregate_verify(pks[0..100].iter().zip([&msg].iter().cycle()), &agg_sig);
     c.bench_function("bls_cache.aggregate_verify, 10% cache hits", |b| {
         let mut cache = bls_cache.clone();
         b.iter(|| {
-            assert!(cache.aggregate_verify(&pks, [&msg].iter().cycle(), &agg_sig));
+            assert!(cache.aggregate_verify(pks.iter().zip([&msg].iter().cycle()), &agg_sig));
         });
     });
 
     // populate another 10% of keys
-    bls_cache.aggregate_verify(&pks[100..200], [&msg].iter().cycle(), &agg_sig);
+    bls_cache.aggregate_verify(pks[100..200].iter().zip([&msg].iter().cycle()), &agg_sig);
     c.bench_function("bls_cache.aggregate_verify, 20% cache hits", |b| {
         let mut cache = bls_cache.clone();
         b.iter(|| {
-            assert!(cache.aggregate_verify(&pks, [&msg].iter().cycle(), &agg_sig));
+            assert!(cache.aggregate_verify(pks.iter().zip([&msg].iter().cycle()), &agg_sig));
         });
     });
 
     // populate another 30% of keys
-    bls_cache.aggregate_verify(&pks[200..500], [&msg].iter().cycle(), &agg_sig);
+    bls_cache.aggregate_verify(pks[200..500].iter().zip([&msg].iter().cycle()), &agg_sig);
     c.bench_function("bls_cache.aggregate_verify, 50% cache hits", |b| {
         let mut cache = bls_cache.clone();
         b.iter(|| {
-            assert!(cache.aggregate_verify(&pks, [&msg].iter().cycle(), &agg_sig));
+            assert!(cache.aggregate_verify(pks.iter().zip([&msg].iter().cycle()), &agg_sig));
         });
     });
 
     // populate all other keys
-    bls_cache.aggregate_verify(&pks[500..1000], [&msg].iter().cycle(), &agg_sig);
+    bls_cache.aggregate_verify(pks[500..1000].iter().zip([&msg].iter().cycle()), &agg_sig);
     c.bench_function("bls_cache.aggregate_verify, 100% cache hits", |b| {
         let mut cache = bls_cache.clone();
         b.iter(|| {
-            assert!(cache.aggregate_verify(&pks, [&msg].iter().cycle(), &agg_sig));
+            assert!(cache.aggregate_verify(pks.iter().zip([&msg].iter().cycle()), &agg_sig));
         });
     });
 

diff --git a/crates/chia-bls/src/bls_cache.rs b/crates/chia-bls/src/bls_cache.rs
@@ -44,13 +44,12 @@ impl BlsCache {
         self.cache.is_empty()
     }
 
-    pub fn aggregate_verify(
+    pub fn aggregate_verify<Pk: Borrow<PublicKey>, Msg: AsRef<[u8]>>(
         &mut self,
-        pks: impl IntoIterator<Item = impl Borrow<PublicKey>>,
-        msgs: impl IntoIterator<Item = impl AsRef<[u8]>>,
+        pks_msgs: impl IntoIterator<Item = (Pk, Msg)>,
         sig: &Signature,
     ) -> bool {
-        let iter = pks.into_iter().zip(msgs).map(|(pk, msg)| -> GTElement {
+        let iter = pks_msgs.into_iter().map(|(pk, msg)| -> GTElement {
             // Hash pubkey + message
             let mut hasher = Sha256::new();
             hasher.update(pk.borrow().to_bytes());
@@ -124,7 +123,7 @@ impl BlsCache {
             .map(|item| item?.extract())
             .collect::<PyResult<Vec<PyBackedBytes>>>()?;
 
-        Ok(self.aggregate_verify(pks, msgs, sig))
+        Ok(self.aggregate_verify(pks.into_iter().zip(msgs), sig))
     }
 
     #[pyo3(name = "len")]
@@ -180,18 +179,17 @@ pub mod tests {
         let msg = [106; 32];
 
         let sig = sign(&sk, msg);
-        let pk_list = [pk];
-        let msg_list = [msg];
+        let pks_msgs = [(pk, msg)];
 
         // Before we cache anything, it should be empty.
         assert!(bls_cache.is_empty());
 
         // Verify the signature and add to the cache.
-        assert!(bls_cache.aggregate_verify(pk_list, msg_list, &sig));
+        assert!(bls_cache.aggregate_verify(pks_msgs, &sig));
         assert_eq!(bls_cache.len(), 1);
 
         // Now that it's cached, it shouldn't cache it again.
-        assert!(bls_cache.aggregate_verify(pk_list, msg_list, &sig));
+        assert!(bls_cache.aggregate_verify(pks_msgs, &sig));
         assert_eq!(bls_cache.len(), 1);
     }
 
@@ -204,14 +202,13 @@ pub mod tests {
         let msg1 = [106; 32];
 
         let mut agg_sig = sign(&sk1, msg1);
-        let mut pk_list = vec![pk1];
-        let mut msg_list = vec![msg1];
+        let mut pks_msgs = vec![(pk1, msg1)];
 
         // Before we cache anything, it should be empty.
         assert!(bls_cache.is_empty());
 
         // Add the first signature to cache.
-        assert!(bls_cache.aggregate_verify(&pk_list, &msg_list, &agg_sig));
+        assert!(bls_cache.aggregate_verify(pks_msgs.clone(), &agg_sig));
         assert_eq!(bls_cache.len(), 1);
 
         // Try with the first key message pair in the cache but not the second.
@@ -220,21 +217,19 @@ pub mod tests {
         let msg2 = [107; 32];
 
         agg_sig += &sign(&sk2, msg2);
-        pk_list.push(pk2);
-        msg_list.push(msg2);
+        pks_msgs.push((pk2, msg2));
 
-        assert!(bls_cache.aggregate_verify(&pk_list, &msg_list, &agg_sig));
+        assert!(bls_cache.aggregate_verify(pks_msgs.clone(), &agg_sig));
         assert_eq!(bls_cache.len(), 2);
 
         // Try reusing a public key.
         let msg3 = [108; 32];
 
         agg_sig += &sign(&sk2, msg3);
-        pk_list.push(pk2);
-        msg_list.push(msg3);
+        pks_msgs.push((pk2, msg3));
 
         // Verify this signature and add to the cache as well (since it's still a different aggregate).
-        assert!(bls_cache.aggregate_verify(pk_list, msg_list, &agg_sig));
+        assert!(bls_cache.aggregate_verify(pks_msgs, &agg_sig));
         assert_eq!(bls_cache.len(), 3);
     }
 
@@ -253,11 +248,10 @@ pub mod tests {
             let msg = [106; 32];
 
             let sig = sign(&sk, msg);
-            let pk_list = [pk];
-            let msg_list = [msg];
+            let pks_msgs = [(pk, msg)];
 
             // Add to cache by validating them one at a time.
-            assert!(bls_cache.aggregate_verify(pk_list.iter(), msg_list.iter(), &sig));
+            assert!(bls_cache.aggregate_verify(pks_msgs, &sig));
         }
 
         // The cache should be full now.
@@ -282,9 +276,8 @@ pub mod tests {
     fn test_empty_sig() {
         let mut bls_cache = BlsCache::default();
 
-        let pks: [&PublicKey; 0] = [];
-        let msgs: [&[u8]; 0] = [];
+        let pks_msgs: [(&PublicKey, &[u8]); 0] = [];
 
-        assert!(bls_cache.aggregate_verify(pks, msgs, &Signature::default()));
+        assert!(bls_cache.aggregate_verify(pks_msgs, &Signature::default()));
     }
 }