checkpoint

halo2_debug/src/lib.rs

@@ -1,6 +1,9 @@
-use ff::{Field, PrimeField};
+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.
@@ -35,17 +38,63 @@ impl<F: PrimeField> fmt::Display for FDisp<'_, F> {
     }
 }
 
-pub struct ExprDisp<'a, F: PrimeField, V: Variable>(pub &'a Expression<F, 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,
+);
 
-impl<F: PrimeField, V: Variable> fmt::Display for ExprDisp<'_, F, V> {
+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))?;
+                write!(f, "{}", ExprDisp(e, self.1, self.2))?;
                 if parens {
                     write!(f, ")")?;
                 }
@@ -54,7 +103,7 @@ impl<F: PrimeField, V: Variable> fmt::Display for ExprDisp<'_, F, V> {
 
         match self.0 {
             Expression::Constant(c) => write!(f, "{}", FDisp(c)),
-            Expression::Var(v) => write!(f, "{}", v),
+            Expression::Var(v) => self.1(v, f, self.2),
             Expression::Negated(a) => {
                 write!(f, "-")?;
                 fmt_expr(&a, f, is_sum(&a))
@@ -78,50 +127,204 @@ impl<F: PrimeField, V: Variable> fmt::Display for ExprDisp<'_, F, V> {
     }
 }
 
+/// 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 {
-            column_index: 0,
-            column_type: Any::Advice,
-            rotation: Rotation(0),
-        });
-        let a1 = VarMid::Query(QueryMid {
-            column_index: 1,
-            column_type: Any::Advice,
-            rotation: Rotation(0),
-        });
+        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!("{}", ExprDisp(&e)));
+        assert_eq!("a0 + a1", format!("{}", expr_disp(&e)));
         let e = a0.clone() + a1.clone() + a0.clone();
-        assert_eq!("a0 + a1 + a0", format!("{}", ExprDisp(&e)));
+        assert_eq!("a0 + a1 + a0", format!("{}", expr_disp(&e)));
 
         let e = a0.clone() * a1.clone();
-        assert_eq!("a0 * a1", format!("{}", ExprDisp(&e)));
+        assert_eq!("a0 * a1", format!("{}", expr_disp(&e)));
         let e = a0.clone() * a1.clone() * a0.clone();
-        assert_eq!("a0 * a1 * a0", format!("{}", ExprDisp(&e)));
+        assert_eq!("a0 * a1 * a0", format!("{}", expr_disp(&e)));
 
         let e = a0.clone() - a1.clone();
-        assert_eq!("a0 - a1", format!("{}", ExprDisp(&e)));
+        assert_eq!("a0 - a1", format!("{}", expr_disp(&e)));
         let e = (a0.clone() - a1.clone()) - a0.clone();
-        assert_eq!("a0 - a1 - a0", format!("{}", ExprDisp(&e)));
+        assert_eq!("a0 - a1 - a0", format!("{}", expr_disp(&e)));
         let e = a0.clone() - (a1.clone() - a0.clone());
-        assert_eq!("a0 - (a1 - a0)", format!("{}", ExprDisp(&e)));
+        assert_eq!("a0 - (a1 - a0)", format!("{}", expr_disp(&e)));
 
         let e = a0.clone() * a1.clone() + a0.clone();
-        assert_eq!("a0 * a1 + a0", format!("{}", ExprDisp(&e)));
+        assert_eq!("a0 * a1 + a0", format!("{}", expr_disp(&e)));
         let e = a0.clone() * (a1.clone() + a0.clone());
-        assert_eq!("a0 * (a1 + a0)", format!("{}", ExprDisp(&e)));
+        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]

halo2_frontend/src/plonk/circuit/constraint_system.rs

@@ -790,6 +790,16 @@ impl<F: Field> ConstraintSystem<F> {
 
     /// Annotate an Instance column.
     pub fn annotate_lookup_any_column<A, AR, T>(&mut self, column: T, annotation: A)
+    where
+        A: Fn() -> AR,
+        AR: Into<String>,
+        T: Into<Column<Any>>,
+    {
+        self.annotate_column(column, annotation)
+    }
+
+    /// Annotate a column.
+    pub fn annotate_column<A, AR, T>(&mut self, column: T, annotation: A)
     where
         A: Fn() -> AR,
         AR: Into<String>,

halo2_middleware/src/circuit.rs

@@ -34,6 +34,16 @@ pub struct QueryMid {
     pub rotation: Rotation,
 }
 
+impl QueryMid {
+    pub fn new(column_type: Any, column_index: usize, rotation: Rotation) -> Self {
+        Self {
+            column_index,
+            column_type,
+            rotation,
+        }
+    }
+}
+
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub enum VarMid {
     /// This is a generic column query

halo2_proofs/Cargo.toml

@@ -61,6 +61,7 @@ gumdrop = "0.8"
 proptest = "1"
 dhat = "0.3.2"
 serde_json = "1"
+halo2_debug = { path = "../halo2_debug" }
 
 [target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dev-dependencies]
 getrandom = { version = "0.2", features = ["js"] }