feat: add pretty debug information for common items (#111)

Resolve #85 - [x] TraceWitness - [x] Assignments --------- Co-authored-by: Leo Lara <[email protected]>
privacy-scaling-explorations · Oct 14, 2023 · 9e1afe8 · 9e1afe8
1 parent 2a0a8f0
commit 9e1afe8
Show file tree

Hide file tree

Showing 4 changed files with 126 additions and 7 deletions.
diff --git a/src/plonkish/backend/halo2.rs b/src/plonkish/backend/halo2.rs
@@ -188,7 +188,7 @@ impl<F: Field + From<u64> + Hash> ChiquitoHalo2<F> {
     }
 
     fn assign_advice(&self, region: &mut Region<F>, witness: &Assignments<F>) -> Result<(), Error> {
-        for (column, assignments) in witness {
+        for (column, assignments) in witness.iter() {
             let column = self.convert_advice_column(column);
 
             for (offset, value) in assignments.iter().enumerate() {
@@ -200,7 +200,7 @@ impl<F: Field + From<u64> + Hash> ChiquitoHalo2<F> {
     }
 
     fn assign_fixed(&self, region: &mut Region<F>, fixed: &Assignments<F>) -> Result<(), Error> {
-        for (column, values) in fixed {
+        for (column, values) in fixed.iter() {
             let column = self.convert_fixed_column(column);
 
             for (offset, value) in values.iter().enumerate() {

diff --git a/src/plonkish/backend/plaf.rs b/src/plonkish/backend/plaf.rs
@@ -116,7 +116,7 @@ impl<F: PrimeField<Repr = [u8; 32]>> ChiquitoPlaf<F> {
             fixed.push(vec![None; plaf.info.num_rows]);
         }
 
-        for (column, values) in self.circuit.fixed_assignments.clone().into_iter() {
+        for (column, values) in self.circuit.fixed_assignments.clone().0.into_iter() {
             let column = self
                 .c_column_id_to_p_column_index
                 .get(&column.uuid())
@@ -306,7 +306,7 @@ impl ChiquitoPlafWitGen {
         };
 
         if let Some(witness) = &witness {
-            for (column, assignments) in witness {
+            for (column, assignments) in witness.iter() {
                 let p_column_index = self
                     .c_column_id_to_p_column_index
                     .get(&column.uuid())

diff --git a/src/plonkish/ir/assignments.rs b/src/plonkish/ir/assignments.rs
@@ -1,4 +1,8 @@
-use std::collections::HashMap;
+use std::{
+    collections::HashMap,
+    fmt,
+    ops::{Deref, DerefMut},
+};
 
 use crate::field::Field;
 
@@ -13,7 +17,51 @@ use crate::{
 
 use super::{Column, PolyExpr};
 
-pub type Assignments<F> = HashMap<Column, Vec<F>>;
+#[derive(Debug, Clone)]
+pub struct Assignments<F>(pub HashMap<Column, Vec<F>>);
+
+impl<F: fmt::Debug> fmt::Display for Assignments<F> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // get the decimal width based on the step_instances size, add extra one leading zero
+        let decimal_width = self.0.len().checked_ilog10().unwrap_or(0) + 2;
+        // offset(col_uuid): value0, value1, value2,...
+        for (i, (col, vals)) in self.0.iter().enumerate() {
+            let vals = vals.iter().fold(String::new(), |mut acc, val| {
+                acc.push_str(&format!("{:?}, ", val));
+                acc
+            });
+            writeln!(
+                f,
+                "{:0>width$}({}): {}",
+                i,
+                col.id,
+                vals,
+                width = decimal_width as usize,
+            )?;
+        }
+        Ok(())
+    }
+}
+
+impl<F> Default for Assignments<F> {
+    fn default() -> Self {
+        Self(HashMap::default())
+    }
+}
+
+impl<F> Deref for Assignments<F> {
+    type Target = HashMap<Column, Vec<F>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<F> DerefMut for Assignments<F> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
 
 pub struct AssignmentGenerator<F, TraceArgs> {
     columns: Vec<Column>,
@@ -227,3 +275,20 @@ impl<F: Field, TraceArgs> AssignmentGenerator<F, TraceArgs> {
         None
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn pretty_print_assignments() {
+        let display = format!(
+            "{}",
+            Assignments::<i32>(HashMap::from([
+                (Column::advice("a", 1), vec![1, 2, 3]),
+                (Column::fixed("a"), vec![4, 5, 6]),
+            ])),
+        );
+        println!("{}", display);
+    }
+}
diff --git a/src/wit_gen.rs b/src/wit_gen.rs
@@ -1,4 +1,4 @@
-use std::{collections::HashMap, hash::Hash, rc::Rc};
+use std::{collections::HashMap, fmt, hash::Hash, rc::Rc};
 
 use crate::{
     ast::{query::Queriable, StepTypeUUID},
@@ -14,6 +14,16 @@ pub struct StepInstance<F> {
     pub assignments: HashMap<Queriable<F>, F>,
 }
 
+impl<F: fmt::Debug> fmt::Display for StepInstance<F> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "({}): ", self.step_type_uuid)?;
+        for (queriable, value) in self.assignments.iter() {
+            write!(f, "{:?} = {:?}, ", queriable, value)?;
+        }
+        Ok(())
+    }
+}
+
 impl<F> StepInstance<F> {
     pub fn new(step_type_uuid: StepTypeUUID) -> StepInstance<F> {
         StepInstance {
@@ -38,6 +48,24 @@ pub struct TraceWitness<F> {
     pub step_instances: Witness<F>,
 }
 
+impl<F: fmt::Debug> fmt::Display for TraceWitness<F> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // get the decimal width based on the step_instances size, add extra one leading zero
+        let decimal_width = self.step_instances.len().checked_ilog10().unwrap_or(0) + 2;
+        // offset(step_uuid): assignations
+        for (i, step_instance) in self.step_instances.iter().enumerate() {
+            writeln!(
+                f,
+                "{:0>width$}{}",
+                i,
+                step_instance,
+                width = decimal_width as usize,
+            )?;
+        }
+        Ok(())
+    }
+}
+
 #[derive(Debug)]
 pub struct TraceContext<F> {
     witness: TraceWitness<F>,
@@ -200,6 +228,32 @@ mod tests {
         assert_eq!(ctx.witness.step_instances.len(), 5);
     }
 
+    #[test]
+    fn test_trace_witness_display() {
+        let display = format!(
+            "{}",
+            TraceWitness::<i32> {
+                step_instances: vec![
+                    StepInstance {
+                        step_type_uuid: 9,
+                        assignments: HashMap::from([
+                            (Queriable::Fixed(FixedSignal::new("a".into()), 0), 1),
+                            (Queriable::Fixed(FixedSignal::new("b".into()), 0), 2)
+                        ]),
+                    },
+                    StepInstance {
+                        step_type_uuid: 10,
+                        assignments: HashMap::from([
+                            (Queriable::Fixed(FixedSignal::new("a".into()), 0), 1),
+                            (Queriable::Fixed(FixedSignal::new("b".into()), 0), 2)
+                        ]),
+                    }
+                ]
+            }
+        );
+        println!("{}", display);
+    }
+
     #[test]
     fn test_fixed_gen_context() {
         let mut ctx = FixedGenContext::new(3);