Add export preprocessed roots dev util.

Author: gilbens-starkware

stwo_cairo_prover/Cargo.lock

@@ -7,6 +7,9 @@ edition.workspace = true
 name = "extract_mem_trace"
 required-features = ["stwo-cairo-adapter/extract-mem-trace"]
 
+[[bin]]
+name = "export_preprocessed_roots"
+
 [dependencies]
 anyhow.workspace = true
 cairo-air.workspace = true
@@ -25,6 +28,7 @@ log.workspace = true
 regex = "1.11.1"
 num-bigint = "0.4.6"
 bincode = "1.3.3"
+itertools.workspace = true
 
 serde_json.workspace = true
 cairo-vm.workspace = true

stwo_cairo_prover/crates/dev_utils/Cargo.toml

@@ -0,0 +1,54 @@
+use clap::{Parser, ValueEnum};
+use dev_utils::preprocessed_roots::{export_preprocessed_roots, HashSelection};
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
+enum HashArg {
+    Blake,
+    Poseidon,
+    Both,
+}
+
+impl From<HashArg> for HashSelection {
+    fn from(value: HashArg) -> Self {
+        match value {
+            HashArg::Blake => HashSelection::Blake,
+            HashArg::Poseidon => HashSelection::Poseidon,
+            HashArg::Both => HashSelection::Both,
+        }
+    }
+}
+
+/// Export preprocessed commitment roots for Cairo preprocessed traces.
+///
+/// Computes the commitment roots for Blake2s and/or Poseidon252 Merkle channels for all
+/// log blowup factors in the range 1..=max.
+///
+/// Performance note: This computation is slow. It is strongly recommended to run with
+/// release optimizations:
+///
+///   cargo run --release --bin export_preprocessed_roots -- [OPTIONS]
+///
+/// Examples:
+///   # Compute both sets of roots up to the default max (2)
+///   cargo run --release --bin export_preprocessed_roots
+///
+///   # Compute only Blake roots up to 4
+///   cargo run --release --bin export_preprocessed_roots -- -m 4 --hash blake
+#[derive(Parser, Debug)]
+#[command(name = "export_preprocessed_roots")]
+struct Args {
+    /// Maximum log blowup factor (inclusive), computes for 1..=max
+    #[arg(short = 'm', long = "max-log-blowup-factor", default_value_t = 2)]
+    max_log_blowup_factor: u32,
+
+    /// Hash function to compute roots for: blake, poseidon, or both
+    #[arg(long = "hash", value_enum, default_value_t = HashArg::Both)]
+    hash: HashArg,
+}
+
+fn main() {
+    let args = Args::parse();
+    export_preprocessed_roots(args.max_log_blowup_factor, args.hash.into());
+}
+
+

stwo_cairo_prover/crates/dev_utils/src/bin/export_preprocessed_roots.rs

@@ -1 +1,2 @@
 pub mod utils;
+pub mod preprocessed_roots;

stwo_cairo_prover/crates/dev_utils/src/lib.rs

@@ -0,0 +1,73 @@
+use cairo_air::PreProcessedTraceVariant;
+use itertools::Itertools;
+use stwo::core::channel::MerkleChannel;
+use stwo::core::vcs::blake2_merkle::Blake2sMerkleChannel;
+use stwo::core::vcs::poseidon252_merkle::Poseidon252MerkleChannel;
+use stwo::core::vcs::MerkleHasher;
+use stwo::prover::backend::BackendForChannel;
+use stwo_cairo_prover::witness::preprocessed_trace::generate_preprocessed_commitment_root;
+
+#[derive(Clone, Copy, Debug)]
+pub enum HashSelection {
+    Blake,
+    Poseidon,
+    Both,
+}
+
+fn get_preprocessed_roots<MC: MerkleChannel>(
+    max_log_blowup_factor: u32,
+    preprocessed_trace: PreProcessedTraceVariant,
+) -> Vec<<MC::H as MerkleHasher>::Hash>
+where
+    stwo::prover::backend::simd::SimdBackend: BackendForChannel<MC>,
+{
+    (1..=max_log_blowup_factor)
+        .map(|i| generate_preprocessed_commitment_root::<MC>(i, preprocessed_trace))
+        .collect_vec()
+}
+
+pub fn export_preprocessed_roots(max_log_blowup_factor: u32, selection: HashSelection) {
+    match selection {
+        HashSelection::Blake => {
+            let roots = get_preprocessed_roots::<Blake2sMerkleChannel>(
+                max_log_blowup_factor,
+                PreProcessedTraceVariant::Canonical,
+            );
+            roots.iter().enumerate().for_each(|(i, root)| {
+                let root_bytes = root.0;
+                let u32s_hex = root_bytes
+                    .as_slice()
+                    .chunks_exact(4)
+                    .map(|bytes| {
+                        let arr: [u8; 4] = bytes.try_into().expect("chunk size");
+                        format!("{:#010x}", u32::from_le_bytes(arr))
+                    })
+                    .collect_vec()
+                    .join(", ");
+                println!("log_blowup_factor: {}, blake root: [{}]", i + 1, u32s_hex);
+            });
+        }
+        HashSelection::Poseidon => {
+            println!("Starting Poseidon roots");
+            get_preprocessed_roots::<Poseidon252MerkleChannel>(
+                max_log_blowup_factor,
+                PreProcessedTraceVariant::CanonicalWithoutPedersen,
+            )
+            .into_iter()
+            .enumerate()
+            .for_each(|(i, root)| {
+                println!(
+                    "log_blowup_factor: {}, poseidon root: [{:#010x}]",
+                    i + 1,
+                    root
+                );
+            });
+        }
+        HashSelection::Both => {
+            export_preprocessed_roots(max_log_blowup_factor, HashSelection::Blake);
+            export_preprocessed_roots(max_log_blowup_factor, HashSelection::Poseidon);
+        }
+    }
+}
+
+

