perf: build PAGE trace from dense memory store

prover/src/paged_mem.rs

@@ -103,6 +103,17 @@ impl<T: Copy> PagedMem<T> {
         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

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<GoldilocksField, GoldilocksExtension> {
+    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
 // =========================================================================

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<I: ImageSource>(
     // Derive ALL page bases from memory_state (includes ELF + runtime pages)
     let page_bases: BTreeSet<u64> = 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<I: ImageSource>(
     // Derive ALL page bases from memory_state (includes ELF + runtime pages)
     let page_bases: BTreeSet<u64> = 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<I: ImageSource>(
             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);
     }

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<F>: PartialEq` only
+/// holds without `disk-spill`, and the whole `TraceTable` can't be compared
+/// because its `aux_table: Table<E>` 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}"
+        );
+    }
+}