Derive page configs without a per-byte image

crypto/stark/src/grinding.rs

@@ -1,6 +1,6 @@
 use crypto::hash::platform_keccak::PlatformKeccak256 as Keccak256;
 use digest::Digest;
-#[cfg(feature = "parallel")]
+#[cfg(all(feature = "parallel", not(test)))]
 use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 
 const PREFIX: [u8; 8] = [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xed];
@@ -55,10 +55,26 @@ pub fn generate_nonce(seed: &[u8; 32], grinding_factor: u8) -> Option<u64> {
         is_valid_nonce_for_inner_hash(&inner_hash, candidate_nonce, limit)
     });
 
-    #[cfg(feature = "parallel")]
+    #[cfg(all(feature = "parallel", not(test)))]
     return (0..u64::MAX).into_par_iter().find_any(|&candidate_nonce| {
         is_valid_nonce_for_inner_hash(&inner_hash, candidate_nonce, limit)
     });
+
+    // Tests that compare two proofs byte for byte — `retire_lde_proof_is_byte_identical`
+    // is the one — are comparing two grindings as much as two provers, and
+    // `find_any` returns an arbitrary valid nonce of the many that exist. The
+    // nonce reaches the transcript, so a different one moves every challenge and
+    // query index after it: the proofs differ with nothing wrong. The serial
+    // search returns the smallest, which is stable across runs.
+    //
+    // Deliberately not `find_first` in production: which nonce comes back is not
+    // a contract (see `generate_nonce_maybe_gpu`), and this search is the
+    // prover's dominant CPU cost — ordering it would be paying for a property
+    // only the tests want.
+    #[cfg(all(feature = "parallel", test))]
+    return (0..u64::MAX).find(|&candidate_nonce| {
+        is_valid_nonce_for_inner_hash(&inner_hash, candidate_nonce, limit)
+    });
 }
 
 /// Checks if the leftmost 8 bytes of `Hash(inner_hash || candidate_nonce)` are less than `limit`

prover/src/tables/trace_builder.rs

@@ -5417,21 +5417,52 @@ impl Traces {
     /// init data populated. Used by the verifier to reconstruct the ELF
     /// portion of the PAGE table layout.
     pub fn page_configs_from_elf(elf: &Elf) -> Vec<PageConfig> {
-        use std::collections::BTreeSet;
+        use std::collections::BTreeMap;
 
-        let init_page_data = build_init_page_data(&build_initial_image(elf, &[]));
-
-        let page_bases: BTreeSet<u64> = init_page_data.keys().copied().collect();
+        let page_size = page::DEFAULT_PAGE_SIZE;
+        // Written straight into the pages rather than through a per-byte
+        // `HashMap` (`build_initial_image` + `build_init_page_data`): the
+        // verifier runs this, the recursion guest runs the verifier, and one
+        // SipHash insert per ELF byte is most of that guest's startup. Order
+        // and values are what the map produced — ascending bases, last write
+        // wins within a segment run.
+        let mut pages: BTreeMap<u64, Vec<u8>> = BTreeMap::new();
+        for segment in &elf.data {
+            if !segment.base_addr.is_multiple_of(4) {
+                // A word can straddle a page boundary; take the byte path.
+                for (i, &word) in segment.values.iter().enumerate() {
+                    let word_addr = segment.base_addr.wrapping_add(i as u64 * 4);
+                    for byte_offset in 0..4u64 {
+                        let addr = word_addr.wrapping_add(byte_offset);
+                        let page = pages
+                            .entry(page::page_base_for_address(addr))
+                            .or_insert_with(|| vec![0u8; page_size]);
+                        page[page::offset_in_page(addr)] =
+                            ((word >> (byte_offset * 8)) & 0xFF) as u8;
+                    }
+                }
+                continue;
+            }
+            let mut i = 0usize;
+            while i < segment.values.len() {
+                let addr = segment.base_addr.wrapping_add(i as u64 * 4);
+                let offset = page::offset_in_page(addr);
+                let words_in_page = (page_size - offset) / 4;
+                let take = words_in_page.min(segment.values.len() - i);
+                let page = pages
+                    .entry(page::page_base_for_address(addr))
+                    .or_insert_with(|| vec![0u8; page_size]);
+                for (j, &word) in segment.values[i..i + take].iter().enumerate() {
+                    let at = offset + j * 4;
+                    page[at..at + 4].copy_from_slice(&word.to_le_bytes());
+                }
+                i += take;
+            }
+        }
 
-        page_bases
+        pages
             .into_iter()
-            .map(|base| {
-                if let Some(init_data) = init_page_data.get(&base) {
-                    PageConfig::with_data(base, init_data.clone())
-                } else {
-                    PageConfig::zero_init(base)
-                }
-            })
+            .map(|(base, data)| PageConfig::with_data(base, data))
             .collect()
     }
 

prover/src/tests/trace_builder_tests.rs

@@ -2360,3 +2360,43 @@ fn the_end_of_run_tables_match_the_ordinary_build() {
         assert_eq!(flat(a), flat(b), "PAGE {i} differs");
     }
 }
+
+#[test]
+fn page_configs_match_the_byte_image_derivation() {
+    use crate::tables::trace_builder::{Traces, build_init_page_data, build_initial_image};
+    use std::collections::BTreeSet;
+
+    let mut checked = 0;
+    for name in ["fibonacci", "keccak", "allocator", "ethrex"] {
+        let Ok(bytes) = std::fs::read(format!(
+            "{}/executor/program_artifacts/rust/{name}.elf",
+            env!("CARGO_MANIFEST_DIR").trim_end_matches("/prover")
+        )) else {
+            continue;
+        };
+        let elf = executor::elf::Elf::load(&bytes).expect("load elf");
+
+        let init_page_data = build_init_page_data(&build_initial_image(&elf, &[]));
+        let bases: BTreeSet<u64> = init_page_data.keys().copied().collect();
+        let expected: Vec<_> = bases
+            .into_iter()
+            .map(|base| (base, init_page_data.get(&base).cloned().unwrap()))
+            .collect();
+
+        let got = Traces::page_configs_from_elf(&elf);
+        checked += 1;
+        assert_eq!(got.len(), expected.len(), "{name}: page count");
+        for (cfg, (base, data)) in got.iter().zip(expected.iter()) {
+            assert_eq!(cfg.page_base, *base, "{name}: page base");
+            assert_eq!(
+                cfg.init_values.as_deref(),
+                Some(data.as_slice()),
+                "{name}: init values at 0x{base:x}"
+            );
+        }
+    }
+    assert!(
+        checked > 0,
+        "no Rust ELF artifacts found to compare against"
+    );
+}