fmt for cairo

stwo_cairo_prover/crates/adapted_prover/src/main.rs

@@ -28,6 +28,8 @@ struct Args {
     proof_path: PathBuf,
     #[structopt(long = "debug_lookup")]
     debug_lookup: bool,
+    #[structopt(long = "display_components")]
+    display_components: bool,
 }
 
 #[derive(Debug, Error)]
@@ -69,7 +71,7 @@ fn run(args: impl Iterator<Item = String>) -> Result<CairoProof<Blake2sMerkleHas
     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:?}");
 
-    let proof = prove_cairo(vm_output, args.debug_lookup)?;
+    let proof = prove_cairo(vm_output, args.debug_lookup, args.display_components)?;
 
     // TODO(yuval): This is just some serialization for the sake of serialization. Find the right
     // way to serialize the proof.

stwo_cairo_prover/crates/prover/src/cairo_air/air.rs

@@ -17,6 +17,7 @@ use stwo_prover::core::vcs::blake2_merkle::Blake2sMerkleChannel;
 use stwo_prover::core::vcs::ops::MerkleHasher;
 use tracing::{span, Level};
 
+use super::debug_tools::indented_component_display;
 use super::opcodes_air::{
     OpcodeClaim, OpcodeComponents, OpcodeInteractionClaim, OpcodesClaimGenerator,
     OpcodesInteractionClaimGenerator,
@@ -548,3 +549,51 @@ impl CairoComponents {
             .collect()
     }
 }
+
+impl std::fmt::Display for CairoComponents {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        writeln!(f, "CairoComponents")?;
+        writeln!(f, "Opcodes: {}", self.opcodes)?;
+        writeln!(
+            f,
+            "VerifyInstruction: {}",
+            indented_component_display(&self.verify_instruction)
+        )?;
+        writeln!(
+            f,
+            "MemoryAddressToId: {}",
+            indented_component_display(&self.memory_address_to_id)
+        )?;
+        writeln!(
+            f,
+            "MemoryIdToValue: {}",
+            indented_component_display(&self.memory_id_to_value.0)
+        )?;
+        writeln!(
+            f,
+            "SmallMemoryIdToValue: {}",
+            indented_component_display(&self.memory_id_to_value.1)
+        )?;
+        writeln!(
+            f,
+            "RangeCheck19: {}",
+            indented_component_display(&self.range_check_19)
+        )?;
+        writeln!(
+            f,
+            "RangeCheck9_9: {}",
+            indented_component_display(&self.range_check9_9)
+        )?;
+        writeln!(
+            f,
+            "RangeCheck7_2_5: {}",
+            indented_component_display(&self.range_check7_2_5)
+        )?;
+        writeln!(
+            f,
+            "RangeCheck4_3: {}",
+            indented_component_display(&self.range_check4_3)
+        )?;
+        Ok(())
+    }
+}

stwo_cairo_prover/crates/prover/src/cairo_air/debug.rs → stwo_cairo_prover/crates/prover/src/cairo_air/debug_tools.rs

@@ -3,7 +3,9 @@ use num_traits::One;
 use stwo_prover::constraint_framework::relation_tracker::{
     RelationTrackerComponent, RelationTrackerEntry,
 };
-use stwo_prover::constraint_framework::TraceLocationAllocator;
+use stwo_prover::constraint_framework::{
+    FrameworkComponent, FrameworkEval, TraceLocationAllocator,
+};
 use stwo_prover::core::backend::simd::SimdBackend;
 use stwo_prover::core::fields::m31::M31;
 use stwo_prover::core::pcs::CommitmentSchemeProver;
@@ -583,3 +585,20 @@ pub fn track_cairo_relations(
 
     entries
 }
+
+pub(super) fn indented_component_display<E: FrameworkEval>(
+    component: &FrameworkComponent<E>,
+) -> String {
+    let component_display = &format!("\n{}", component);
+    component_display
+        .lines()
+        .map(|line| format!("\t{}", line))
+        .join("\n")
+}
+
+pub(super) fn display_components<E: FrameworkEval>(components: &[FrameworkComponent<E>]) -> String {
+    components
+        .iter()
+        .map(|component| indented_component_display(component))
+        .join("\n")
+}

stwo_cairo_prover/crates/prover/src/cairo_air/mod.rs

@@ -1,9 +1,9 @@
 pub mod air;
-mod debug;
+mod debug_tools;
 pub mod opcodes_air;
 
 use air::{lookup_sum, CairoClaimGenerator, CairoComponents, CairoInteractionElements, CairoProof};
-use debug::track_cairo_relations;
+use debug_tools::track_cairo_relations;
 use num_traits::Zero;
 use stwo_prover::constraint_framework::preprocessed_columns::gen_is_first;
 use stwo_prover::constraint_framework::relation_tracker::RelationSummary;
@@ -24,9 +24,12 @@ const LOG_MAX_ROWS: u32 = 22;
 const IS_FIRST_LOG_SIZES: [u32; 19] = [
     22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4,
 ];