stwo_cairo_prover/crates/dev_utils/src/preprocessed_roots.rs

@@ -8,25 +8,24 @@ use stwo::prover::poly::circle::PolyOps;
 use stwo::prover::CommitmentTreeProver;
 
 /// Generates the root of the preprocessed trace commitment tree for a given `log_blowup_factor`.
-// TODO(Shahars): remove allow.
-#[allow(unused)]
 pub fn generate_preprocessed_commitment_root<MC: MerkleChannel>(
     log_blowup_factor: u32,
     preprocessed_trace: PreProcessedTraceVariant,
 ) -> <<MC as MerkleChannel>::H as MerkleHasher>::Hash
 where
     SimdBackend: BackendForChannel<MC>,
 {
+    println!("Generating preprocessed trace root for log blowup factor: {}", log_blowup_factor);
     let preprocessed_trace = preprocessed_trace.to_preprocessed_trace();
-
+    println!("Computing twiddles");
     // Precompute twiddles for the commitment scheme.
     let max_log_size = preprocessed_trace.log_sizes().into_iter().max().unwrap();
     let twiddles = SimdBackend::precompute_twiddles(
         CanonicCoset::new(max_log_size + log_blowup_factor)
             .circle_domain()
             .half_coset,
     );
-
+    println!("Generating commitment tree");
     // Generate the commitment tree.
     let polys = SimdBackend::interpolate_columns(preprocessed_trace.gen_trace(), &twiddles);
     let commitment_scheme = CommitmentTreeProver::<SimdBackend, MC>::new(

stwo_cairo_prover/crates/prover/src/witness/preprocessed_trace.rs

@@ -1,15 +1,11 @@
 use std::sync::atomic::{AtomicU32, Ordering};
 
 use cairo_air::air::CairoClaim;
-use cairo_air::PreProcessedTraceVariant;
-use itertools::Itertools;
 use num_traits::{One, Zero};
 use stwo::core::channel::MerkleChannel;
 use stwo::core::fields::m31::M31;
 use stwo::core::pcs::{TreeSubspan, TreeVec};
-use stwo::core::vcs::blake2_merkle::Blake2sMerkleChannel;
-use stwo::core::vcs::poseidon252_merkle::Poseidon252MerkleChannel;
-use stwo::core::vcs::MerkleHasher;
+ 
 use stwo::prover::backend::simd::column::BaseColumn;
 use stwo::prover::backend::simd::conversion::Pack;
 use stwo::prover::backend::simd::m31::{PackedM31, N_LANES};
@@ -19,7 +15,7 @@ use stwo::prover::poly::BitReversedOrder;
 use stwo_cairo_common::preprocessed_columns::preprocessed_trace::PreProcessedTrace;
 use stwo_constraint_framework::PREPROCESSED_TRACE_IDX;
 
-use crate::witness::preprocessed_trace::generate_preprocessed_commitment_root;
+ 
 
 pub fn pack_values<T: Pack>(values: &[T]) -> Vec<T::SimdType> {
     values
@@ -130,55 +126,7 @@ pub fn witness_trace_cells(claim: &CairoClaim, pp_trace: &PreProcessedTrace) ->
     tree_trace_cells(log_sizes)
 }
 
-fn get_preprocessed_roots<MC: MerkleChannel>(
-    max_log_blowup_factor: u32,
-    preprocessed_trace: PreProcessedTraceVariant,
-) -> Vec<<MC::H as MerkleHasher>::Hash>
-where
-    stwo::prover::backend::simd::SimdBackend: BackendForChannel<MC>,
-{
-    (1..=max_log_blowup_factor)
-        .map(|i| generate_preprocessed_commitment_root::<MC>(i, preprocessed_trace))
-        .collect_vec()
-}
-
-/// Exports the preprocessed roots for both Blake2s and Poseidon252 channels.
-/// Note: This function is very slow and is intended for generating the preprocessed roots when
-/// needed.
-pub fn export_preprocessed_roots() {
-    let max_log_blowup_factor = 5;
-
-    // Blake2s roots.
-    let blake_roots = get_preprocessed_roots::<Blake2sMerkleChannel>(
-        max_log_blowup_factor,
-        PreProcessedTraceVariant::Canonical,
-    );
-    blake_roots.iter().enumerate().for_each(|(i, root)| {
-        let root_bytes = root.0;
-        let u32s_hex = root_bytes
-            .array_chunks::<4>()
-            .map(|&bytes| format!("{:#010x}", u32::from_le_bytes(bytes)))
-            .collect_vec()
-            .join(", ");
-
-        println!("log_blowup_factor: {}, blake root: [{}]", i + 1, u32s_hex);
-    });
-
-    // Poseidon252 roots.
-    get_preprocessed_roots::<Poseidon252MerkleChannel>(
-        max_log_blowup_factor,
-        PreProcessedTraceVariant::CanonicalWithoutPedersen,
-    )
-    .into_iter()
-    .enumerate()
-    .for_each(|(i, root)| {
-        println!(
-            "log_blowup_factor: {}, poseidon root: [{:#010x}]",
-            i + 1,
-            root
-        );
-    });
-}
+// export_preprocessed_roots moved to `dev_utils` binary.
 
 #[cfg(test)]
 mod tests {