From dc9c32c93e5076e833ac8876afd3a75c511d6ada Mon Sep 17 00:00:00 2001
From: arnaucube <git@arnaucube.com>
Date: Fri, 27 Sep 2024 07:37:18 +0200
Subject: [PATCH] Update aggregate_signals, Add test_recursion

---
 src/recursion.rs | 86 ++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 83 insertions(+), 3 deletions(-)

diff --git a/src/recursion.rs b/src/recursion.rs
index bf0e836..3e5d3f7 100644
--- a/src/recursion.rs
+++ b/src/recursion.rs
@@ -40,6 +40,7 @@ impl AccessSet {
             .collect();
 
         let proof_target0 = builder.add_virtual_proof_with_pis(&verifier_data.common);
+        builder.register_public_inputs(&proof_target0.public_inputs);
         pw.set_proof_with_pis_target(
             &proof_target0,
             &ProofWithPublicInputs {
@@ -47,7 +48,12 @@ impl AccessSet {
                 public_inputs: public_inputs0,
             },
         )?;
+        let vd_target =
+            builder.add_virtual_verifier_data(verifier_data.common.config.fri_config.cap_height);
+        pw.set_verifier_data_target(&vd_target, &verifier_data.verifier_only)?;
+
         let proof_target1 = builder.add_virtual_proof_with_pis(&verifier_data.common);
+        builder.register_public_inputs(&proof_target1.public_inputs);
         pw.set_proof_with_pis_target(
             &proof_target1,
             &ProofWithPublicInputs {
@@ -55,9 +61,10 @@ impl AccessSet {
                 public_inputs: public_inputs1,
             },
         )?;
-
         let vd_target =
             builder.add_virtual_verifier_data(verifier_data.common.config.fri_config.cap_height);
+        pw.set_verifier_data_target(&vd_target, &verifier_data.verifier_only)?;
+
         pw.set_cap_target(
             &vd_target.constants_sigmas_cap,
             &verifier_data.verifier_only.constants_sigmas_cap,
@@ -67,9 +74,9 @@ impl AccessSet {
         builder.verify_proof::<C>(&proof_target1, &vd_target, &verifier_data.common);
 
         let data = builder.build();
-        let recursive_proof = data.prove(pw).unwrap();
+        let recursive_proof = data.prove(pw)?;
 
-        data.verify(recursive_proof.clone()).unwrap();
+        data.verify(recursive_proof.clone())?;
 
         Ok((
             signal0.nullifier,
@@ -79,3 +86,76 @@ impl AccessSet {
         ))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use plonky2::field::types::{Field, Sample};
+    use plonky2::hash::merkle_tree::MerkleTree;
+    use plonky2::hash::poseidon::PoseidonHash;
+    use plonky2::plonk::config::Hasher;
+    use plonky2::plonk::proof::ProofWithPublicInputs;
+    use std::time::Instant;
+
+    use crate::access_set::AccessSet;
+    use crate::signal::{Digest, F};
+
+    #[test]
+    fn test_recursion() -> Result<()> {
+        let n = 1 << 20;
+        let private_keys: Vec<Digest> = (0..n).map(|_| F::rand_array()).collect();
+        let public_keys: Vec<Vec<F>> = private_keys
+            .iter()
+            .map(|&sk| {
+                PoseidonHash::hash_no_pad(&[sk, [F::ZERO; 4]].concat())
+                    .elements
+                    .to_vec()
+            })
+            .collect();
+        let access_set = AccessSet(MerkleTree::new(public_keys, 0));
+
+        // first proof
+        let i0 = 12;
+        let topic0 = F::rand_array();
+        let start = Instant::now();
+        let (signal0, vd0) = access_set.make_signal(private_keys[i0], topic0, i0)?;
+        println!("generate proof: {:?}", start.elapsed());
+        access_set.verify_signal(topic0, signal0.clone(), &vd0)?;
+
+        // second proof
+        let i1 = 42;
+        let topic1 = F::rand_array();
+        let start = Instant::now();
+        let (signal1, vd1) = access_set.make_signal(private_keys[i1], topic1, i1)?;
+        println!("generate proof: {:?}", start.elapsed());
+        access_set.verify_signal(topic1, signal1.clone(), &vd1)?;
+
+        // generate recursive proof
+        let start = Instant::now();
+        let (nullifier0, nullifier1, recursive_proof, vd2) =
+            access_set.aggregate_signals(topic0, signal0, topic1, signal1, &vd0)?;
+        println!("aggregate_signals (recursive prove): {:?}", start.elapsed());
+
+        // verify recursive proof
+        let public_inputs: Vec<F> = access_set
+            .0
+            .cap
+            .0
+            .iter()
+            .flat_map(|h| h.elements)
+            .chain(nullifier0)
+            .chain(topic0)
+            .chain(access_set.0.cap.0.iter().flat_map(|h| h.elements))
+            .chain(nullifier1)
+            .chain(topic1)
+            .collect();
+
+        let start = Instant::now();
+        vd2.verify(ProofWithPublicInputs {
+            proof: recursive_proof,
+            public_inputs,
+        })?;
+        println!("verify recursive proof: {:?}", start.elapsed());
+        Ok(())
+    }
+}