From bd1ee8e0bd614185a5f9c62160603a6f923ddd25 Mon Sep 17 00:00:00 2001 From: Angell Li Date: Fri, 12 Jul 2024 20:19:47 +0800 Subject: [PATCH] feat: unify use of split_elf_into_segs --- emulator/src/state.rs | 1 - emulator/src/tests.rs | 85 ++++-------------------------- emulator/src/utils.rs | 47 +++++++++++++++++ prover/examples/README.md | 15 ++++-- prover/examples/zkmips.rs | 89 ++++++-------------------------- prover/tests/simple_recursion.rs | 57 +++----------------- 6 files changed, 89 insertions(+), 205 deletions(-) diff --git a/emulator/src/state.rs b/emulator/src/state.rs index 8a7efe88..ee38596e 100644 --- a/emulator/src/state.rs +++ b/emulator/src/state.rs @@ -15,7 +15,6 @@ pub const FD_STDOUT: u32 = 1; pub const FD_STDERR: u32 = 2; pub const MIPS_EBADF: u32 = 9; -pub const SEGMENT_STEPS: usize = 1024; pub const REGISTERS_START: u32 = 0x81020400u32; // image_id = keccak(page_hash_root || end_pc) diff --git a/emulator/src/tests.rs b/emulator/src/tests.rs index fad9af24..728e3dc5 100644 --- a/emulator/src/tests.rs +++ b/emulator/src/tests.rs @@ -1,16 +1,13 @@ #[allow(clippy::module_inception)] #[cfg(test)] mod tests { - use elf::{endian::AnyEndian, ElfBytes}; - use std::fs::File; use std::{ fs, path::{Path, PathBuf}, }; - use crate::state::SEGMENT_STEPS; use crate::state::{InstrumentedState, State}; - use crate::utils::get_block_path; + use crate::utils::{get_block_path, load_elf_with_patch, split_prog_into_segs, SEGMENT_STEPS}; const END_ADDR: u32 = 0xa7ef00d0; const OUTPUT: &str = "/tmp/segment"; @@ -54,14 +51,7 @@ mod tests { #[test] fn test_execute_hello() { - let path = PathBuf::from("test-vectors/hello"); - let data = fs::read(path).expect("could not read file"); - let file = - ElfBytes::::minimal_parse(data.as_slice()).expect("opening elf file failed"); - let mut state = State::load_elf(&file); - - state.patch_elf(&file); - state.patch_stack(vec!["aab", "ccd"]); + let state = load_elf_with_patch("test-vectors/hello", vec!["aab", "ccd"]); let mut instrumented_state = InstrumentedState::new(state, String::from("")); @@ -75,13 +65,7 @@ mod tests { #[test] fn test_execute_rust_fib() { - let path = PathBuf::from("test-vectors/rust_fib"); - let data = fs::read(path).expect("could not read file"); - let file = - ElfBytes::::minimal_parse(data.as_slice()).expect("opening elf file failed"); - let mut state = State::load_elf(&file); - state.patch_elf(&file); - state.patch_stack(vec![]); + let state = load_elf_with_patch("test-vectors/rust_fib", vec![]); let mut instrumented_state = InstrumentedState::new(state, String::from("")); log::debug!("begin execute\n"); @@ -94,70 +78,19 @@ mod tests { } #[test] - #[ignore] + #[ignore = "Two slow"] fn test_execute_minigeth() { - let path = PathBuf::from("test-vectors/minigeth"); - let data = fs::read(path).expect("could not read file"); - let file = - ElfBytes::::minimal_parse(data.as_slice()).expect("opening elf file failed"); - let mut state = State::load_elf(&file); - - state.patch_elf(&file); - state.patch_stack(vec![]); + let mut state = load_elf_with_patch("test-vectors/minigeth", vec![]); - let block_path = get_block_path("../test-vectors", "13284491", ""); + let block_path = get_block_path("test-vectors", "13284491", ""); state.load_input(&block_path); - let mut instrumented_state = InstrumentedState::new(state, block_path); - std::fs::create_dir_all(OUTPUT).unwrap(); - let new_writer = |_: &str| -> Option { None }; - instrumented_state.split_segment(false, OUTPUT, new_writer); - let mut segment_step = SEGMENT_STEPS; - let new_writer = |name: &str| -> Option { File::create(name).ok() }; - loop { - if instrumented_state.state.exited { - break; - } - instrumented_state.step(); - segment_step -= 1; - if segment_step == 0 { - segment_step = SEGMENT_STEPS; - instrumented_state.split_segment(true, OUTPUT, new_writer); - } - } - - instrumented_state.split_segment(true, OUTPUT, new_writer); + let _ = split_prog_into_segs(state, OUTPUT, &block_path, SEGMENT_STEPS); } #[test] fn test_execute_split_hello() { - let path = PathBuf::from("test-vectors/hello"); - let data = fs::read(path).expect("could not read file"); - let file = - ElfBytes::::minimal_parse(data.as_slice()).expect("opening elf file failed"); - let mut state = State::load_elf(&file); - - state.patch_elf(&file); - state.patch_stack(vec![]); - - let mut instrumented_state = InstrumentedState::new(state, String::from("")); - std::fs::create_dir_all(OUTPUT).unwrap(); - let new_writer = |_: &str| -> Option { None }; - instrumented_state.split_segment(false, OUTPUT, new_writer); - let mut segment_step = SEGMENT_STEPS; - let new_writer = |name: &str| -> Option { File::create(name).ok() }; - loop { - if instrumented_state.state.exited { - break; - } - instrumented_state.step(); - segment_step -= 1; - if segment_step == 0 { - segment_step = SEGMENT_STEPS; - instrumented_state.split_segment(true, OUTPUT, new_writer); - } - } - - instrumented_state.split_segment(true, OUTPUT, new_writer); + let state = load_elf_with_patch("test-vectors/hello", vec![]); + let _ = split_prog_into_segs(state, OUTPUT, "", SEGMENT_STEPS); } } diff --git a/emulator/src/utils.rs b/emulator/src/utils.rs index 70cb5ce7..f86262fb 100644 --- a/emulator/src/utils.rs +++ b/emulator/src/utils.rs @@ -1,4 +1,51 @@ +use crate::state::{InstrumentedState, State}; +use elf::{endian::AnyEndian, ElfBytes}; +use std::fs; +use std::fs::File; + +pub const SEGMENT_STEPS: usize = 1024; + /// From the minigeth's rule, the `block` starts with `0_` pub fn get_block_path(basedir: &str, block: &str, file: &str) -> String { format!("{basedir}/0_{block}/{file}") } + +pub fn load_elf_with_patch(elf_path: &str, args: Vec<&str>) -> Box { + let data = fs::read(elf_path).expect("could not read file"); + let file = + ElfBytes::::minimal_parse(data.as_slice()).expect("opening elf file failed"); + let mut state = State::load_elf(&file); + state.patch_elf(&file); + state.patch_stack(args); + state +} + +pub fn split_prog_into_segs( + state: Box, + seg_path: &str, + block_path: &str, + seg_size: usize, +) -> usize { + let mut instrumented_state = InstrumentedState::new(state, block_path.to_string()); + std::fs::create_dir_all(seg_path).unwrap(); + let new_writer = |_: &str| -> Option { None }; + instrumented_state.split_segment(false, seg_path, new_writer); + let mut segment_step: usize = seg_size; + let new_writer = |name: &str| -> Option { File::create(name).ok() }; + loop { + if instrumented_state.state.exited { + break; + } + instrumented_state.step(); + segment_step -= 1; + if segment_step == 0 { + segment_step = seg_size; + instrumented_state.split_segment(true, seg_path, new_writer); + } + } + instrumented_state.split_segment(true, seg_path, new_writer); + log::info!("Split done {}", instrumented_state.state.step); + + instrumented_state.dump_memory(); + instrumented_state.state.step as usize +} diff --git a/prover/examples/README.md b/prover/examples/README.md index d5c48046..8eb6a235 100644 --- a/prover/examples/README.md +++ b/prover/examples/README.md @@ -13,28 +13,33 @@ GOOS=linux GOARCH=mips GOMIPS=softfloat go build hello.go * Split the ELF hello into segments. Note that the flag `BLOCK_NO` is only necessary for minigeth. ``` -BASEDIR=test-vectors RUST_LOG=info ELF_PATH=test-vectors/hello BLOCK_NO=13284491 SEG_OUTPUT=/tmp/output SEG_SIZE=1024 ARGS="" \ +BASEDIR=./emulator/test-vectors RUST_LOG=info ELF_PATH=./emulator/test-vectors/minigeth BLOCK_NO=13284491 SEG_OUTPUT=/tmp/output SEG_SIZE=1024 ARGS="" \ cargo run --release --example zkmips split + +OR + +RUST_LOG=info ELF_PATH=/emulator/test-vectors/hello SEG_OUTPUT=/tmp/output SEG_SIZE=1024 ARGS="" \ + cargo run --release --example zkmips split_without_preimage ``` * Generate proof for each segment ``` -BASEDIR=test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE="/tmp/output/0" SEG_SIZE=1024 \ +BASEDIR=./emulator/test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE="/tmp/output/0" SEG_SIZE=1024 \ cargo run --release --example zkmips prove ``` * Aggregate proof ``` -BASEDIR=test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE="/tmp/output/0" SEG_FILE2="/tmp/output/1" SEG_SIZE=1024 \ +BASEDIR=./emulator/test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE="/tmp/output/0" SEG_FILE2="/tmp/output/1" SEG_SIZE=1024 \ cargo run --release --example zkmips aggregate_proof ``` * Aggregate proof all ``` -BASEDIR=test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE_DIR="/tmp/output" SEG_FILE_NUM=299 SEG_SIZE=1024 \ +BASEDIR=./emulator/test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE_DIR="/tmp/output" SEG_FILE_NUM=299 SEG_SIZE=1024 \ cargo run --release --example zkmips aggregate_proof_all ``` @@ -60,7 +65,7 @@ rustflags = ["--cfg", 'target_os="zkvm"',"-C", "target-feature=+crt-static", "-C ``` cd examples/sha2 cargo build --target=mips-unknown-linux-musl -cd ../../ +cd ././ ``` * Run the host program diff --git a/prover/examples/zkmips.rs b/prover/examples/zkmips.rs index 3c6858ad..00729112 100644 --- a/prover/examples/zkmips.rs +++ b/prover/examples/zkmips.rs @@ -1,6 +1,5 @@ -use elf::{endian::AnyEndian, ElfBytes}; use std::env; -use std::fs::{self, File}; +use std::fs::File; use std::io::BufReader; use std::ops::Range; use std::time::Duration; @@ -12,8 +11,9 @@ use plonky2x::backend::circuit::Groth16WrapperParameters; use plonky2x::backend::wrapper::wrap::WrappedCircuit; use plonky2x::frontend::builder::CircuitBuilder as WrapperBuilder; use plonky2x::prelude::DefaultParameters; -use zkm_emulator::state::{InstrumentedState, State, SEGMENT_STEPS}; -use zkm_emulator::utils::get_block_path; +use zkm_emulator::utils::{ + get_block_path, load_elf_with_patch, split_prog_into_segs, SEGMENT_STEPS, +}; use zkm_prover::all_stark::AllStark; use zkm_prover::config::StarkConfig; use zkm_prover::cpu::kernel::assembler::segment_kernel; @@ -25,54 +25,23 @@ use zkm_prover::verifier::verify_proof; const DEGREE_BITS_RANGE: [Range; 6] = [10..21, 12..22, 12..21, 8..21, 6..21, 13..23]; -fn split_elf_into_segs() { +fn split_segs(load_preimage: bool) { // 1. split ELF into segs let basedir = env::var("BASEDIR").unwrap_or("/tmp/cannon".to_string()); let elf_path = env::var("ELF_PATH").expect("ELF file is missing"); - let block_no = env::var("BLOCK_NO"); + let block_no = env::var("BLOCK_NO").unwrap_or("".to_string()); let seg_path = env::var("SEG_OUTPUT").expect("Segment output path is missing"); let seg_size = env::var("SEG_SIZE").unwrap_or(format!("{SEGMENT_STEPS}")); let seg_size = seg_size.parse::<_>().unwrap_or(SEGMENT_STEPS); let args = env::var("ARGS").unwrap_or("".to_string()); let args = args.split_whitespace().collect(); - let data = fs::read(elf_path).expect("could not read file"); - let file = - ElfBytes::::minimal_parse(data.as_slice()).expect("opening elf file failed"); - let mut state = State::load_elf(&file); - state.patch_elf(&file); - state.patch_stack(args); - - let block_path = match block_no { - Ok(no) => { - let block_path = get_block_path(&basedir, &no, ""); - state.load_input(&block_path); - block_path - } - _ => "".to_string(), - }; - - let mut instrumented_state = InstrumentedState::new(state, block_path); - std::fs::create_dir_all(&seg_path).unwrap(); - let new_writer = |_: &str| -> Option { None }; - instrumented_state.split_segment(false, &seg_path, new_writer); - let mut segment_step: usize = seg_size; - let new_writer = |name: &str| -> Option { File::create(name).ok() }; - loop { - if instrumented_state.state.exited { - break; - } - instrumented_state.step(); - segment_step -= 1; - if segment_step == 0 { - segment_step = seg_size; - instrumented_state.split_segment(true, &seg_path, new_writer); - } + let mut state = load_elf_with_patch(&elf_path, args); + let block_path = get_block_path(&basedir, &block_no, ""); + if load_preimage { + state.load_input(&block_path); } - instrumented_state.split_segment(true, &seg_path, new_writer); - log::info!("Split done {}", instrumented_state.state.step); - - instrumented_state.dump_memory(); + let _ = split_prog_into_segs(state, &seg_path, &block_path, seg_size); } fn prove_sha2_bench() { @@ -80,41 +49,16 @@ fn prove_sha2_bench() { let elf_path = env::var("ELF_PATH").expect("ELF file is missing"); let seg_path = env::var("SEG_OUTPUT").expect("Segment output path is missing"); - let data = fs::read(elf_path).expect("could not read file"); - let file = - ElfBytes::::minimal_parse(data.as_slice()).expect("opening elf file failed"); - let mut state = State::load_elf(&file); - state.patch_elf(&file); - state.patch_stack(vec![]); - + let mut state = load_elf_with_patch(&elf_path, vec![]); // load input let input = [5u8; 32]; state.add_input_stream(&input.to_vec()); - let mut instrumented_state: Box = - InstrumentedState::new(state, "".to_string()); - std::fs::create_dir_all(&seg_path).unwrap(); - let new_writer = |_: &str| -> Option { None }; - instrumented_state.split_segment(false, &seg_path, new_writer); - let new_writer = |name: &str| -> Option { File::create(name).ok() }; - loop { - if instrumented_state.state.exited { - break; - } - instrumented_state.step(); - } - instrumented_state.split_segment(true, &seg_path, new_writer); - log::info!("Split done {}", instrumented_state.state.step); + let total_steps = split_prog_into_segs(state, &seg_path, "", 0); let seg_file = format!("{seg_path}/{}", 0); let seg_reader = BufReader::new(File::open(seg_file).unwrap()); - let kernel = segment_kernel( - "", - "", - "", - seg_reader, - instrumented_state.state.step as usize, - ); + let kernel = segment_kernel("", "", "", seg_reader, total_steps); const D: usize = 2; type C = PoseidonGoldilocksConfig; type F = >::F; @@ -176,7 +120,7 @@ fn main() { let args: Vec = env::args().collect(); let helper = || { log::info!( - "Help: {} split | prove | aggregate_proof | aggregate_proof_all | prove_groth16 | bench", + "Help: {} split | split_without_preimage | prove | aggregate_proof | aggregate_proof_all | prove_groth16 | bench", args[0] ); std::process::exit(-1); @@ -185,7 +129,8 @@ fn main() { helper(); } match args[1].as_str() { - "split" => split_elf_into_segs(), + "split" => split_segs(true), + "split_without_preimage" => split_segs(false), "prove" => prove_single_seg(), "aggregate_proof" => aggregate_proof().unwrap(), "aggregate_proof_all" => aggregate_proof_all().unwrap(), diff --git a/prover/tests/simple_recursion.rs b/prover/tests/simple_recursion.rs index a3b2f8df..52450748 100644 --- a/prover/tests/simple_recursion.rs +++ b/prover/tests/simple_recursion.rs @@ -1,6 +1,5 @@ #![allow(clippy::upper_case_acronyms)] -use elf::{endian::AnyEndian, ElfBytes}; -use std::fs::{self, File}; +use std::fs::File; use std::time::Duration; use plonky2::field::goldilocks_field::GoldilocksField; @@ -10,8 +9,7 @@ use plonky2::plonk::circuit_builder::CircuitBuilder; use plonky2::plonk::circuit_data::CircuitConfig; use plonky2::plonk::config::PoseidonGoldilocksConfig; use std::io::BufReader; -use zkm_emulator::state::{InstrumentedState, State}; -use zkm_emulator::utils::get_block_path; +use zkm_emulator::utils::{load_elf_with_patch, split_prog_into_segs}; use zkm_prover::all_stark::AllStark; use zkm_prover::config::StarkConfig; use zkm_prover::cpu::kernel::assembler::segment_kernel; @@ -24,48 +22,6 @@ type F = GoldilocksField; const D: usize = 2; type C = PoseidonGoldilocksConfig; -fn split_elf_into_segs( - basedir: &str, - elf_path: &str, - block_no: &str, - seg_path: &str, - seg_size: usize, -) { - // 1. split ELF into segs - let data = fs::read(elf_path) - .map_err(|_| panic!("could not read file, {}", elf_path)) - .unwrap(); - let file = - ElfBytes::::minimal_parse(data.as_slice()).expect("opening elf file failed"); - let mut state = State::load_elf(&file); - state.patch_elf(&file); - state.patch_stack(vec![]); - - let block_path = get_block_path(basedir, block_no, ""); - state.load_input(&block_path); - - let mut instrumented_state = InstrumentedState::new(state, block_path); - std::fs::create_dir_all(seg_path).unwrap(); - let new_writer = |_: &str| -> Option { None }; - instrumented_state.split_segment(false, seg_path, new_writer); - let mut segment_step: usize = seg_size; - let new_writer = |name: &str| -> Option { File::create(name).ok() }; - loop { - if instrumented_state.state.exited { - break; - } - instrumented_state.step(); - segment_step -= 1; - if segment_step == 0 { - segment_step = seg_size; - instrumented_state.split_segment(true, seg_path, new_writer); - } - } - - instrumented_state.split_segment(true, seg_path, new_writer); - log::info!("Split done"); -} - // Tests proving two transactions, one of which with logs, and aggregating them. #[test] fn test_mips_with_aggreg_fibo() -> anyhow::Result<()> { @@ -146,13 +102,12 @@ fn test_mips_with_aggreg() -> anyhow::Result<()> { use plonky2x::frontend::builder::CircuitBuilder as WrapperBuilder; use plonky2x::prelude::DefaultParameters; - let basedir = "../emualtor/test-vectors"; let seg_output = "/tmp/mips_output"; let elf_path = "../emualtor/test-vectors/hello"; - let block_no = "13284491"; let seg_size = 65536; - split_elf_into_segs(basedir, elf_path, block_no, seg_output, seg_size); + let state = load_elf_with_patch(elf_path, vec![]); + let _ = split_prog_into_segs(state, seg_output, "", seg_size); type InnerParameters = DefaultParameters; type OuterParameters = Groth16WrapperParameters; @@ -170,7 +125,7 @@ fn test_mips_with_aggreg() -> anyhow::Result<()> { let seg_file = format!("{}/0", seg_output); let seg_reader = BufReader::new(File::open(seg_file)?); - let input_first = segment_kernel(basedir, block_no, "", seg_reader, seg_size); + let input_first = segment_kernel("", "", "", seg_reader, seg_size); let mut timing = TimingTree::new("prove root first", log::Level::Info); let (root_proof_first, first_public_values) = all_circuits.prove_root(&all_stark, &input_first, &config, &mut timing)?; @@ -180,7 +135,7 @@ fn test_mips_with_aggreg() -> anyhow::Result<()> { let seg_file = format!("{}/1", seg_output); let seg_reader = BufReader::new(File::open(seg_file)?); - let input = segment_kernel(basedir, block_no, "", seg_reader, seg_size); + let input = segment_kernel("", "", "", seg_reader, seg_size); let mut timing = TimingTree::new("prove root second", log::Level::Info); let (root_proof, public_values) = all_circuits.prove_root(&all_stark, &input, &config, &mut timing)?;