Expr formatting.

starkware-libs · Nov 18, 2024 · e55c200 · e55c200
1 parent 96f7025
commit e55c200
Showing 1 changed file with 53 additions and 0 deletions.
diff --git 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 {}