adapter better print

starkware-libs · Jan 8, 2025 · ad04851 · ad04851
1 parent 4a9e8f6
commit ad04851
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/stwo_cairo_prover/crates/adapted_prover/src/main.rs b/stwo_cairo_prover/crates/adapted_prover/src/main.rs
@@ -67,8 +67,10 @@ fn run(args: impl Iterator<Item = String>) -> Result<(), Error> {
         .display_components(args.display_components)
         .build();
 
-    let casm_states_by_opcode_count = &vm_output.state_transitions.casm_states_by_opcode.counts();
-    log::info!("Casm states by opcode count: {casm_states_by_opcode_count:?}");
+    log::info!(
+        "Casm states by opcode:\n {},",
+        vm_output.state_transitions.casm_states_by_opcode
+    );
 
     // TODO(Ohad): Propagate hash from CLI args.
     let proof = prove_cairo::<Blake2sMerkleChannel>(vm_output, prover_config)?;

diff --git a/stwo_cairo_prover/crates/prover/src/input/state_transitions.rs b/stwo_cairo_prover/crates/prover/src/input/state_transitions.rs
@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::fmt::Display;
 
 use prover_types::cpu::CasmState;
 use stwo_prover::core::fields::m31::M31;
@@ -158,6 +159,18 @@ impl CasmStatesByOpcode {
     }
 }
 
+impl Display for CasmStatesByOpcode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let counts = self.counts();
+        let total_steps = counts.iter().map(|(_, count)| count).sum::<usize>();
+        writeln!(f, "Total steps: {total_steps}")?;
+        for (name, count) in &counts {
+            writeln!(f, "{name}: {count}")?;
+        }
+        Ok(())
+    }
+}
+
 impl From<TraceEntry> for CasmState {
     fn from(entry: TraceEntry) -> Self {
         Self {