Improve mock prover errors printing (#182)

- fmt `Goldilocks(18446744069414584319)` as `-2` as well as `Sum`
`Product` expressions.
- include instance id and relevant wit in.

### Before

```
AssertEqualError "test_assert_zero_1/require_equal/a + 1 == b": Left != Right
Left: GoldilocksExt2([Goldilocks(4), Goldilocks(0)])
Right: GoldilocksExt2([Goldilocks(5), Goldilocks(0)])
Left Expression: ScaledSum(WitIn(1), Constant(Goldilocks(1)), Constant(Goldilocks(0)))
Right Expression: Sum(ScaledSum(WitIn(0), Constant(Goldilocks(1)), Constant(Goldilocks(0))), Constant(Goldilocks(1)))

AssertZeroError "test_assert_zero_1/require_zero/c - 2 == 0": Evaluated expression is not zero
Expression: Sum(ScaledSum(WitIn(2), Constant(Goldilocks(1)), Constant(Goldilocks(0))), Constant(Goldilocks(18446744069414584319)))
Evaluation: GoldilocksExt2([Goldilocks(18446744069414584322), Goldilocks(0)])

LookupError "test_lookup_1/assert_u5/assert u5": Evaluated expression does not exist in T vector
Expression: ScaledSum(WitIn(0), Challenge(1, 1, GoldilocksExt2([Goldilocks(1), Goldilocks(0)]), GoldilocksExt2([Goldilocks(0), Goldilocks(0)])), Challenge(0, 1, GoldilocksExt2([Goldilocks(1), Goldilocks(0)]), GoldilocksExt2([Goldilocks(0), Goldilocks(0)])))
Evaluation: GoldilocksExt2([Goldilocks(1000001), Goldilocks(0)])
```

### After

```
AssertEqualError "test_assert_zero_1/require_equal/a + 1 == b"
Left: GoldilocksExt2[4,0] != Right: GoldilocksExt2[5,0]
Left Expression: 1 * WitIn(1) + 0
Right Expression: 1 * WitIn(0) + 0 + 1
Inst[0]: WitIn(1)=4,WitIn(0)=4

AssertZeroError "test_assert_zero_1/require_zero/c - 2 == 0": Evaluated expression is not zero
Expression: 1 * WitIn(2) + 0 + (-2)
Evaluation: GoldilocksExt2[1,0] != 0
Inst[1]: WitIn(2)=3

LookupError "test_lookup_1/assert_u5/assert u5": Evaluated expression does not exist in T vector
Expression: Challenge(1) * WitIn(0) + Challenge(0)
Evaluation: GoldilocksExt2[1000001,0]
Inst[1]: WitIn(0)=1000
```

---------

Co-authored-by: Ming <[email protected]>

ceno_zkvm/src/scheme/mock_prover.rs

@@ -1,6 +1,11 @@
 use super::utils::{eval_by_expr, wit_infer_by_expr};
-use crate::{circuit_builder::CircuitBuilder, expression::Expression, structs::ROMType};
+use crate::{
+    circuit_builder::CircuitBuilder,
+    expression::Expression,
+    structs::{ROMType, WitnessId},
+};
 use ff_ext::ExtensionField;
+use goldilocks::SmallField;
 use itertools::Itertools;
 use multilinear_extensions::virtual_poly_v2::ArcMultilinearExtension;
 use std::{marker::PhantomData, ops::Neg};
@@ -12,35 +17,46 @@ pub(crate) enum MockProverError<E: ExtensionField> {
         expression: Expression<E>,
         evaluated: E,
         name: String,
+        inst_id: usize,
     },
     AssertEqualError {
         left_expression: Expression<E>,
         right_expression: Expression<E>,
         left: E,
         right: E,
         name: String,
+        inst_id: usize,
     },
     LookupError {
         expression: Expression<E>,
         evaluated: E,
         name: String,
+        inst_id: usize,
     },
     // TODO later
     // r_expressions
     // w_expressions
 }
 
 impl<E: ExtensionField> MockProverError<E> {
-    pub fn print(&self) {
+    pub fn print(&self, wits_in: &[ArcMultilinearExtension<E>]) {
+        let mut wtns = vec![];
+
         match self {
             Self::AssertZeroError {
                 expression,
                 evaluated,
                 name,
+                inst_id,
             } => {
+                let expression_fmt = fmt_expr(expression, &mut wtns, false);
+                let wtns_fmt = fmt_wtns::<E>(&wtns, wits_in, *inst_id);
+                let eval_fmt = fmt_field::<E>(evaluated);
                 println!(
-                    "\nAssertZeroError {name:#?}: Evaluated expression is not zero\nExpression: \
-                    {expression:?}\nEvaluation: {evaluated:?}\n",
+                    "\nAssertZeroError {name:?}: Evaluated expression is not zero\n\
+                    Expression: {expression_fmt}\n\
+                    Evaluation: {eval_fmt} != 0\n\
+                    Inst[{inst_id}]: {wtns_fmt}\n",
                 );
             }
             Self::AssertEqualError {
@@ -49,25 +65,126 @@ impl<E: ExtensionField> MockProverError<E> {
                 left,
                 right,
                 name,
+                inst_id,
             } => {
+                let left_expression_fmt = fmt_expr(left_expression, &mut wtns, false);
+                let right_expression_fmt = fmt_expr(right_expression, &mut wtns, false);
+                let wtns_fmt = fmt_wtns::<E>(&wtns, wits_in, *inst_id);
+                let left_eval_fmt = fmt_field::<E>(left);
+                let right_eval_fmt = fmt_field::<E>(right);
                 println!(
-                    "\nAssertEqualError {name:#?}: Left != Right\n\
-                    Left: {left:?}\nRight: {right:?}\n\
-                    Left Expression: {left_expression:?}\n\
-                    Right Expression: {right_expression:?}\n",
+                    "\nAssertEqualError {name:?}\n\
+                    Left: {left_eval_fmt} != Right: {right_eval_fmt}\n\
+                    Left Expression: {left_expression_fmt}\n\
+                    Right Expression: {right_expression_fmt}\n\
+                    Inst[{inst_id}]: {wtns_fmt}\n",
                 );
             }
             Self::LookupError {
                 expression,
                 evaluated,
                 name,
+                inst_id,
             } => {
+                let expression_fmt = fmt_expr(expression, &mut wtns, false);
+                let wtns_fmt = fmt_wtns::<E>(&wtns, wits_in, *inst_id);
+                let eval_fmt = fmt_field::<E>(evaluated);
                 println!(
-                    "\nLookupError {name:#?}: Evaluated expression does not exist in T \
-                    vector\nExpression: {expression:?}\nEvaluation: {evaluated:?}\n",
+                    "\nLookupError {name:#?}: Evaluated expression does not exist in T vector\n\
+                    Expression: {expression_fmt}\n\
+                    Evaluation: {eval_fmt}\n\
+                    Inst[{inst_id}]: {wtns_fmt}\n",
                 );
             }
         }
+
+        fn fmt_expr<E: ExtensionField>(
+            expression: &Expression<E>,
+            wtns: &mut Vec<WitnessId>,
+            add_prn_sum: bool,
+        ) -> String {
+            match expression {
+                Expression::WitIn(wit_in) => {
+                    wtns.push(*wit_in);
+                    format!("WitIn({})", wit_in)
+                }
+                Expression::Challenge(id, _, _, _) => format!("Challenge({})", id),
+                Expression::Constant(constant) => fmt_base_field::<E>(constant).to_string(),
+                Expression::Fixed(fixed) => format!("{:?}", fixed),
+                Expression::Sum(left, right) => {
+                    let s = format!(
+                        "{} + {}",
+                        fmt_expr(left, wtns, false),
+                        fmt_expr(right, wtns, false)
+                    );
+                    if add_prn_sum { format!("({})", s) } else { s }
+                }
+                Expression::Product(left, right) => {
+                    format!(
+                        "{} * {}",
+                        fmt_expr(left, wtns, true),
+                        fmt_expr(right, wtns, true)
+                    )
+                }
+                Expression::ScaledSum(x, a, b) => {
+                    let s = format!(
+                        "{} * {} + {}",
+                        fmt_expr(a, wtns, true),
+                        fmt_expr(x, wtns, true),
+                        fmt_expr(b, wtns, false)
+                    );
+                    if add_prn_sum { format!("({})", s) } else { s }
+                }
+            }
+        }
+
+        fn fmt_field<E: ExtensionField>(field: &E) -> String {
+            let name = format!("{:?}", field);
+            let name = name.split('(').next().unwrap_or("ExtensionField");
+            format!(
+                "{name}[{}]",
+                field
+                    .as_bases()
+                    .iter()
+                    .map(fmt_base_field::<E>)
+                    .collect::<Vec<String>>()
+                    .join(",")
+            )
+        }
+
+        fn fmt_base_field<E: ExtensionField>(base_field: &E::BaseField) -> String {
+            let value = base_field.to_canonical_u64();
+
+            if value > E::BaseField::MODULUS_U64 - u16::MAX as u64 {
+                // beautiful format for negative number > -65536
+                format!("(-{})", E::BaseField::MODULUS_U64 - value)
+            } else if value < u16::MAX as u64 {
+                format!("{value}")
+            } else {
+                // hex
+                format!("{value:#x}")
+            }
+        }
+
+        fn fmt_wtns<E: ExtensionField>(
+            wtns: &[WitnessId],
+            wits_in: &[ArcMultilinearExtension<E>],
+            inst_id: usize,
+        ) -> String {
+            wtns.iter()
+                .map(|wt_id| {
+                    let wit = &wits_in[*wt_id as usize];
+                    let value_fmt = if let Some(e) = wit.get_ext_field_vec_optn() {
+                        fmt_field(&e[inst_id])
+                    } else if let Some(bf) = wit.get_base_field_vec_optn() {
+                        fmt_base_field::<E>(&bf[inst_id])
+                    } else {
+                        "Unknown".to_string()
+                    };
+                    format!("WitIn({wt_id})={value_fmt}")
+                })
+                .join(",")
+        }
     }
 }
 
@@ -111,28 +228,33 @@ impl<'a, E: ExtensionField> MockProver<E> {
                 let right_evaluated = wit_infer_by_expr(&[], wits_in, &challenge, &right);
                 let right_evaluated = right_evaluated.get_ext_field_vec();
 
-                for (left_element, right_element) in left_evaluated.iter().zip_eq(right_evaluated) {
+                for (inst_id, (left_element, right_element)) in
+                    left_evaluated.iter().zip_eq(right_evaluated).enumerate()
+                {
                     if *left_element != *right_element {
                         errors.push(MockProverError::AssertEqualError {
                             left_expression: left.clone(),
                             right_expression: right.clone(),
                             left: *left_element,
                             right: *right_element,
                             name: name.clone(),
+                            inst_id,
                         });
                     }
                 }
             } else {
+                // contains require_zero
                 let expr = expr.clone().neg().neg(); // TODO get_ext_field_vec doesn't work without this
                 let expr_evaluated = wit_infer_by_expr(&[], wits_in, &challenge, &expr);
                 let expr_evaluated = expr_evaluated.get_ext_field_vec();
 
-                for element in expr_evaluated {
+                for (inst_id, element) in expr_evaluated.iter().enumerate() {
                     if *element != E::ZERO {
                         errors.push(MockProverError::AssertZeroError {
                             expression: expr.clone(),
                             evaluated: *element,
                             name: name.clone(),
+                            inst_id,
                         });
                     }
                 }
@@ -154,12 +276,13 @@ impl<'a, E: ExtensionField> MockProver<E> {
             let expr_evaluated = expr_evaluated.get_ext_field_vec();
 
             // Check each lookup expr exists in t vec
-            for element in expr_evaluated {
+            for (inst_id, element) in expr_evaluated.iter().enumerate() {
                 if !table_vec.contains(element) {
                     errors.push(MockProverError::LookupError {
                         expression: expr.clone(),
                         evaluated: *element,
                         name: name.clone(),
+                        inst_id,
                     });
                 }
             }
@@ -182,9 +305,13 @@ impl<'a, E: ExtensionField> MockProver<E> {
         match result {
             Ok(_) => {}
             Err(errors) => {
+                println!("======================================================");
+                println!("Error: {} constraints not satisfied", errors.len());
+
                 for error in errors {
-                    error.print();
+                    error.print(wits_in);
                 }
+                println!("======================================================");
                 panic!("Constraints not satisfied");
             }
         }
@@ -344,6 +471,7 @@ mod tests {
                 ),
                 evaluated: 123002.into(), // 123 * 1000 + 2
                 name: "test_lookup_error/assert_u5/assert u5".to_string(),
+                inst_id: 0,
             }]
         );
     }

multilinear_extensions/src/mle.rs

@@ -44,22 +44,32 @@ pub trait MultilinearExtension<E: ExtensionField>: Send + Sync {
     fn name(&self) -> &'static str;
 
     fn get_ext_field_vec(&self) -> &[E] {
+        self.get_ext_field_vec_optn()
+            .unwrap_or_else(|| unreachable!())
+    }
+
+    fn get_base_field_vec(&self) -> &[E::BaseField] {
+        self.get_base_field_vec_optn()
+            .unwrap_or_else(|| unreachable!())
+    }
+
+    fn get_ext_field_vec_optn(&self) -> Option<&[E]> {
         match &self.evaluations() {
             FieldType::Ext(evaluations) => {
                 let (start, offset) = self.evaluations_range().unwrap_or((0, evaluations.len()));
-                &evaluations[start..][..offset]
+                Some(&evaluations[start..][..offset])
             }
-            _ => unreachable!(),
+            _ => None,
         }
     }
 
-    fn get_base_field_vec(&self) -> &[E::BaseField] {
+    fn get_base_field_vec_optn(&self) -> Option<&[E::BaseField]> {
         match &self.evaluations() {
             FieldType::Base(evaluations) => {
                 let (start, offset) = self.evaluations_range().unwrap_or((0, evaluations.len()));
-                &evaluations[start..][..offset]
+                Some(&evaluations[start..][..offset])
             }
-            _ => unreachable!(),
+            _ => None,
         }
     }
 }