diff --git a/prover/src/paged_mem.rs b/prover/src/paged_mem.rs index 196d077bf..61bdedbe6 100644 --- a/prover/src/paged_mem.rs +++ b/prover/src/paged_mem.rs @@ -103,6 +103,17 @@ impl PagedMem { self.pages.iter().map(|(b, _)| *b) } + /// The dense per-offset data slice for the page at `base` (page-aligned), or + /// `None` if that page holds no `set` cell. One indexed read per offset (no + /// hashing) — lets PAGE trace generation read each offset's final value + /// straight from the store instead of a sparse `FinalStateMap` lookup. + pub fn page_data(&self, base: u64) -> Option<&[T]> { + match self.pages.binary_search_by_key(&base, |(b, _)| *b) { + Ok(i) => Some(&self.pages[i].1.data), + Err(_) => None, + } + } + /// Number of cells that were explicitly `set`. pub fn len(&self) -> usize { self.pages diff --git a/prover/src/tables/page.rs b/prover/src/tables/page.rs index 18ce6b52b..0613ed865 100644 --- a/prover/src/tables/page.rs +++ b/prover/src/tables/page.rs @@ -284,6 +284,78 @@ pub fn generate_page_trace( trace } +/// The final `(value, timestamp)` a PAGE offset contributes, from its init byte +/// and the dense-store cell `(value, timestamp)`. An untouched/image offset +/// (`timestamp == 0`, where value already equals init) or a runtime write dropped +/// by `exclude_touched` collapses to `(init, 0)`; otherwise it keeps the written +/// `(value, timestamp)`. +/// +/// Single source of truth for the PAGE table's FINI/TIMESTAMP columns +/// ([`generate_page_trace_from_dense`]) AND the ARE_BYTES bitwise multiplicities +/// (`collect_bitwise_from_page`), so the two cannot drift and the AreBytes bus +/// stays balanced. +pub fn page_final(init: u8, value: u8, timestamp: u64, exclude_touched: bool) -> (u8, u64) { + if timestamp == 0 || (exclude_touched && timestamp > 0) { + (init, 0) + } else { + (value, timestamp) + } +} + +/// Like [`generate_page_trace`] but reads each offset's final `(value, timestamp)` +/// straight from the dense per-page memory store ([`crate::paged_mem::PagedMem::page_data`]) +/// — one indexed read per offset, no hashing. Equivalent to looking each byte up in a +/// `FinalStateMap` built from the same cells, but avoids the sparse map and its +/// `page_size` (mostly-miss) lookups per page — the dominant cost of PAGE generation. +/// +/// `final_page` is `None` for a page with no runtime cells (every offset falls back to +/// init, ts 0). `exclude_touched` drops runtime-written cells (ts > 0) so PAGE +/// self-cancels them — the continuation-epoch case where the L2G table owns them. +/// +/// `ts == 0` marks an offset that is either untouched or an initial-image byte; either +/// way its final value equals its init value, so we emit `(init, 0)` (matching the +/// `FinalStateMap`-miss branch of [`generate_page_trace`]). +pub fn generate_page_trace_from_dense( + config: &PageConfig, + final_page: Option<&[(u8, u64)]>, + exclude_touched: bool, +) -> TraceTable { + let page_size = DEFAULT_PAGE_SIZE; + assert!( + config.page_base.is_multiple_of(page_size as u64), + "Page base must be page-aligned" + ); + if let Some(page) = final_page { + debug_assert_eq!(page.len(), page_size, "dense page slice must span the page"); + } + + let num_rows = page_size; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), + cols::NUM_COLUMNS, + 1, + ); + let table = &mut trace.main_table; + + for offset in 0..page_size { + table.set_u64(offset, cols::OFFSET, offset as u64); + + let init_value = config + .init_values + .as_ref() + .and_then(|v| v.get(offset).copied()) + .unwrap_or(0); + table.set_byte(offset, cols::INIT, init_value); + + let (value, timestamp) = final_page.map_or((0u8, 0u64), |p| p[offset]); + let (fini_value, fini_ts) = page_final(init_value, value, timestamp, exclude_touched); + table.set_byte(offset, cols::FINI, fini_value); + table.set_dword_wl(offset, cols::TIMESTAMP_LO, fini_ts); + } + + trace +} + // ========================================================================= // Preprocessed commitment // ========================================================================= diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 5c6b3085e..862bf22ab 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -61,7 +61,7 @@ use super::memw::{self, MemwOperation}; use super::memw_aligned; use super::memw_register::{self, RegRow}; use super::mul::{self, MulOperation}; -use super::page::{self, FinalByteState, FinalStateMap, PageConfig}; +use super::page::{self, PageConfig}; use super::register::{self, FinalRegisterStateMap, FinalRegisterWordState}; use super::shift::{self, ShiftOperation}; use super::store; @@ -2125,30 +2125,23 @@ fn collect_bitwise_from_page( // Derive ALL page bases from memory_state (includes ELF + runtime pages) let page_bases: BTreeSet = memory_state.cells.page_bases().collect(); - // Build final state map from memory_state, matching `generate_page_tables`: - // when `exclude_touched`, touched cells (timestamp > 0) are dropped so PAGE - // emits `fini == init` for them, and the ARE_BYTES multiplicities here must - // agree (otherwise the AreBytes bus would not balance). - let final_state: FinalStateMap = memory_state - .cells - .iter() - .filter(|(_, cell)| !exclude_touched || cell.1 == 0) - .map(|(addr, (value, timestamp))| (addr, FinalByteState { timestamp, value })) - .collect(); - - // For each page and each byte, add ARE_BYTES lookups for init and fini + // Read each offset's final `(value, timestamp)` straight from the dense + // per-page store instead of a sparse `FinalStateMap` lookup per offset + // (mostly-miss) — same optimization as `generate_page_tables`. `page_final` + // computes `(init, fini)` identically to `generate_page_trace_from_dense`, so + // the ARE_BYTES multiplicities match the PAGE table's FINI column and the + // AreBytes bus stays balanced. for &page_base in &page_bases { let init_data = init_page_data.get(&page_base); + let final_page = memory_state.cells.page_data(page_base); for offset in 0..page_size { - let addr = page_base + offset as u64; - - // Get init value (from ELF or 0). `.get().unwrap_or(0)` to match the - // relaxed `init_values` contract: a shorter vec reads as trailing zeros. + // Init value (from ELF or 0). `.get().unwrap_or(0)` matches the relaxed + // `init_values` contract: a shorter vec reads as trailing zeros. let init = init_data.map_or(0u8, |data| data.get(offset).copied().unwrap_or(0)); - // Get fini value (from final_state or init if never accessed) - let fini = final_state.get(&addr).map_or(init, |state| state.value); + let (value, timestamp) = final_page.map_or((0u8, 0u64), |p| p[offset]); + let (fini, _) = page::page_final(init, value, timestamp, exclude_touched); // C1+C2: ARE_BYTES[init, fini] — batched range check for both bytes. // Bumped straight into the histogram: this loop visits every byte of @@ -2614,16 +2607,13 @@ fn generate_page_tables( // Derive ALL page bases from memory_state (includes ELF + runtime pages) let page_bases: BTreeSet = memory_state.cells.page_bases().collect(); - // Build final state map from memory_state. When `exclude_touched` (continuation - // epoch with L2G bookend), drop touched cells (timestamp > 0) so PAGE self- - // cancels them (init == fini, ts == 0) and the local-to-global table owns their - // Memory-bus init/fini instead. - let final_state: FinalStateMap = memory_state - .cells - .iter() - .filter(|(_, cell)| !exclude_touched || cell.1 == 0) - .map(|(addr, (value, timestamp))| (addr, FinalByteState { timestamp, value })) - .collect(); + // The per-page final `(value, timestamp)` is read straight from the dense + // `memory_state.cells` store (one indexed read per offset) rather than routing + // through a sparse `FinalStateMap` whose per-page (mostly-miss) lookups dominated + // PAGE generation. `exclude_touched` (continuation epoch with L2G bookend) drops + // runtime-written cells (ts > 0) so PAGE self-cancels them (init == fini, ts == 0) + // and the local-to-global table owns their Memory-bus init/fini instead — applied + // per offset inside `generate_page_trace_from_dense`. // Generate PAGE tables and configs let mut pages = Vec::new(); @@ -2643,7 +2633,8 @@ fn generate_page_tables( PageConfig::zero_init(page_base) }; - let trace = page::generate_page_trace(&config, &final_state); + let final_page = memory_state.cells.page_data(page_base); + let trace = page::generate_page_trace_from_dense(&config, final_page, exclude_touched); pages.push(trace); page_configs.push(config); } diff --git a/prover/src/tests/page_tests.rs b/prover/src/tests/page_tests.rs index f1bcc6933..15edf84ae 100644 --- a/prover/src/tests/page_tests.rs +++ b/prover/src/tests/page_tests.rs @@ -249,3 +249,51 @@ fn page_commitments_empty_list_matches_none() { "empty page_commitments slice must behave like None — every page falls through to recompute", ); } + +/// Differential test: the dense PAGE generator must produce a byte-identical +/// trace to the sparse `FinalStateMap` generator it replaces, for both the +/// monolithic (`exclude_touched = false`) and continuation-epoch +/// (`exclude_touched = true`) cases. Locks the PR's "exact same PAGE trace" +/// claim directly, instead of relying only on full prove+verify integration. +/// +/// Compares `main_table` (PAGE is a main-only table): `Table: PartialEq` only +/// holds without `disk-spill`, and the whole `TraceTable` can't be compared +/// because its `aux_table: Table` param lacks `PartialEq`. +#[cfg(not(feature = "disk-spill"))] +#[test] +fn generate_page_trace_dense_matches_sparse() { + use crate::paged_mem::PagedMem; + + let page_base = 0u64; + let init_bytes = vec![0x01u8, 0x02, 0x03, 0x04]; + let config = PageConfig::with_data(page_base, init_bytes.clone()); + + // Mirror production (`MemoryState::from_image` + replay): seed the initial + // image as (byte, ts=0), then apply runtime writes at ts > 0 (production + // uses ts = i*4 + 4, always >= 4). + let mut cells: PagedMem<(u8, u64)> = PagedMem::new((0, 0)); + for (off, &b) in init_bytes.iter().enumerate() { + cells.set(page_base + off as u64, (b, 0)); // image seed (ts 0) + } + cells.set(page_base, (0xFF, 100)); // overwrite an init byte at runtime + cells.set(page_base + 10, (0x77, 48)); // write a previously-untouched cell + + for exclude_touched in [false, true] { + // Old path: sparse FinalStateMap, filtered exactly as trace_builder did. + let final_state: FinalStateMap = cells + .iter() + .filter(|(_, cell)| !exclude_touched || cell.1 == 0) + .map(|(addr, (value, timestamp))| (addr, FinalByteState { timestamp, value })) + .collect(); + let sparse = generate_page_trace(&config, &final_state); + + // New path: dense per-page slice. + let dense = + generate_page_trace_from_dense(&config, cells.page_data(page_base), exclude_touched); + + assert_eq!( + sparse.main_table, dense.main_table, + "mismatch with exclude_touched = {exclude_touched}" + ); + } +}