aggregator added

README.md

@@ -102,3 +102,17 @@ aligned submit \
 --network holesky \
 --rpc_url https://ethereum-holesky-rpc.publicnode.com
 ```
+
+### Sending an Aggregated Proof
+
+```sh
+aligned submit \
+--proving_system SP1 \
+--proof ./data/foods-med.sp1.proof \
+--public_input ./data/foods-med.sp1.pub \
+--vm_program ./elf/riscv32im-succinct-aggregator-elf \
+--batcher_url wss://batcher.alignedlayer.com \
+--keystore_path ./secrets/wallet.json \
+--network holesky \
+--rpc_url https://ethereum-holesky-rpc.publicnode.com
+```

aligned_verification_data/7ed9ac7f_0.json

@@ -0,0 +1,9 @@
+{
+  "batch_merkle_root": "7ed9ac7f9ee3271eb194871bf13945a3eac9e54ae3c64f0af1497fd36f46c98e",
+  "merkle_proof": "3d203d795d2a9595be3a27f60af35d9ea28b915de48db93406e880ea061d203c1b5847e001b3c6e8acaaaeac03f7691fadef7a593d07c2e8534aeb93e5e7cc2191d8a57de663e45c238dc51c7ad6ebfd33fef2a1869bc048ee8b9c2328fc78d8",
+  "program_id_commitment": "4491254d8ff491a773e60e423039e35f008d38eb84bbc51d478839d635622feb",
+  "proof_commitment": "a8287c41e2ea521ec4e017fa8ba28e94f08f8c5325e4e0caaab7a61cb9feb276",
+  "proof_generator_addr": "f39fd6e51aad88f6f4ce6ab8827279cfffb92266",
+  "pub_input_commitment": "0000000000000000000000000000000000000000000000000000000000000000",
+  "verification_data_batch_index": 0
+}

aligned_verification_data/README.md

@@ -1,4 +1,5 @@
 # Verifications
 
-- `6cd22ec2`: Fibonacci
-- `56a9b870`: Foods (smol)
+- `6cd22ec2`: Fibonacci at <https://explorer.alignedlayer.com/batches/0x6cd22ec2d108ced3dcfcf25c2ef3f33e90aa390d74d87d0cbe1c46ea21a1b7b3>
+- `56a9b870`: Foods (smol) at <https://explorer.alignedlayer.com/batches/0x56a9b8708e2e236f394bcae84706c0a97ff13b83bc920917155b3b73c311c7e4>
+- `7ed9ac7f`: Foods (med & aggregated) at <https://explorer.alignedlayer.com/batches/0x7ed9ac7f9ee3271eb194871bf13945a3eac9e54ae3c64f0af1497fd36f46c98e>

lib/src/lib.rs

@@ -5,13 +5,14 @@ alloy_sol_types::sol! {
     }
 }
 
-pub fn compute_best_sample(samples: &[Vec<f32>], query: &[f32]) -> (usize, f32) {
+pub fn compute_best_sample(samples: &[Vec<f32>], query: &[f32]) -> usize {
     samples
         .iter()
         .map(|sample| sample.iter().zip(query).map(|(a, b)| a * b).sum::<f32>())
         .enumerate()
         .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
         .unwrap()
+        .0
 }
 
 pub fn iterative_similarity_search(
@@ -27,10 +28,9 @@ pub fn iterative_similarity_search(
         }
         current_samples = best_samples
             .iter()
-            .map(|&(idx, _)| current_samples[idx].clone())
+            .map(|&idx| current_samples[idx].clone())
             .collect::<Vec<_>>();
     }
 
-    let (best_idx, _) = compute_best_sample(&current_samples, &query);
-    best_idx
+    compute_best_sample(&current_samples, &query)
 }

program/src/main.rs

@@ -10,7 +10,7 @@ pub fn main() {
     let samples = sp1_zkvm::io::read::<Vec<Vec<f32>>>();
     let query = sp1_zkvm::io::read::<Vec<f32>>();
 
-    let (idx, _) = compute_best_sample(&samples, &query);
+    let idx = compute_best_sample(&samples, &query);
 
     let bytes = PublicValuesStruct::abi_encode(&PublicValuesStruct { idx: idx as u32 });
     sp1_zkvm::io::commit_slice(&bytes);

script/src/bin/main.rs

@@ -12,7 +12,9 @@
 
 use alloy_sol_types::SolType;
 use clap::Parser;
-use sp1_sdk::{ProverClient, SP1Stdin};
+use sp1_sdk::{
+    HashableKey, ProverClient, SP1Proof, SP1ProofWithPublicValues, SP1Stdin, SP1VerifyingKey,
+};
 use std::path::PathBuf;
 
 use zkvdb_embedder::{Data, EmbeddedData};
@@ -21,6 +23,14 @@ use zkvdb_lib::PublicValuesStruct;
 /// The ELF (executable and linkable format) file for the Succinct RISC-V zkVM.
 pub const PROGRAM_ELF: &[u8] = include_bytes!("../../../elf/riscv32im-succinct-zkvm-elf");
 
+/// The ELF (executable and linkable format) file for the Succinct RISC-V zkVM aggregator.
+pub const AGGREGATOR_ELF: &[u8] = include_bytes!("../../../elf/riscv32im-succinct-aggregator-elf");
+
+struct AggregationInput {
+    pub proof: SP1ProofWithPublicValues,
+    pub vk: SP1VerifyingKey,
+}
+
 /// The arguments for the command.
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about = None)]
@@ -33,19 +43,27 @@ struct Args {
 
     #[clap(short, long, default_value = "../data/foods-smol.json")]
     path: PathBuf,
+}
 
-    #[clap(short, long, default_value = "2")]
-    top_k: usize,
+enum ExecutionType {
+    Execute,
+    Prove,
 }
 
 fn main() {
     sp1_sdk::utils::setup_logger();
-
     let args = Args::parse();
+
+    ////////// Parse execution type.
     if args.execute == args.prove {
         eprintln!("Error: You must specify either --execute or --prove");
         std::process::exit(1);
     }
+    let exec_type = if args.execute {
+        ExecutionType::Execute
+    } else {
+        ExecutionType::Prove
+    };
 
     ///////// Setup the prover client.
     let client = ProverClient::new();
@@ -66,64 +84,115 @@ fn main() {
         std::fs::read(args.path.with_extension("query.json")).expect("Failed to read the file");
     let query: Vec<f32> = serde_json::from_slice(&query_bytes).expect("Failed to parse JSON");
 
-    // Read k from args
-    let k: usize = args.top_k;
-    assert!(
-        k <= samples.len(),
-        "k must be less than or equal to the number of samples"
-    );
-
-    // Pass inputs to zkVM
-    let mut stdin = SP1Stdin::new();
-    stdin.write(&samples);
-    stdin.write(&query);
-    stdin.write(&k);
-    println!("Inputs prepared.");
-
-    if args.execute {
-        println!("Executing program.");
-        // Execute the program
-        let (output, report) = client.execute(PROGRAM_ELF, stdin).run().unwrap();
-        println!("Program executed successfully.");
-
-        // Read the output.
-        let decoded = PublicValuesStruct::abi_decode(output.as_slice(), true).unwrap();
-        let PublicValuesStruct { idx } = decoded;
-        println!("Closest idx: {}", idx);
-
-        let expected_dest = zkvdb_lib::similarity_search(samples, query, k);
-        assert_eq!(dest, expected_dest);
-        println!("Values are correct!");
-
-        // Record the number of cycles executed.
-        println!("Number of cycles: {}", report.total_instruction_count());
-    } else {
-        println!("Proving program.");
-        // setup the program for proving.
-        let (pk, vk) = client.setup(PROGRAM_ELF);
-
-        // generate the proof
-        let proof = client
-            .prove(&pk, stdin)
-            .compressed()
-            .run()
-            .expect("failed to generate proof");
-
-        println!("Successfully generated proof!");
-
-        // verify the proof
-        client.verify(&proof, &vk).expect("failed to verify proof");
-        println!("Successfully verified proof!");
-
-        // create & save proof
-        println!("Saving proof.");
-        let proof_data = bincode::serialize(&proof).expect("failed to serialize proof");
-        std::fs::write(args.path.with_extension("sp1.proof"), proof_data)
-            .expect("failed to save SP1 Proof file");
-
-        // save public input
-        println!("Saving public inputs.");
-        std::fs::write(args.path.with_extension("sp1.pub"), proof.public_values)
-            .expect("failed to save SP1 public input");
+    match exec_type {
+        ExecutionType::Execute => {
+            // pass everything at once for execution
+            let mut stdin = SP1Stdin::new();
+            stdin.write(&samples);
+            stdin.write(&query);
+
+            println!("Executing program.");
+            // Execute the program
+
+            let (output, report) = client.execute(PROGRAM_ELF, stdin).run().unwrap();
+            println!("Program executed successfully.");
+
+            // Read the output.
+            let decoded = PublicValuesStruct::abi_decode(output.as_slice(), true).unwrap();
+            let PublicValuesStruct { idx } = decoded;
+            println!("Closest idx: {}", idx);
+
+            let expected_idx = zkvdb_lib::compute_best_sample(&samples, &query);
+            assert_eq!(idx, expected_idx as u32);
+            println!("Values are correct!");
+
+            // Record the number of cycles executed.
+            println!("Number of cycles: {}", report.total_instruction_count());
+        }
+        ExecutionType::Prove => {
+            const CHUNK_SIZE: usize = 3;
+
+            // setup the program for proving.
+            let (pk, vk) = client.setup(PROGRAM_ELF);
+            let (agg_pk, agg_vk) = client.setup(AGGREGATOR_ELF);
+
+            // generate similarity proofs
+            println!("Proving all chunks (chunk size {})", CHUNK_SIZE);
+            let mut proofs = Vec::new();
+            for chunk in samples.chunks(CHUNK_SIZE) {
+                println!("Generating proof for chunk.");
+                let mut stdin = SP1Stdin::new();
+                stdin.write(&chunk);
+                stdin.write(&query);
+
+                let proof = client
+                    .prove(&pk, stdin)
+                    .compressed()
+                    .run()
+                    .expect("failed to generate proof");
+                proofs.push(proof);
+            }
+
+            // aggregate all proofs
+            println!("Aggregating all {} proofs.", proofs.len());
+            let mut stdin = SP1Stdin::new();
+            let inputs: Vec<AggregationInput> = proofs
+                .into_iter()
+                .map(|proof| AggregationInput {
+                    proof,
+                    vk: vk.clone(),
+                })
+                .collect();
+
+            // Write the verification keys.
+            let vkeys = inputs
+                .iter()
+                .map(|input| input.vk.hash_u32())
+                .collect::<Vec<_>>();
+            stdin.write::<Vec<[u32; 8]>>(&vkeys);
+
+            // Write the public values.
+            let public_values = inputs
+                .iter()
+                .map(|input| input.proof.public_values.to_vec())
+                .collect::<Vec<_>>();
+            stdin.write::<Vec<Vec<u8>>>(&public_values);
+
+            // Write the proofs.
+            //
+            // Note: this data will not actually be read by the aggregation program, instead it will be
+            // witnessed by the prover during the recursive aggregation process inside SP1 itself.
+            for input in inputs {
+                let SP1Proof::Compressed(proof) = input.proof.proof else {
+                    panic!("expected compressed proof");
+                };
+                stdin.write_proof(proof, input.vk.vk);
+            }
+
+            println!("Proving the aggregated proof.");
+            let proof = client
+                .prove(&agg_pk, stdin)
+                .compressed()
+                .run()
+                .expect("failed to generate aggregation proof");
+            println!("Successfully generated aggregation proof!");
+
+            // verify the proof
+            client
+                .verify(&proof, &agg_vk)
+                .expect("failed to verify proof");
+            println!("Successfully verified aggregation proof!");
+
+            // create & save proof
+            println!("Saving proof.");
+            let proof_data = bincode::serialize(&proof).expect("failed to serialize proof");
+            std::fs::write(args.path.with_extension("sp1.proof"), proof_data)
+                .expect("failed to save SP1 Proof file");
+
+            // save public input
+            println!("Saving public inputs.");
+            std::fs::write(args.path.with_extension("sp1.pub"), proof.public_values)
+                .expect("failed to save SP1 public input");
+        }
     }
 }