feat: unify use of split_elf_into_segs

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)

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::<AnyEndian>::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::<AnyEndian>::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::<AnyEndian>::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<std::fs::File> { None };
- instrumented_state.split_segment(false, OUTPUT, new_writer);
- let mut segment_step = SEGMENT_STEPS;
- let new_writer = |name: &str| -> Option<std::fs::File> { 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::<AnyEndian>::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<std::fs::File> { None };
- instrumented_state.split_segment(false, OUTPUT, new_writer);
- let mut segment_step = SEGMENT_STEPS;
- let new_writer = |name: &str| -> Option<std::fs::File> { 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);
  }
 }

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<State> {
+ let data = fs::read(elf_path).expect("could not read file");
+ let file =
+ ElfBytes::<AnyEndian>::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<State>,
+ 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<std::fs::File> { None };
+ instrumented_state.split_segment(false, seg_path, new_writer);
+ let mut segment_step: usize = seg_size;
+ let new_writer = |name: &str| -> Option<std::fs::File> { 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
+}

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 

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,96 +25,40 @@ use zkm_prover::verifier::verify_proof;
 
 const DEGREE_BITS_RANGE: [Range<usize>; 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::<AnyEndian>::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<std::fs::File> { None };
- instrumented_state.split_segment(false, &seg_path, new_writer);
- let mut segment_step: usize = seg_size;
- let new_writer = |name: &str| -> Option<std::fs::File> { 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() {
  // 1. split ELF into segs
  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::<AnyEndian>::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> =
- InstrumentedState::new(state, "".to_string());
- std::fs::create_dir_all(&seg_path).unwrap();
- let new_writer = |_: &str| -> Option<std::fs::File> { None };
- instrumented_state.split_segment(false, &seg_path, new_writer);
- let new_writer = |name: &str| -> Option<std::fs::File> { 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 = <C as GenericConfig<D>>::F;
@@ -176,7 +120,7 @@ fn main() {
  let args: Vec<String> = 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(),