Skip to content

Commit

Permalink
feat: add test-batches-with-each-chunk-num-prove to test batches wi…
Browse files Browse the repository at this point in the history
…th `1 - 15` chunks (#278)

* Add `recursive-prove` to test over 1 - 15 chunks in a batch.

* Rename to `test_batches_with_iter_chunk_num_prove_verify`.

* Rename to `test_batches_with_each_chunk_num_prove_verify`.

* Refactor and fix to use chunk-proofs for further iter batch-proving.
  • Loading branch information
silathdiir authored Oct 3, 2023
1 parent 453251b commit 8e3b21a
Show file tree
Hide file tree
Showing 4 changed files with 469 additions and 37 deletions.
3 changes: 3 additions & 0 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,9 @@ test-agg-prove:
test-batch-prove:
@cargo test --features prove_verify --release test_batch_prove_verify

test-batches-with-each-chunk-num-prove:
@cargo test --features prove_verify --release test_batches_with_each_chunk_num_prove_verify

test-ccc:
@cargo test --release test_capacity_checker

Expand Down
44 changes: 29 additions & 15 deletions integration/tests/aggregation_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,21 +15,10 @@ fn test_agg_prove_verify() {
log::info!("Initialized ENV and created output-dir {output_dir}");

let trace_paths = vec!["./tests/extra_traces/new.json".to_string()];

let chunk_hashes_proofs = gen_chunk_hashes_and_proofs(&output_dir, &trace_paths);
log::info!("Generated chunk hashes and proofs");

env::set_var("AGG_VK_FILENAME", "vk_batch_agg.vkey");
env::set_var("CHUNK_PROTOCOL_FILENAME", "chunk_chunk_0.protocol");
let mut agg_prover = Prover::from_dirs(PARAMS_DIR, &output_dir);
log::info!("Constructed aggregation prover");

// Load or generate aggregation snark (layer-3).
let layer3_snark = agg_prover
.load_or_gen_last_agg_snark("agg", chunk_hashes_proofs, Some(&output_dir))
.unwrap();

gen_and_verify_batch_proofs(&mut agg_prover, layer3_snark, &output_dir);
let mut batch_prover = new_batch_prover(&output_dir);
prove_and_verify_batch(&output_dir, &mut batch_prover, chunk_hashes_proofs);
}

fn gen_chunk_hashes_and_proofs(
Expand All @@ -47,7 +36,7 @@ fn gen_chunk_hashes_and_proofs(
})
.collect();

chunk_traces
let chunk_hashes_proofs = chunk_traces
.into_iter()
.enumerate()
.map(|(i, chunk_trace)| {
Expand All @@ -60,5 +49,30 @@ fn gen_chunk_hashes_and_proofs(

(chunk_hash, proof)
})
.collect()
.collect();

log::info!("Generated chunk hashes and proofs");
chunk_hashes_proofs
}

fn new_batch_prover(assets_dir: &str) -> Prover {
env::set_var("AGG_VK_FILENAME", "vk_batch_agg.vkey");
env::set_var("CHUNK_PROTOCOL_FILENAME", "chunk_chunk_0.protocol");
let prover = Prover::from_dirs(PARAMS_DIR, &assets_dir);
log::info!("Constructed batch prover");

prover
}

fn prove_and_verify_batch(
output_dir: &str,
batch_prover: &mut Prover,
chunk_hashes_proofs: Vec<(ChunkHash, ChunkProof)>,
) {
// Load or generate aggregation snark (layer-3).
let layer3_snark = batch_prover
.load_or_gen_last_agg_snark("agg", chunk_hashes_proofs, Some(&output_dir))
.unwrap();

gen_and_verify_batch_proofs(batch_prover, layer3_snark, &output_dir);
}
93 changes: 71 additions & 22 deletions integration/tests/batch_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,51 +3,100 @@ use prover::{
aggregator::Prover, proof::from_json_file, utils::init_env_and_log, ChunkHash, ChunkProof,
};
use serde_derive::{Deserialize, Serialize};
use std::env;
use std::{env, fs, path::PathBuf};

#[cfg(feature = "prove_verify")]
#[test]
fn test_batch_prove_verify() {
let output_dir = init_env_and_log("batch_tests");
log::info!("Initialized ENV and created output-dir {output_dir}");

let chunk_hashes_proofs = load_chunk_hashes_and_proofs("tests/test_data", "1");
log::info!("Loaded chunk-hashes and chunk-proofs");
let chunk_hashes_proofs = load_chunk_hashes_and_proofs("tests/test_data", "1", &output_dir);
let mut batch_prover = new_batch_prover(&output_dir);
prove_and_verify_batch(&output_dir, &mut batch_prover, chunk_hashes_proofs);
}

chunk_hashes_proofs
.first()
.unwrap()
.1
.dump(&output_dir, "0")
.unwrap();
#[cfg(feature = "prove_verify")]
#[test]
fn test_batches_with_each_chunk_num_prove_verify() {
let output_dir = init_env_and_log("batches_with_each_chunk_num_tests");
log::info!("Initialized ENV and created output-dir {output_dir}");

env::set_var("AGG_VK_FILENAME", "vk_batch_agg.vkey");
env::set_var("CHUNK_PROTOCOL_FILENAME", "chunk_chunk_0.protocol");
let mut agg_prover = Prover::from_dirs(PARAMS_DIR, &output_dir);
log::info!("Constructed aggregation prover");
let chunk_hashes_proofs = load_chunk_hashes_and_proofs("tests/test_data", "2", &output_dir);
let mut batch_prover = new_batch_prover(&output_dir);

// Load or generate aggregation snark (layer-3).
let layer3_snark = agg_prover
.load_or_gen_last_agg_snark("agg", chunk_hashes_proofs, Some(&output_dir))
.unwrap();
// Iterate over chunk proofs to test with 1 - 15 chunks (in a batch).
for i in 0..chunk_hashes_proofs.len() {
let mut output_dir = PathBuf::from(&output_dir);
output_dir.push(format!("batch_{}", i + 1));
fs::create_dir_all(&output_dir).unwrap();

gen_and_verify_batch_proofs(&mut agg_prover, layer3_snark, &output_dir);
prove_and_verify_batch(
&output_dir.to_string_lossy().to_string(),
&mut batch_prover,
chunk_hashes_proofs[..=i].to_vec(),
);
}
}

#[derive(Debug, Deserialize, Serialize)]
struct BatchTaskDetail {
chunk_infos: Vec<ChunkHash>,
chunk_proofs: Vec<ChunkProof>,
}

fn load_chunk_hashes_and_proofs(dir: &str, filename: &str) -> Vec<(ChunkHash, ChunkProof)> {
fn load_chunk_hashes_and_proofs(
dir: &str,
filename: &str,
output_dir: &str,
) -> Vec<(ChunkHash, ChunkProof)> {
let batch_task_detail: BatchTaskDetail = from_json_file(dir, filename).unwrap();
let chunk_hashes = batch_task_detail.chunk_infos;
let chunk_proofs = batch_task_detail.chunk_proofs;

chunk_hashes[..]
let chunk_hashes_proofs: Vec<_> = chunk_hashes[..]
.to_vec()
.into_iter()
.zip(chunk_proofs[..].to_vec().into_iter())
.collect()
.collect();

// Dump chunk-procotol for further batch-proving.
chunk_hashes_proofs
.first()
.unwrap()
.1
.dump(&output_dir, "0")
.unwrap();

log::info!(
"Loaded chunk-hashes and chunk-proofs: total = {}",
chunk_hashes_proofs.len()
);
chunk_hashes_proofs
}

fn new_batch_prover(assets_dir: &str) -> Prover {
env::set_var("AGG_VK_FILENAME", "vk_batch_agg.vkey");
env::set_var("CHUNK_PROTOCOL_FILENAME", "chunk_chunk_0.protocol");
let prover = Prover::from_dirs(PARAMS_DIR, &assets_dir);
log::info!("Constructed batch prover");

prover
}

fn prove_and_verify_batch(
output_dir: &str,
batch_prover: &mut Prover,
chunk_hashes_proofs: Vec<(ChunkHash, ChunkProof)>,
) {
let chunk_num = chunk_hashes_proofs.len();
log::info!("Prove batch BEGIN: chunk_num = {chunk_num}");

// Load or generate aggregation snark (layer-3).
let layer3_snark = batch_prover
.load_or_gen_last_agg_snark("agg", chunk_hashes_proofs, Some(&output_dir))
.unwrap();

gen_and_verify_batch_proofs(batch_prover, layer3_snark, &output_dir);

log::info!("Prove batch END: chunk_num = {chunk_num}");
}
Loading

0 comments on commit 8e3b21a

Please sign in to comment.