checkpoint

halo2_debug/src/display.rs

@@ -0,0 +1,353 @@
+use ff::PrimeField;
+use halo2_middleware::circuit::{ColumnMid, VarMid};
+use halo2_middleware::expression::{Expression, Variable};
+use halo2_middleware::{lookup, shuffle};
+use num_bigint::BigUint;
+use std::collections::HashMap;
+use std::fmt;
+
+/// Wrapper type over `PrimeField` that implements Display with nice output.
+/// - If the value is a power of two, format it as `2^k`
+/// - If the value is smaller than 2^16, format it in decimal
+/// - If the value is bigger than congruent -2^16, format it in decimal as the negative congruent
+/// (between -2^16 and 0).
+/// - Else format it in hex without leading zeros.
+pub struct FDisp<'a, F: PrimeField>(pub &'a F);
+
+impl<F: PrimeField> fmt::Display for FDisp<'_, F> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let v = (*self.0).to_repr();
+        let v = v.as_ref();
+        let v = BigUint::from_bytes_le(v.as_ref());
+        let v_bits = v.bits();
+        if v_bits >= 8 && v.count_ones() == 1 {
+            write!(f, "2^{}", v.trailing_zeros().unwrap_or_default())
+        } else if v_bits < 16 {
+            write!(f, "{}", v)
+        } else {
+            let v_neg = (F::ZERO - self.0).to_repr();
+            let v_neg = v_neg.as_ref();
+            let v_neg = BigUint::from_bytes_le(v_neg.as_ref());
+            let v_neg_bits = v_neg.bits();
+            if v_neg_bits < 16 {
+                write!(f, "-{}", v_neg)
+            } else {
+                write!(f, "0x{:x}", v)
+            }
+        }
+    }
+}
+
+/// Wrapper type over `Expression` that implements Display with nice output.
+/// The formatting of the `Expression::Variable` case is parametrized with the second field, which
+/// take as auxiliary value the third field.
+/// Use the constructor `expr_disp` to format variables using their `Display` implementation.
+/// Use the constructor `expr_disp_names` for an `Expression` with `V: VarMid` to format column
+/// queries according to their string names.
+pub struct ExprDisp<'a, F: PrimeField, V: Variable, A>(
+    /// Expression to display
+    pub &'a Expression<F, V>,
+    /// `V: Variable` formatter method that uses auxiliary value
+    pub fn(&V, &mut fmt::Formatter<'_>, a: &A) -> fmt::Result,
+    /// Auxiliary value to be passed to the `V: Variable` formatter
+    pub &'a A,
+);
+
+fn var_fmt_default<V: Variable>(v: &V, f: &mut fmt::Formatter<'_>, _: &()) -> fmt::Result {
+    write!(f, "{}", v)
+}
+
+fn var_fmt_names(
+    v: &VarMid,
+    f: &mut fmt::Formatter<'_>,
+    names: &HashMap<ColumnMid, String>,
+) -> fmt::Result {
+    if let VarMid::Query(q) = v {
+        if let Some(name) = names.get(&ColumnMid::new(q.column_type, q.column_index)) {
+            return write!(f, "{}", name);
+        }
+    }
+    write!(f, "{}", v)
+}
+
+/// ExprDisp constructor that formats viariables using their `Display` implementation.
+pub fn expr_disp<'a, F: PrimeField, V: Variable>(
+    e: &'a Expression<F, V>,
+) -> ExprDisp<'a, F, V, ()> {
+    ExprDisp(e, var_fmt_default, &())
+}
+
+/// ExprDisp constructor for an `Expression` with `V: VarMid` that formats column queries according
+/// to their string names.
+pub fn expr_disp_names<'a, F: PrimeField>(
+    e: &'a Expression<F, VarMid>,
+    names: &'a HashMap<ColumnMid, String>,
+) -> ExprDisp<'a, F, VarMid, HashMap<ColumnMid, String>> {
+    ExprDisp(e, var_fmt_names, names)
+}
+
+impl<F: PrimeField, V: Variable, A> fmt::Display for ExprDisp<'_, F, V, A> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let is_sum = |e: &Expression<F, V>| -> bool { matches!(e, Expression::Sum(_, _)) };
+        let fmt_expr =
+            |e: &Expression<F, V>, f: &mut fmt::Formatter<'_>, parens: bool| -> fmt::Result {
+                if parens {
+                    write!(f, "(")?;
+                }
+                write!(f, "{}", ExprDisp(e, self.1, self.2))?;
+                if parens {
+                    write!(f, ")")?;
+                }
+                Ok(())
+            };
+
+        match self.0 {
+            Expression::Constant(c) => write!(f, "{}", FDisp(c)),
+            Expression::Var(v) => self.1(v, f, self.2),
+            Expression::Negated(a) => {
+                write!(f, "-")?;
+                fmt_expr(&a, f, is_sum(&a))
+            }
+            Expression::Sum(a, b) => {
+                fmt_expr(&a, f, false)?;
+                if let Expression::Negated(neg) = &**b {
+                    write!(f, " - ")?;
+                    fmt_expr(&neg, f, is_sum(&neg))
+                } else {
+                    write!(f, " + ")?;
+                    fmt_expr(&b, f, false)
+                }
+            }
+            Expression::Product(a, b) => {
+                fmt_expr(&a, f, is_sum(&a))?;
+                write!(f, " * ")?;
+                fmt_expr(&b, f, is_sum(&b))
+            }
+        }
+    }
+}
+
+/// Wrapper type over `lookup::Argument` that implements Display with nice output.
+/// The formatting of the `Expression::Variable` case is parametrized with the second field, which
+/// take as auxiliary value the third field.
+/// Use the constructor `lookup_arg_disp` to format variables using their `Display` implementation.
+/// Use the constructor `lookup_arg_disp_names` for a lookup of `Expression` with `V: VarMid` that
+/// formats column queries according to their string names.
+pub struct LookupArgDisp<'a, F: PrimeField, V: Variable, A>(
+    /// Lookup argument to display
+    pub &'a lookup::Argument<F, V>,
+    /// `V: Variable` formatter method that uses auxiliary value
+    pub fn(&V, &mut fmt::Formatter<'_>, a: &A) -> fmt::Result,
+    /// Auxiliary value to be passed to the `V: Variable` formatter
+    pub &'a A,
+);
+
+/// LookupArgDisp constructor that formats viariables using their `Display` implementation.
+pub fn lookup_arg_disp<'a, F: PrimeField, V: Variable>(
+    a: &'a lookup::Argument<F, V>,
+) -> LookupArgDisp<'a, F, V, ()> {
+    LookupArgDisp(a, var_fmt_default, &())
+}
+
+/// LookupArgDisp constructor for a lookup of `Expression` with `V: VarMid` that formats column
+/// queries according to their string names.
+pub fn lookup_arg_disp_names<'a, F: PrimeField>(
+    a: &'a lookup::Argument<F, VarMid>,
+    names: &'a HashMap<ColumnMid, String>,
+) -> LookupArgDisp<'a, F, VarMid, HashMap<ColumnMid, String>> {
+    LookupArgDisp(a, var_fmt_names, names)
+}
+
+impl<F: PrimeField, V: Variable, A> fmt::Display for LookupArgDisp<'_, F, V, A> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "[")?;
+        for (i, expr) in self.0.input_expressions.iter().enumerate() {
+            if i != 0 {
+                write!(f, ", ")?;
+            }
+            write!(f, "{}", ExprDisp(expr, self.1, self.2))?;
+        }
+        write!(f, "] in [")?;
+        for (i, expr) in self.0.table_expressions.iter().enumerate() {
+            if i != 0 {
+                write!(f, ", ")?;
+            }
+            write!(f, "{}", ExprDisp(expr, self.1, self.2))?;
+        }
+        write!(f, "]")?;
+        Ok(())
+    }
+}
+
+/// Wrapper type over `shuffle::Argument` that implements Display with nice output.
+/// The formatting of the `Expression::Variable` case is parametrized with the second field, which
+/// take as auxiliary value the third field.
+/// Use the constructor `shuffle_arg_disp` to format variables using their `Display`
+/// implementation.
+/// Use the constructor `shuffle_arg_disp_names` for a shuffle of `Expression` with `V: VarMid`
+/// that formats column queries according to their string names.
+pub struct ShuffleArgDisp<'a, F: PrimeField, V: Variable, A>(
+    /// Shuffle argument to display
+    pub &'a shuffle::Argument<F, V>,
+    /// `V: Variable` formatter method that uses auxiliary value
+    pub fn(&V, &mut fmt::Formatter<'_>, a: &A) -> fmt::Result,
+    /// Auxiliary value to be passed to the `V: Variable` formatter
+    pub &'a A,
+);
+
+/// ShuffleArgDisp constructor that formats viariables using their `Display` implementation.
+pub fn shuffle_arg_disp<'a, F: PrimeField, V: Variable>(
+    a: &'a shuffle::Argument<F, V>,
+) -> ShuffleArgDisp<'a, F, V, ()> {
+    ShuffleArgDisp(a, var_fmt_default, &())
+}
+
+/// ShuffleArgDisp constructor for a shuffle of `Expression` with `V: VarMid` that formats column
+/// queries according to their string names.
+pub fn shuffle_arg_disp_names<'a, F: PrimeField>(
+    a: &'a shuffle::Argument<F, VarMid>,
+    names: &'a HashMap<ColumnMid, String>,
+) -> ShuffleArgDisp<'a, F, VarMid, HashMap<ColumnMid, String>> {
+    ShuffleArgDisp(a, var_fmt_names, names)
+}
+
+impl<F: PrimeField, V: Variable, A> fmt::Display for ShuffleArgDisp<'_, F, V, A> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "[")?;
+        for (i, expr) in self.0.input_expressions.iter().enumerate() {
+            if i != 0 {
+                write!(f, ", ")?;
+            }
+            write!(f, "{}", ExprDisp(expr, self.1, self.2))?;
+        }
+        write!(f, "] shuff [")?;
+        for (i, expr) in self.0.shuffle_expressions.iter().enumerate() {
+            if i != 0 {
+                write!(f, ", ")?;
+            }
+            write!(f, "{}", ExprDisp(expr, self.1, self.2))?;
+        }
+        write!(f, "]")?;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use ff::Field;
+    use halo2_middleware::circuit::{Any, QueryMid, VarMid};
+    use halo2_middleware::poly::Rotation;
+    use halo2curves::bn256::Fr;
+
+    #[test]
+    fn test_lookup_shuffle_arg_disp() {
+        type E = Expression<Fr, VarMid>;
+        let a0 = VarMid::Query(QueryMid::new(Any::Advice, 0, Rotation(0)));
+        let a1 = VarMid::Query(QueryMid::new(Any::Advice, 1, Rotation(0)));
+        let f0 = VarMid::Query(QueryMid::new(Any::Fixed, 0, Rotation(0)));
+        let a0: E = Expression::Var(a0);
+        let a1: E = Expression::Var(a1);
+        let f0: E = Expression::Var(f0);
+
+        let names = [
+            (ColumnMid::new(Any::Advice, 0), "a".to_string()),
+            (ColumnMid::new(Any::Advice, 1), "b".to_string()),
+            (ColumnMid::new(Any::Fixed, 0), "s".to_string()),
+        ]
+        .into_iter()
+        .collect();
+
+        let arg = lookup::Argument {
+            name: "lookup".to_string(),
+            input_expressions: vec![f0.clone() * a0.clone(), f0.clone() * a1.clone()],
+            table_expressions: vec![f0.clone(), f0.clone() * (a0.clone() + a1.clone())],
+        };
+        assert_eq!(
+            "[f0 * a0, f0 * a1] in [f0, f0 * (a0 + a1)]",
+            format!("{}", lookup_arg_disp(&arg))
+        );
+        assert_eq!(
+            "[s * a, s * b] in [s, s * (a + b)]",
+            format!("{}", lookup_arg_disp_names(&arg, &names))
+        );
+
+        let arg = shuffle::Argument {
+            name: "shuffle".to_string(),
+            input_expressions: vec![f0.clone() * a0.clone(), f0.clone() * a1.clone()],
+            shuffle_expressions: vec![f0.clone(), f0.clone() * (a0.clone() + a1.clone())],
+        };
+        assert_eq!(
+            "[f0 * a0, f0 * a1] shuff [f0, f0 * (a0 + a1)]",
+            format!("{}", shuffle_arg_disp(&arg))
+        );
+        assert_eq!(
+            "[s * a, s * b] shuff [s, s * (a + b)]",
+            format!("{}", shuffle_arg_disp_names(&arg, &names))
+        );
+    }
+
+    #[test]
+    fn test_expr_disp() {
+        type E = Expression<Fr, VarMid>;
+        let a0 = VarMid::Query(QueryMid::new(Any::Advice, 0, Rotation(0)));
+        let a1 = VarMid::Query(QueryMid::new(Any::Advice, 1, Rotation(0)));
+        let a0: E = Expression::Var(a0);
+        let a1: E = Expression::Var(a1);
+
+        let e = a0.clone() + a1.clone();
+        assert_eq!("a0 + a1", format!("{}", expr_disp(&e)));
+        let e = a0.clone() + a1.clone() + a0.clone();
+        assert_eq!("a0 + a1 + a0", format!("{}", expr_disp(&e)));
+
+        let e = a0.clone() * a1.clone();
+        assert_eq!("a0 * a1", format!("{}", expr_disp(&e)));
+        let e = a0.clone() * a1.clone() * a0.clone();
+        assert_eq!("a0 * a1 * a0", format!("{}", expr_disp(&e)));
+
+        let e = a0.clone() - a1.clone();
+        assert_eq!("a0 - a1", format!("{}", expr_disp(&e)));
+        let e = (a0.clone() - a1.clone()) - a0.clone();
+        assert_eq!("a0 - a1 - a0", format!("{}", expr_disp(&e)));
+        let e = a0.clone() - (a1.clone() - a0.clone());
+        assert_eq!("a0 - (a1 - a0)", format!("{}", expr_disp(&e)));
+
+        let e = a0.clone() * a1.clone() + a0.clone();
+        assert_eq!("a0 * a1 + a0", format!("{}", expr_disp(&e)));
+        let e = a0.clone() * (a1.clone() + a0.clone());
+        assert_eq!("a0 * (a1 + a0)", format!("{}", expr_disp(&e)));
+
+        let e = a0.clone() + a1.clone();
+        let names = [
+            (ColumnMid::new(Any::Advice, 0), "a".to_string()),
+            (ColumnMid::new(Any::Advice, 1), "b".to_string()),
+        ]
+        .into_iter()
+        .collect();
+        assert_eq!("a + b", format!("{}", expr_disp_names(&e, &names)));
+    }
+
+    #[test]
+    fn test_f_disp() {
+        let v = Fr::ZERO;
+        assert_eq!("0", format!("{}", FDisp(&v)));
+
+        let v = Fr::ONE;
+        assert_eq!("1", format!("{}", FDisp(&v)));
+
+        let v = Fr::from(12345u64);
+        assert_eq!("12345", format!("{}", FDisp(&v)));
+
+        let v = Fr::from(0x10000);
+        assert_eq!("2^16", format!("{}", FDisp(&v)));
+
+        let v = Fr::from(0x12345);
+        assert_eq!("0x12345", format!("{}", FDisp(&v)));
+
+        let v = Fr::from(-Fr::ONE);
+        assert_eq!("-1", format!("{}", FDisp(&v)));
+
+        let v = Fr::from(-Fr::from(12345u64));
+        assert_eq!("-12345", format!("{}", FDisp(&v)));
+    }
+}