Avoid duplicate L2G trace work in continuations

prover/src/continuation.rs

@@ -39,7 +39,6 @@ use std::collections::HashMap;
 use crypto::fiat_shamir::default_transcript::DefaultTranscript;
 use executor::elf::Elf;
 use executor::vm::execution::Executor;
-use executor::vm::logs::Log;
 use math::field::element::FieldElement;
 use stark::config::Commitment;
 use stark::lookup::{AirWithBuses, AuxiliaryTraceBuildData, NullBoundaryConstraintBuilder};
@@ -50,14 +49,12 @@ use stark::trace::TraceTable;
 use stark::traits::AIR;
 use stark::verifier::{IsStarkVerifier, Verifier};
 
-use crate::paged_mem::PagedMem;
 use crate::statement::{StatementKind, absorb_continuation_global_statement, absorb_statement};
 use crate::tables::local_to_global::{self, CellBoundary};
 use crate::tables::page::{self, PageConfig};
 use crate::tables::register;
 use crate::tables::trace_builder::{
     Traces, build_init_page_data, build_initial_image, build_initial_image_paged,
-    epoch_touched_cells,
 };
 use crate::tables::types::{GoldilocksExtension, GoldilocksField};
 use crate::tables::{MaxRowsConfig, global_memory};
@@ -220,12 +217,9 @@ fn global_memory_configs(
         .collect()
 }
 
-/// Per-epoch starting state: the memory image and register image the epoch begins from.
-/// `image` is borrowed from the persistent cross-epoch image (init = previous fini), so
-/// it is not re-snapshotted or cloned per epoch.
+/// Per-epoch register state and label.
 struct EpochStart<'a> {
-    image: &'a PagedMem<u8>,
-    register_init: HashMap<u64, u32>,
+    register_init: &'a HashMap<u64, u32>,
     is_first: bool,
     /// This epoch's 1-based table label (the `fini_epoch` constant).
     label: u64,
