From 89c469c60b977ca801d0b41f6dad0682ec17cb7b Mon Sep 17 00:00:00 2001
From: Alon Titelman <alon.titelman@starkware.co>
Date: Thu, 14 Nov 2024 16:12:25 +0200
Subject: [PATCH] Expr formatting.

---
 .../prover/src/constraint_framework/expr.rs   | 53 +++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/crates/prover/src/constraint_framework/expr.rs b/crates/prover/src/constraint_framework/expr.rs
index fb86edac8..70ceaa700 100644
--- a/crates/prover/src/constraint_framework/expr.rs
+++ b/crates/prover/src/constraint_framework/expr.rs
@@ -36,6 +36,35 @@ enum Expr {
     Inv(Box<Expr>),
 }
 
+impl Expr {
+    #[allow(dead_code)]
+    pub fn format_expr(&self) -> String {
+        match self {
+            Expr::Col(ColumnExpr {
+                interaction,
+                idx,
+                offset,
+            }) => {
+                format!("col_{interaction}_{idx}[{offset}]")
+            }
+            Expr::SecureCol([a, b, c, d]) => format!(
+                "SecureCol({}, {}, {}, {})",
+                a.format_expr(),
+                b.format_expr(),
+                c.format_expr(),
+                d.format_expr()
+            ),
+            Expr::Const(c) => c.0.to_string(),
+            Expr::Param(v) => v.to_string(),
+            Expr::Add(a, b) => format!("{} + {}", a.format_expr(), b.format_expr()),
+            Expr::Sub(a, b) => format!("{} - ({})", a.format_expr(), b.format_expr()),
+            Expr::Mul(a, b) => format!("({}) * ({})", a.format_expr(), b.format_expr()),
+            Expr::Neg(a) => format!("-({})", a.format_expr()),
+            Expr::Inv(a) => format!("1/({})", a.format_expr()),
+        }
+    }
+}
+
 impl From<BaseField> for Expr {
     fn from(val: BaseField) -> Self {
         Expr::Const(val)
@@ -256,6 +285,7 @@ mod tests {
     use crate::core::fields::m31::M31;
     use crate::core::fields::qm31::SecureField;
     use crate::core::fields::FieldExpOps;
+
     #[test]
     fn test_expr_eval() {
         let test_struct = TestStruct {};
@@ -414,6 +444,29 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_format_expr() {
+        let test_struct = TestStruct {};
+        let eval = test_struct.evaluate(ExprEvaluator::new(16, (SecureField::zero(), None)));
+        let constraint0_str =  "(1) * ((((col_1_0[0]) * (col_1_1[0])) * (col_1_2[0])) * (1/(col_1_0[0] + col_1_1[0])))";
+        assert_eq!(eval.constraints[0].format_expr(), constraint0_str);
+        let constraint1_str = "(1) \
+            * ((SecureCol(col_2_4[0], col_2_6[0], col_2_8[0], col_2_10[0]) \
+                - (SecureCol(\
+                    col_2_5[-1], \
+                    col_2_7[-1], \
+                    col_2_9[-1], \
+                    col_2_11[-1]\
+                ) - ((col_0_3[0]) * (SecureCol(0, 0, 0, 0)))) \
+                - (0)) \
+                * (0 + (TestRelation_alpha0) * (col_1_0[0]) \
+                    + (TestRelation_alpha1) * (col_1_1[0]) \
+                    + (TestRelation_alpha2) * (col_1_2[0]) \
+                    - (TestRelation_z)) \
+                - (1))";
+        assert_eq!(eval.constraints[1].format_expr(), constraint1_str);
+    }
+
     relation!(TestRelation, 3);
 
     struct TestStruct {}