+
 pub fn prove_cairo(
     input: CairoInput,
+    // TODO(Ohad): wrap these flags in a struct.
     track_relations: bool,
+    display_components: bool,
 ) -> Result<CairoProof<Blake2sMerkleHasher>, ProvingError> {
     let _span = span!(Level::INFO, "prove_cairo").entered();
     let config = PcsConfig::default();
@@ -94,6 +97,10 @@ pub fn prove_cairo(
     // Prove stark.
     let proof = prove::<SimdBackend, _>(&components, channel, commitment_scheme)?;
 
+    if display_components {
+        println!("{}", component_builder);
+    }
+
     Ok(CairoProof {
         claim,
         interaction_claim,
@@ -183,14 +190,14 @@ mod tests {
 
     #[test]
     fn test_basic_cairo_air() {
-        let cairo_proof = prove_cairo(test_input(), true).unwrap();
+        let cairo_proof = prove_cairo(test_input(), true, true).unwrap();
         verify_cairo(cairo_proof).unwrap();
     }
 
     #[ignore]
     #[test]
     fn test_full_cairo_air() {
-        let cairo_proof = prove_cairo(small_cairo_input(), true).unwrap();
+        let cairo_proof = prove_cairo(small_cairo_input(), true, true).unwrap();
         verify_cairo(cairo_proof).unwrap();
     }
 }

stwo_cairo_prover/crates/prover/src/cairo_air/opcodes_air.rs

@@ -10,6 +10,7 @@ use stwo_prover::core::pcs::{TreeBuilder, TreeVec};
 use stwo_prover::core::vcs::blake2_merkle::Blake2sMerkleChannel;
 
 use super::air::CairoInteractionElements;
+use super::debug_tools::display_components;
 use crate::components::{
     add_ap_opcode_is_imm_f_op1_base_fp_f, add_ap_opcode_is_imm_f_op1_base_fp_t,
     add_ap_opcode_is_imm_t_op1_base_fp_f, add_opcode_is_small_f_is_imm_f,
@@ -2260,3 +2261,59 @@ impl OpcodeComponents {
         vec
     }
 }
+
+impl std::fmt::Display for OpcodeComponents {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        writeln!(f, "add_f_f:")?;
+        writeln!(f, "{}", display_components(&self.add_f_f))?;
+        writeln!(f, "add_f_t:")?;
+        writeln!(f, "{}", display_components(&self.add_f_t))?;
+        writeln!(f, "add_t_f:")?;
+        writeln!(f, "{}", display_components(&self.add_t_f))?;
+        writeln!(f, "add_t_t:")?;
+        writeln!(f, "{}", display_components(&self.add_t_t))?;
+        writeln!(f, "add_ap_f_f:")?;
+        writeln!(f, "{}", display_components(&self.add_ap_f_f))?;
+        writeln!(f, "add_ap_f_t:")?;
+        writeln!(f, "{}", display_components(&self.add_ap_f_t))?;
+        writeln!(f, "add_ap_t_f:")?;
+        writeln!(f, "{}", display_components(&self.add_ap_t_f))?;
+        writeln!(f, "assert_eq_f_f:")?;
+        writeln!(f, "{}", display_components(&self.assert_eq_f_f))?;
+        writeln!(f, "assert_eq_f_t:")?;
+        writeln!(f, "{}", display_components(&self.assert_eq_f_t))?;
+        writeln!(f, "assert_eq_t_f:")?;
+        writeln!(f, "{}", display_components(&self.assert_eq_t_f))?;
+        writeln!(f, "call_f_f:")?;
+        writeln!(f, "{}", display_components(&self.call_f_f))?;
+        writeln!(f, "call_f_t:")?;
+        writeln!(f, "{}", display_components(&self.call_f_t))?;
+        writeln!(f, "call_t_f:")?;
+        writeln!(f, "{}", display_components(&self.call_t_f))?;
+        writeln!(f, "generic:")?;
+        writeln!(f, "{}", display_components(&self.generic))?;
+        writeln!(f, "jnz_f_f:")?;
+        writeln!(f, "{}", display_components(&self.jnz_f_f))?;
+        writeln!(f, "jnz_f_t:")?;
+        writeln!(f, "{}", display_components(&self.jnz_f_t))?;
+        writeln!(f, "jnz_t_f:")?;
+        writeln!(f, "{}", display_components(&self.jnz_t_f))?;
+        writeln!(f, "jnz_t_t:")?;
+        writeln!(f, "{}", display_components(&self.jnz_t_t))?;
+        writeln!(f, "jump_f_f_f:")?;
+        writeln!(f, "{}", display_components(&self.jump_f_f_f))?;
+        writeln!(f, "jump_f_f_t:")?;
+        writeln!(f, "{}", display_components(&self.jump_f_f_t))?;
+        writeln!(f, "jump_t_f_f:")?;
+        writeln!(f, "{}", display_components(&self.jump_t_f_f))?;
+        writeln!(f, "jump_t_t_f:")?;
+        writeln!(f, "{}", display_components(&self.jump_t_t_f))?;
+        writeln!(f, "mul_f_f:")?;
+        writeln!(f, "{}", display_components(&self.mul_f_f))?;
+        writeln!(f, "mul_f_t:")?;
+        writeln!(f, "{}", display_components(&self.mul_f_t))?;
+        writeln!(f, "ret:")?;
+        writeln!(f, "{}", display_components(&self.ret))?;
+        Ok(())
+    }
+}