@@ -331,25 +325,11 @@ fn prove_epoch(
     elf: &Elf,
     elf_bytes: &[u8],
     start: &EpochStart,
-    logs: &[Log],
+    mut traces: Traces,
     is_final: bool,
     boundary: &[CellBoundary],
-    private_inputs: &[u8],
     opts: &ProofOptions,
 ) -> Result<EpochProof, Error> {
-    let mut traces = Traces::from_image_and_logs(
-        elf,
-        start.image,
-        &start.register_init,
-        logs,
-        &MaxRowsConfig::default(),
-        private_inputs,
-        is_final,
-        true,
-        #[cfg(feature = "disk-spill")]
-        stark::storage_mode::StorageMode::Ram,
-    )?;
-
     // Use the cross-epoch boundary so this epoch's L2G table is identical to the
     // one the global proof commits (the commitment binding compares their roots).
     traces.local_to_global = local_to_global::generate_local_to_global_trace(boundary);
@@ -386,7 +366,7 @@ fn prove_epoch(
         opts,
         &[],
         &table_counts,
-        &start.register_init,
+        start.register_init,
         &reg_fini,
         start.is_first,
         is_final,
@@ -689,25 +669,27 @@ pub fn prove_continuation(
         );
 
         let label = local_to_global::epoch_label(index);
-        let touched = epoch_touched_cells(&elf, &image, &register_init, &logs)?;
-        let boundary = local_to_global::epoch_boundary(&mut provenance, label, &touched);
+        let traces = Traces::from_image_and_logs(
+            &elf,
+            &image,
+            &register_init,
+            &logs,
+            &MaxRowsConfig::default(),
+            private_inputs,
+            is_final,
+            true,
+            #[cfg(feature = "disk-spill")]
+            stark::storage_mode::StorageMode::Ram,
+        )?;
+        let boundary =
+            local_to_global::epoch_boundary(&mut provenance, label, &traces.touched_memory_cells);
 
         let start = EpochStart {
-            image: &image,
-            register_init,
+            register_init: &register_init,
             is_first: index == 0,
             label,
         };
-        let epoch = prove_epoch(
-            &elf,
-            elf_bytes,
-            &start,
-            &logs,
-            is_final,
-            &boundary,
-            private_inputs,
-            opts,
-        )?;
+        let epoch = prove_epoch(&elf, elf_bytes, &start, traces, is_final, &boundary, opts)?;
         prev_fini = Some(epoch.reg_fini.clone());
 
         // Carry the image forward: this epoch's fini is the next epoch's init.

prover/src/tables/trace_builder.rs

@@ -1929,14 +1929,18 @@ pub fn epoch_touched_cells<I: ImageSource>(
     let mut register_state = RegisterState::from_init_map(register_init);
     let _ = collect_ops_from_cpu(&cpu_ops, &mut memory_state, &mut register_state);
 
+    Ok(touched_cells_from_memory_state(&memory_state))
+}
+
+fn touched_cells_from_memory_state(memory_state: &MemoryState) -> local_to_global::EpochTouches {
     let mut touched: Vec<(u64, u64, u64)> = memory_state
         .cells
         .iter()
         .filter(|(_, cell)| cell.1 > 0)
         .map(|(addr, cell)| (addr, cell.0 as u64, cell.1))
         .collect();
     touched.sort_by_key(|&(addr, _, _)| addr);
-    Ok(touched)
+    touched
 }
 
 /// Bucket an initial-memory image into per-page byte arrays for PAGE init columns.
@@ -2576,9 +2580,13 @@ pub struct Traces {
     /// MEMW_R register-only fast-path traces (split into chunks of max_rows::MEMW_R)
     pub memw_registers: Vec<TraceTable<GoldilocksField, GoldilocksExtension>>,
     /// Local-to-global boundary table for continuation epochs. Empty unless the
-    /// epoch is built with `l2g_memory_bookend` (then it bookends the Memory bus
-    /// for touched RAM bytes; see [`local_to_global`]).
+    /// continuation driver fills it with the boundary derived from
+    /// `touched_memory_cells`.
     pub local_to_global: TraceTable<GoldilocksField, GoldilocksExtension>,
+    /// Touched cells observed while replaying this epoch's logs, each as
+    /// `(address, end_value, end_timestamp)`. Populated only for continuation
+    /// epochs that use the L2G memory bookend.
+    pub touched_memory_cells: local_to_global::EpochTouches,
     // Auxiliary ALU / memory / CPU32 dispatch chips (split into chunks of their max_rows)
     pub eqs: Vec<TraceTable<GoldilocksField, GoldilocksExtension>>,
     pub bytewises: Vec<TraceTable<GoldilocksField, GoldilocksExtension>>,
@@ -3156,26 +3164,15 @@ fn build_traces<I: ImageSource + Sync>(
         }
     }
 
-    // Local-to-global boundary table. Built only for continuation epochs that use
-    // L2G as the Memory-bus bookend; it claims each touched RAM byte's epoch-start
-    // value (init, at ts 0) and epoch-end value/timestamp (fini), derived from the
-    // SAME `memory_state.cells` (timestamp > 0) that PAGE just excluded.
-    let local_to_global = match (l2g_memory_bookend, initial_image) {
-        (true, Some(image)) => {
-            let mut touched: Vec<(u64, u64, u64)> = memory_state
-                .cells
-                .iter()
-                .filter(|(_, cell)| cell.1 > 0)
-                .map(|(addr, (value, ts))| (addr, value as u64, ts))
-                .collect();
-            touched.sort_by_key(|&(addr, _, _)| addr);
-            let initial_memory: HashMap<u64, u64> =
-                image.image_iter().map(|(a, v)| (a, v as u64)).collect();
-            let boundaries = local_to_global::epoch_boundaries(&initial_memory, &[touched]);
-            local_to_global::generate_local_to_global_trace(&boundaries[0])
-        }
-        _ => local_to_global::generate_local_to_global_trace(&[]),
+    // Continuation callers derive the real cross-epoch boundary from this set and
+    // install its L2G trace after provenance is applied. Avoid building a
+    // throwaway genesis-only L2G trace here.
+    let touched_memory_cells = if l2g_memory_bookend {
+        touched_cells_from_memory_state(memory_state)
+    } else {
+        Vec::new()
     };
+    let local_to_global = local_to_global::generate_local_to_global_trace(&[]);
 
     Ok(Traces {
         cpus,
@@ -3203,6 +3200,7 @@ fn build_traces<I: ImageSource + Sync>(
         ecdas: ecdas_trace,
         memw_registers,
         local_to_global,
+        touched_memory_cells,
         eqs,
         bytewises,
         stores,
@@ -3506,6 +3504,7 @@ impl Traces {
             page_configs: _,
             public_output_bytes: _,
             local_to_global: _,
+            touched_memory_cells: _,
         } = self;
 
         let mut total: u64 = 0;
@@ -3638,6 +3637,7 @@ impl Traces {
             page_configs: _,
             public_output_bytes: _,
             local_to_global: _,
+            touched_memory_cells: _,
         } = self;
 
         let mut total: u64 = 0;

prover/src/tests/prove_elfs_tests.rs

@@ -3349,6 +3349,7 @@ fn test_epoch_memory_bus_with_l2g_bookend() {
     use crate::tables::register;
     use crate::tables::trace_builder::build_initial_image;
     use crate::test_utils::asm_elf_bytes;
+    use std::collections::HashMap;
 
     let _ = env_logger::builder().is_test(true).try_init();
     let elf_bytes = asm_elf_bytes("all_loadstore_32");
@@ -3379,6 +3380,10 @@ fn test_epoch_memory_bus_with_l2g_bookend() {
         stark::storage_mode::StorageMode::Ram,
     )
     .unwrap();
+    let initial_memory: HashMap<u64, u64> = image.iter().map(|(&a, &v)| (a, v as u64)).collect();
+    let boundaries =
+        local_to_global::epoch_boundaries(&initial_memory, &[traces.touched_memory_cells.clone()]);
+    traces.local_to_global = local_to_global::generate_local_to_global_trace(&boundaries[0]);
 
     let proof_options = ProofOptions::default_test_options();
     let table_counts = traces.table_counts();