diff --git a/halo2_debug/src/lib.rs b/halo2_debug/src/lib.rs index e5b5f428d3..426cfa5b9d 100644 --- a/halo2_debug/src/lib.rs +++ b/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,9 +38,55 @@ impl fmt::Display for FDisp<'_, F> { } } -pub struct ExprDisp<'a, F: PrimeField, V: Variable>(pub &'a Expression); +/// 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, + /// `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 fmt::Display for ExprDisp<'_, F, V> { +fn var_fmt_default(v: &V, f: &mut fmt::Formatter<'_>, _: &()) -> fmt::Result { + write!(f, "{}", v) +} + +fn var_fmt_names( + v: &VarMid, + f: &mut fmt::Formatter<'_>, + names: &HashMap, +) -> 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, +) -> 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, + names: &'a HashMap, +) -> ExprDisp<'a, F, VarMid, HashMap> { + ExprDisp(e, var_fmt_names, names) +} + +impl fmt::Display for ExprDisp<'_, F, V, A> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let is_sum = |e: &Expression| -> bool { matches!(e, Expression::Sum(_, _)) }; let fmt_expr = @@ -45,7 +94,7 @@ impl fmt::Display for ExprDisp<'_, F, V> { 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 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 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, + /// `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, +) -> 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, + names: &'a HashMap, +) -> LookupArgDisp<'a, F, VarMid, HashMap> { + LookupArgDisp(a, var_fmt_names, names) +} + +impl 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, + /// `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, +) -> 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, + names: &'a HashMap, +) -> ShuffleArgDisp<'a, F, VarMid, HashMap> { + ShuffleArgDisp(a, var_fmt_names, names) +} + +impl 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; + 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; - 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] diff --git a/halo2_frontend/src/plonk/circuit/constraint_system.rs b/halo2_frontend/src/plonk/circuit/constraint_system.rs index 7a393a6ac0..cabe718042 100644 --- a/halo2_frontend/src/plonk/circuit/constraint_system.rs +++ b/halo2_frontend/src/plonk/circuit/constraint_system.rs @@ -790,6 +790,16 @@ impl ConstraintSystem { /// Annotate an Instance column. pub fn annotate_lookup_any_column(&mut self, column: T, annotation: A) + where + A: Fn() -> AR, + AR: Into, + T: Into>, + { + self.annotate_column(column, annotation) + } + + /// Annotate a column. + pub fn annotate_column(&mut self, column: T, annotation: A) where A: Fn() -> AR, AR: Into, diff --git a/halo2_middleware/src/circuit.rs b/halo2_middleware/src/circuit.rs index c7b0b611a9..4c8d4ee0d9 100644 --- a/halo2_middleware/src/circuit.rs +++ b/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 diff --git a/halo2_proofs/Cargo.toml b/halo2_proofs/Cargo.toml index a68f422934..4302dd7276 100644 --- a/halo2_proofs/Cargo.toml +++ b/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"] } diff --git a/halo2_proofs/tests/compress_selectors.rs b/halo2_proofs/tests/compress_selectors.rs index ec87354fc2..96a8abc629 100644 --- a/halo2_proofs/tests/compress_selectors.rs +++ b/halo2_proofs/tests/compress_selectors.rs @@ -3,6 +3,8 @@ use std::marker::PhantomData; use ff::PrimeField; +use halo2_debug::expr_disp_names; +use halo2_frontend::circuit::compile_circuit; use halo2_frontend::plonk::Error; use halo2_proofs::circuit::{Cell, Layouter, SimpleFloorPlanner, Value}; use halo2_proofs::poly::Rotation; @@ -10,6 +12,7 @@ use halo2_proofs::poly::Rotation; use halo2_backend::transcript::{ Blake2bRead, Blake2bWrite, Challenge255, TranscriptReadBuffer, TranscriptWriterBuffer, }; +use halo2_middleware::circuit::{Any, ColumnMid}; use halo2_middleware::zal::impls::{H2cEngine, PlonkEngineConfig}; use halo2_proofs::arithmetic::Field; use halo2_proofs::plonk::{ @@ -101,12 +104,16 @@ impl MyCircuitChip { let l = meta.advice_column(); let r = meta.advice_column(); let o = meta.advice_column(); + meta.annotate_column(l, || "l"); + meta.annotate_column(r, || "r"); + meta.annotate_column(o, || "o"); let s_add = meta.selector(); let s_mul = meta.selector(); let s_cubed = meta.selector(); let PI = meta.instance_column(); + meta.annotate_column(PI, || "pi"); meta.enable_equality(l); meta.enable_equality(r); @@ -435,6 +442,63 @@ How the `compress_selectors` works in `MyCircuit` under the hood: */ +#[test] +fn test_compress_gates() { + let k = 4; + let circuit: MyCircuit = MyCircuit { + x: Value::known(Fr::one()), + y: Value::known(Fr::one()), + constant: Fr::one(), + }; + + // Without compression + + let (mut compress_false, _, _) = compile_circuit(k, &circuit, false).unwrap(); + + let names = &mut compress_false.cs.general_column_annotations; + names.insert(ColumnMid::new(Any::Fixed, 0), "s_add".to_string()); + names.insert(ColumnMid::new(Any::Fixed, 1), "s_mul".to_string()); + names.insert(ColumnMid::new(Any::Fixed, 2), "s_cubed".to_string()); + let cs = &compress_false.cs; + let names = &cs.general_column_annotations; + assert_eq!(3, cs.gates.len()); + assert_eq!( + "s_add * (l + r - o)", + format!("{}", expr_disp_names(&cs.gates[0].poly, names)) + ); + assert_eq!( + "s_mul * (l * r - o)", + format!("{}", expr_disp_names(&cs.gates[1].poly, names)) + ); + assert_eq!( + "s_cubed * (l * l * l - o)", + format!("{}", expr_disp_names(&cs.gates[2].poly, names)) + ); + + // With compression + + let (mut compress_true, _, _) = compile_circuit(k, &circuit, true).unwrap(); + + let names = &mut compress_true.cs.general_column_annotations; + names.insert(ColumnMid::new(Any::Fixed, 0), "s_add_mul".to_string()); + names.insert(ColumnMid::new(Any::Fixed, 1), "s_cubed".to_string()); + let cs = &compress_true.cs; + let names = &cs.general_column_annotations; + assert_eq!(3, cs.gates.len()); + assert_eq!( + "s_add_mul * (2 - s_add_mul) * (l + r - o)", + format!("{}", expr_disp_names(&cs.gates[0].poly, names)) + ); + assert_eq!( + "s_add_mul * (1 - s_add_mul) * (l * r - o)", + format!("{}", expr_disp_names(&cs.gates[1].poly, names)) + ); + assert_eq!( + "s_cubed * (l * l * l - o)", + format!("{}", expr_disp_names(&cs.gates[2].poly, names)) + ); +} + #[test] fn test_success() { // vk & pk keygen both WITH compress