Merge pull request #228 from zkMIPS/split-large-seg

feat: add support for splitting large segment into small segments

Author: xander42280

emulator/src/memory.rs

@@ -322,6 +322,25 @@ impl Memory {
         self.count = count;
     }
 
+    pub fn init_memory<'a>(&mut self, addr: u32, v: u32) {
+        let page_index = addr >> PAGE_ADDR_SIZE;
+        let page_addr = (addr as usize) & PAGE_ADDR_MASK;
+        let cached_page = match self.page_lookup(page_index) {
+            None => {
+                // allocate the page if we have not already
+                // Golang may mmap relatively large ranges, but we only allocate just in time.
+                self.alloc_page(page_index)
+            }
+            Some(cached_page) => {
+                // self.invalidate(addr);
+                cached_page
+            }
+        };
+
+        let mut cached_page = cached_page.borrow_mut();
+        cached_page.data[page_addr..page_addr + 4].copy_from_slice(&v.to_le_bytes());
+    }
+
     pub fn set_memory_range<'a>(
         &mut self,
         mut addr: u32,

emulator/src/state.rs

@@ -8,6 +8,8 @@ use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 use std::fmt::{Display, Formatter};
 use std::fs;
+use std::fs::File;
+use std::io::BufReader;
 use std::io::{stderr, stdout, Read, Write};
 use std::path::Path;
 
@@ -138,6 +140,72 @@ impl State {
         })
     }
 
