feat(aggregation_mode): aggregate proofs in chunks (#1896)

Co-authored-by: MauroFab <[email protected]>
Co-authored-by: Julian Arce <[email protected]>

Author: 3 people

Makefile

@@ -218,9 +218,7 @@ install_aggregation_mode: ## Install the aggregation mode with proving enabled
 	cargo install --path aggregation_mode --features prove,gpu --bin proof_aggregator
 
 agg_mode_write_program_ids: ## Write proof aggregator zkvm programs ids 
-	@cd aggregation_mode && \
-	cargo run --release --bin write_program_image_id_vk_hash
-	
+	@cd aggregation_mode && ./scripts/build_programs.sh
 
 _AGGREGATOR_:
 

aggregation_mode/README.md

@@ -61,11 +61,22 @@ make start_proof_aggregator_gpu AGGREGATOR="sp1|risc0"
 1. Get latest aggregated proof:
 
 ```shell
-cast call 0xc351628EB244ec633d5f21fBD6621e1a683B1181 "currentAggregatedProofNumber()" --rpc-url http://localhost:8545
+cast logs 0xc351628EB244ec633d5f21fBD6621e1a683B1181 'AggregatedProofVerified(bytes32,bytes32)' --from-block 0 --to-block latest --rpc-url http://localhost:8545
 ```
 
-2. Get aggregated proof info:
+## Compiling programs
+
+Whenever any of the programs change, you must recompile them and update their corresponding program ids in `aggregation_mode/program_ids.json`. To do this, run the following command:
 
 ```shell
-cast call 0xc351628EB244ec633d5f21fBD6621e1a683B1181 "getAggregatedProof(uint64)(uint8,bytes32,bytes32)" <AGG_PROOF_NUMBER>  --rpc-url http://localhost:8545
+make agg_mode_write_program_ids
 ```
+
+We are using docker to produce deterministic builds so that the program ids are the same for all systems.
+
+### Updating the program id in `AlignedProofAggregationService` contract
+
+If the program ids have changed, you will also need to update them in the `AlignedProofAggregationService` contract.
+
+-   Risc0: call `setRisc0AggregatorProgramImageId` method with the value of `risc0_root_aggregator_image_id` from `aggregation_mode/program_ids.json`.
+-   SP1: call: `setSP1AggregatorProgramVKHash` method with the value of `sp1_root_aggregator_vk_hash` from `aggregation_mode/program_ids.json`.

aggregation_mode/aggregation_programs/Cargo.toml

@@ -8,4 +8,3 @@ members = ["sp1", "risc0"]
 [patch.crates-io]
 # Adding RISC Zero keccak precompile support
 tiny-keccak = { git = "https://github.com/risc0/tiny-keccak", tag = "tiny-keccak/v2.0.2-risczero.0" }
-

aggregation_mode/aggregation_programs/risc0/Cargo.toml

@@ -16,5 +16,9 @@ lambdaworks-crypto = { git = "https://github.com/lambdaclass/lambdaworks.git", r
 path = "./src/lib.rs"
 
 [[bin]]
-name = "risc0_aggregator_program"
-path = "./src/main.rs"
+name = "risc0_user_proofs_aggregator_program"
+path = "./src/user_proofs_aggregator_main.rs"
+
+[[bin]]
+name = "risc0_chunk_aggregator_program"
+path = "./src/chunk_aggregator_main.rs"

aggregation_mode/aggregation_programs/risc0/src/chunk_aggregator_main.rs

@@ -0,0 +1,50 @@
+#![no_main]
+
+use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
+use risc0_aggregation_program::{ChunkAggregatorInput, Hash32};
+use risc0_zkvm::guest::env;
+
+risc0_zkvm::guest::entry!(main);
+
+// Generated with `make agg_mode_write_program_ids` and copied from program_ids.json
+pub const USER_PROOFS_AGGREGATOR_PROGRAM_IMAGE_ID: [u8; 32] = [
+    83, 145, 39, 254, 127, 217, 146, 127, 63, 217, 69, 190, 11, 204, 170, 138, 215, 35, 175, 246,
+    209, 154, 52, 243, 85, 37, 177, 147, 22, 153, 155, 156,
+];
+
+fn main() {
+    let input = env::read::<ChunkAggregatorInput>();
+
+    let mut leaves: Vec<Hash32> = vec![];
+
+    for (proof, leaves_commitment) in input.proofs_and_leaves_commitment {
+        let image_id = proof.image_id;
+
+        // Ensure the aggregated chunk originates from the L1 aggregation program.
+        // This validation step guarantees that the proof was genuinely verified
+        // by this program. Without this check, a different program using the
+        // same public inputs could bypass verification.
+        assert!(image_id == USER_PROOFS_AGGREGATOR_PROGRAM_IMAGE_ID);
+
+        // Ensure the committed root matches the root of the provided leaves
+        let merkle_root: [u8; 32] = proof
+            .public_inputs
+            .clone()
+            .try_into()
+            .expect("Public input to be the chunk merkle root");
+
+        let leaves_commitment: Vec<Hash32> =
+            leaves_commitment.into_iter().map(|el| Hash32(el)).collect();
+        let merkle_tree = MerkleTree::<Hash32>::build(&leaves_commitment).unwrap();
+        assert!(merkle_root == merkle_tree.root);
+
+        leaves.extend(leaves_commitment);
+
+        // finally verify the proof
+        env::verify(image_id, &proof.public_inputs).expect("proof to be verified correctly");
+    }
+
+    let merkle_tree = MerkleTree::<Hash32>::build(&leaves).unwrap();
+
+    env::commit_slice(&merkle_tree.root);
+}

aggregation_mode/aggregation_programs/risc0/src/lib.rs

@@ -67,6 +67,58 @@ impl IsMerkleTreeBackend for Risc0ImageIdAndPubInputs {
 }
 
 #[derive(Serialize, Deserialize)]
-pub struct Input {
+pub struct Hash32(pub [u8; 32]);
+
+// Note: this MerkleTreeBackend is defined in three locations
+// - aggregation_mode/src/aggregators/mod.rs
+// - aggregation_mode/src/aggregators/risc0_aggregator.rs and
+// - aggregation_mode/src/aggregators/sp1_aggregator.rs
+// All 3 implementations should match
+// The definition on aggregator/mod.rs supports taking proofs from both Risc0 and SP1,
+// Additionally, a version that takes the leaves as already hashed data is defined on:
+// - batcher/aligned-sdk/src/sdk/aggregation.rs
+// This one is used in the SDK since
+// the user may not have access to the proofs that they didn't submit
+impl IsMerkleTreeBackend for Hash32 {
+    type Data = Hash32;
+    type Node = [u8; 32];
+
+    fn hash_data(leaf: &Self::Data) -> Self::Node {
+        leaf.0
+    }
+
+    /// Computes a commutative Keccak256 hash, ensuring H(a, b) == H(b, a).
+    ///
+    /// See: https://docs.openzeppelin.com/contracts/5.x/api/utils#Hashes
+    ///
+    /// Source: https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/utils/cryptography/Hashes.sol#L17-L19
+    ///
+    /// Compliant with OpenZeppelin's `processProofCalldata` function from MerkleProof.sol.
+    ///
+    /// See: https://docs.openzeppelin.com/contracts/5.x/api/utils#MerkleProof
+    ///
+    /// Source: https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/utils/cryptography/MerkleProof.sol#L114-L128
+    fn hash_new_parent(child_1: &Self::Node, child_2: &Self::Node) -> Self::Node {
+        let mut hasher = Keccak::v256();
+        if child_1 < child_2 {
+            hasher.update(child_1);
+            hasher.update(child_2);
+        } else {
+            hasher.update(child_2);
+            hasher.update(child_1);
+        }
+        let mut hash = [0u8; 32];
+        hasher.finalize(&mut hash);
+        hash
+    }
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct UserProofsAggregatorInput {
     pub proofs_image_id_and_pub_inputs: Vec<Risc0ImageIdAndPubInputs>,
 }
+
+#[derive(Serialize, Deserialize)]
+pub struct ChunkAggregatorInput {
+    pub proofs_and_leaves_commitment: Vec<(Risc0ImageIdAndPubInputs, Vec<[u8; 32]>)>,
+}

aggregation_mode/aggregation_programs/risc0/src/main.rs

@@ -0,0 +1,22 @@
+#![no_main]
+
+use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
+use risc0_aggregation_program::{Risc0ImageIdAndPubInputs, UserProofsAggregatorInput};
+use risc0_zkvm::guest::env;
+
+risc0_zkvm::guest::entry!(main);
+
+fn main() {
+    let input = env::read::<UserProofsAggregatorInput>();
+
+    for proof in &input.proofs_image_id_and_pub_inputs {
+        env::verify(proof.image_id.clone(), &proof.public_inputs)
+            .expect("proof to be verified correctly");
+    }
+
+    let merkle_tree =
+        MerkleTree::<Risc0ImageIdAndPubInputs>::build(&input.proofs_image_id_and_pub_inputs)
+            .unwrap();
+
+    env::commit_slice(&merkle_tree.root);
+}

aggregation_mode/aggregation_programs/risc0/src/user_proofs_aggregator_main.rs

@@ -16,5 +16,9 @@ lambdaworks-crypto = { git = "https://github.com/lambdaclass/lambdaworks.git", r
 path = "./src/lib.rs"
 
 [[bin]]
-name = "sp1_aggregator_program"
-path = "./src/main.rs"
+name = "sp1_user_proofs_aggregator_program"
+path = "./src/user_proofs_aggregator_main.rs"
+
+[[bin]]
+name = "sp1_chunk_aggregator_program"
+path = "./src/chunk_aggregator_main.rs"

aggregation_mode/aggregation_programs/sp1/Cargo.toml

@@ -0,0 +1,50 @@
+#![no_main]
+sp1_zkvm::entrypoint!(main);
+
+use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
+use sha2::{Digest, Sha256};
+use sp1_aggregation_program::{ChunkAggregatorInput, Hash32};
+
+// Generated with `make agg_mode_write_program_ids` and copied from program_ids.json
+pub const USER_PROOFS_AGGREGATOR_PROGRAM_VK_HASH: [u32; 8] = [
+    684911098, 272834847, 1514192666, 1104122402, 1853418149, 488480116, 2005139814, 1901405498,
+];
+
+pub fn main() {
+    let input = sp1_zkvm::io::read::<ChunkAggregatorInput>();
+
+    let mut leaves = vec![];
+
+    // Verify the proofs.
+    for (proof, leaves_commitment) in input.proofs_and_leaves_commitment {
+        let vkey = proof.vk;
+        let public_values_digest = Sha256::digest(&proof.public_inputs);
+
+        // Ensure the aggregated chunk originates from the user proofs aggregation program.
+        // This validation step guarantees that the proof was genuinely verified
+        // by this program. Without this check, a different program using the
+        // same public inputs could bypass verification.
+        assert!(proof.vk == USER_PROOFS_AGGREGATOR_PROGRAM_VK_HASH);
+
+        let merkle_root: [u8; 32] = proof
+            .public_inputs
+            .clone()
+            .try_into()
+            .expect("Public input to be the hash of the chunk tree");
+
+        // Reconstruct the merkle tree and verify that the roots match
+        let leaves_commitment: Vec<Hash32> =
+            leaves_commitment.into_iter().map(|el| Hash32(el)).collect();
+        let merkle_tree: MerkleTree<Hash32> = MerkleTree::build(&leaves_commitment).unwrap();
+        assert!(merkle_tree.root == merkle_root);
+
+        leaves.extend(leaves_commitment);
+
+        sp1_zkvm::lib::verify::verify_sp1_proof(&vkey, &public_values_digest.into());
+    }
+
+    // Finally, compute the final merkle root with all the leaves
+    let merkle_tree: MerkleTree<Hash32> = MerkleTree::build(&leaves).unwrap();
+
+    sp1_zkvm::io::commit_slice(&merkle_tree.root);
+}