+    pub fn load_seg(seg_path: &str) -> (Box<Self>, u64) {
+        let reader = BufReader::new(File::open(seg_path).unwrap());
+        let segment: Segment = serde_json::from_reader(reader).unwrap();
+        let mut s = Box::new(Self {
+            memory: Box::new(Memory::new()),
+            registers: Default::default(),
+            pc: segment.pc,
+            next_pc: 0,
+            hi: 0,
+            lo: 0,
+            heap: 0,
+            local_user: 0,
+            step: 0,
+            total_step: 0,
+            cycle: 0,
+            total_cycle: 0,
+            brk: 0,
+            input_stream: segment.input_stream,
+            input_stream_ptr: segment.input_stream_ptr,
+            public_values_stream: segment.public_values_stream,
+            public_values_stream_ptr: segment.public_values_stream_ptr,
+            exited: false,
+            exit_code: 0,
+            dump_info: false,
+        });
+
+        let image = segment.mem_image;
+
+        for i in 0..32 {
+            let data = image.get(&(REGISTERS_START + (i << 2) as u32)).unwrap();
+            s.registers[i] = data.to_be();
+        }
+
+        s.lo = image
+            .get(&(REGISTERS_START + (32 << 2) as u32))
+            .unwrap()
+            .to_be();
+        s.hi = image
+            .get(&(REGISTERS_START + (33 << 2) as u32))
+            .unwrap()
+            .to_be();
+        s.heap = image
+            .get(&(REGISTERS_START + (34 << 2) as u32))
+            .unwrap()
+            .to_be();
+        s.next_pc = image
+            .get(&(REGISTERS_START + (36 << 2) as u32))
+            .unwrap()
+            .to_be();
+
+        s.brk = image
+            .get(&(REGISTERS_START + (37 << 2) as u32))
+            .unwrap()
+            .to_be();
+
+        s.local_user = image
+            .get(&(REGISTERS_START + (38 << 2) as u32))
+            .unwrap()
+            .to_be();
+
+        for (addr, data) in image {
+            s.memory.init_memory(addr, data);
+        }
+        (s, segment.step)
+    }
+
     pub fn load_elf(f: &elf::ElfBytes<AnyEndian>) -> Box<Self> {
         let mut s = Box::new(Self {
             memory: Box::new(Memory::new()),

emulator/src/utils.rs

@@ -32,10 +32,10 @@ pub fn split_prog_into_segs(
     instrumented_state.split_segment(false, seg_path, new_writer);
     let new_writer = |name: &str| -> Option<std::fs::File> { File::create(name).ok() };
     loop {
+        let cycles = instrumented_state.step();
         if instrumented_state.state.exited {
             break;
         }
-        let cycles = instrumented_state.step();
         if cycles > (seg_size as isize - 1) as u64 {
             instrumented_state.split_segment(true, seg_path, new_writer);
         }
@@ -54,3 +54,51 @@ pub fn split_prog_into_segs(
         instrumented_state.state,
     )
 }
+
+pub fn load_segment(seg_file: &str) -> (Box<State>, u64) {
+    State::load_seg(seg_file)
+}
+
+pub fn split_seg_into_segs(
+    seg_file: &str,
+    seg_path: &str,
+    block_path: &str,
+    seg_size: usize,
+) -> (usize, usize, Box<State>) {
+    let (state, final_step) = load_segment(seg_file);
+    let mut instrumented_state = InstrumentedState::new(state, block_path.to_string());
+    log::debug!("start pc: {:X} {}", instrumented_state.state.pc, final_step);
+    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 {
+        let cycles = instrumented_state.step();
+        if instrumented_state.state.total_step + instrumented_state.state.step == final_step {
+            break;
+        }
+        if cycles > (seg_size as isize - 1) as u64 {
+            instrumented_state.split_segment(true, seg_path, new_writer);
+            log::debug!(
+                "Split at {} : {} into {}",
+                instrumented_state.state.total_step,
+                instrumented_state.state.total_cycle,
+                instrumented_state.pre_segment_id
+            );
+        }
+    }
+    instrumented_state.split_segment(true, seg_path, new_writer);
+    log::info!(
+        "Split done {} : {} into {}",
+        instrumented_state.state.total_step,
+        instrumented_state.state.total_cycle,
+        instrumented_state.pre_segment_id
+    );
+
+    instrumented_state.dump_memory();
+    (
+        instrumented_state.state.total_step as usize,
+        instrumented_state.pre_segment_id as usize,
+        instrumented_state.state,
+    )
+}

prover/examples/Cargo.toml

@@ -7,7 +7,8 @@ members = [
   "sha2-syscall/host",
   "keccak/host",
   "split-seg",
-  "prove-seg"
+  "prove-seg",
+  "prove-large-seg"
 ]
 resolver = "2"
 

prover/examples/README.md

@@ -33,6 +33,14 @@ BASEDIR=../../../emulator/test-vectors RUST_LOG=info BLOCK_NO=13284491 SEG_FILE_
     cargo run --release
 ```
 
+* Generate proof for specific large segment (Split large segment into small segments)
+
+```
+cd ../prove-large-seg
+BASEDIR=../../../emulator/test-vectors RUST_LOG=info SEG_FILE=/tmp/output/8 BLOCK_NO=13284491 SEG_OUTPUT=/tmp/output1 SEG_SIZE=16384 cargo run --release
+
+```
+
 ### Prove Go sdk code
 The SDK provide Read and Commit interface to read input and commit output.
 Take sha2-go for example:

prover/examples/prove-large-seg/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "prove-large-seg"
+version = { workspace = true }
+edition = { workspace = true }
+publish = false
+
+[dependencies]
+zkm-prover = { workspace = true }
+zkm-emulator = { workspace = true }
+zkm-utils = { path = "../utils" }
+
+log = { version = "0.4.14", default-features = false }
+serde = { version = "1.0.144", features = ["derive"] }
+serde_json = "1.0"
+byteorder = "1.5.0"
+hex = "0.4"
+env_logger = "0.11.5"
+anyhow = "1.0.75"
+
+[build-dependencies]
+zkm-build = { workspace = true }

prover/examples/prove-large-seg/src/main.rs

@@ -0,0 +1,23 @@
+use std::env;
+use zkm_emulator::utils::{get_block_path, split_seg_into_segs};
+use zkm_utils::utils;
+
+fn prove_large_segment() {
+    let basedir = env::var("BASEDIR").unwrap_or("/tmp/cannon".to_string());
+    let block = env::var("BLOCK_NO").unwrap_or("".to_string());
+    let file = env::var("BLOCK_FILE").unwrap_or(String::from(""));
+    let seg_file = env::var("SEG_FILE").expect("big segment file is missing");
+    let seg_dir = env::var("SEG_OUTPUT").expect("segment output dir is missing");
+    let seg_size = env::var("SEG_SIZE").unwrap_or("1024".to_string());
+    let seg_size = seg_size.parse::<_>().unwrap_or(1usize);
+
+    let block_path = get_block_path(&basedir, &block, "");
+    let (_, seg_num, _) = split_seg_into_segs(&seg_file, &seg_dir, &block_path, seg_size);
+
+    let _ = utils::prove_segments(&seg_dir, &basedir, &block, &file, seg_num, 0, vec![]);
+}
+
+fn main() {
+    env_logger::try_init().unwrap_or_default();
+    prove_large_segment();
+}

prover/examples/prove-seg/src/main.rs

@@ -1,5 +1,5 @@
 use std::env;
-use zkm_prover::utils;
+use zkm_utils::utils;
 
 fn prove_segments() {
     let basedir = env::var("BASEDIR").unwrap_or("/tmp/cannon".to_string());