From 42ea7116a7b3ee8eaa5d19c77207bd5bfce552da Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@electriccoin.co>
Date: Tue, 20 Jul 2021 10:43:12 +0100
Subject: [PATCH] dev: Introduce Cell helper struct to CircuitLayout internals

---
 halo2_proofs/src/dev/cost.rs         | 33 +++++++++++++++++++++-------
 halo2_proofs/src/dev/graph/layout.rs | 24 ++++++++++----------
 2 files changed, 37 insertions(+), 20 deletions(-)

diff --git a/halo2_proofs/src/dev/cost.rs b/halo2_proofs/src/dev/cost.rs
index 0c3c7efa..b0a9171e 100644
--- a/halo2_proofs/src/dev/cost.rs
+++ b/halo2_proofs/src/dev/cost.rs
@@ -54,6 +54,12 @@ pub struct CircuitCost<G: PrimeGroup, ConcreteCircuit: Circuit<G::Scalar>> {
     _marker: PhantomData<(G, ConcreteCircuit)>,
 }
 
+#[derive(Debug)]
+pub(crate) struct Cell {
+    pub(crate) column: RegionColumn,
+    pub(crate) row: usize,
+}
+
 /// Region implementation used by Layout
 #[allow(dead_code)]
 #[derive(Debug)]
@@ -66,8 +72,9 @@ pub(crate) struct LayoutRegion {
     pub(crate) offset: Option<usize>,
     /// The number of rows that this region takes up.
     pub(crate) rows: usize,
-    /// The cells assigned in this region.
-    pub(crate) cells: Vec<(RegionColumn, usize)>,
+    /// The cells assigned in this region. We store this as a `Vec` to track multiple
+    /// assignments to a cell.
+    pub(crate) cells: Vec<Cell>,
 }
 
 /// Cost and graphing layouter
@@ -84,10 +91,11 @@ pub(crate) struct Layout {
     pub(crate) total_advice_rows: usize,
     /// Total fixed rows
     pub(crate) total_fixed_rows: usize,
-    /// Any cells assigned outside of a region.
-    pub(crate) loose_cells: Vec<(RegionColumn, usize)>,
+    /// Any cells assigned outside of a region. We store this as a `Vec` to track multiple
+    /// assignments to a cell.
+    pub(crate) loose_cells: Vec<Cell>,
     /// Pairs of cells between which we have equality constraints.
-    pub(crate) equality: Vec<(Column<Any>, usize, Column<Any>, usize)>,
+    pub(crate) equality: Vec<(Cell, Cell)>,
     /// Selector assignments used for optimization pass
     pub(crate) selectors: Vec<Vec<bool>>,
 }
@@ -139,9 +147,9 @@ impl Layout {
             region.rows = cmp::max(region.rows, row - offset + 1);
             region.offset = Some(offset);
 
-            region.cells.push((column, row));
+            region.cells.push(Cell { column, row });
         } else {
-            self.loose_cells.push((column, row));
+            self.loose_cells.push(Cell { column, row });
         }
     }
 }
@@ -228,7 +236,16 @@ impl<F: Field> Assignment<F> for Layout {
         r_col: Column<Any>,
         r_row: usize,
     ) -> Result<(), crate::plonk::Error> {
-        self.equality.push((l_col, l_row, r_col, r_row));
+        self.equality.push((
+            Cell {
+                column: l_col.into(),
+                row: l_row,
+            },
+            Cell {
+                column: r_col.into(),
+                row: r_row,
+            },
+        ));
         Ok(())
     }
 
diff --git a/halo2_proofs/src/dev/graph/layout.rs b/halo2_proofs/src/dev/graph/layout.rs
index 965e0683..b1060d92 100644
--- a/halo2_proofs/src/dev/graph/layout.rs
+++ b/halo2_proofs/src/dev/graph/layout.rs
@@ -8,7 +8,7 @@ use std::ops::Range;
 
 use crate::{
     circuit::layouter::RegionColumn,
-    dev::cost::Layout,
+    dev::cost::{Cell, Layout},
     plonk::{Any, Circuit, Column, ConstraintSystem, FloorPlanner},
 };
 
@@ -242,26 +242,26 @@ impl CircuitLayout {
 
         // Darken the cells of the region that have been assigned to.
         for region in layout.regions {
-            for (column, row) in region.cells {
+            for Cell { column, row } in region.cells {
                 draw_cell(&root, column_index(&cs, column), row)?;
             }
         }
 
         // Darken any loose cells that have been assigned to.
-        for (column, row) in layout.loose_cells {
+        for Cell { column, row } in layout.loose_cells {
             draw_cell(&root, column_index(&cs, column), row)?;
         }
 
         // Mark equality-constrained cells.
         if self.mark_equality_cells {
             let mut cells = HashSet::new();
-            for (l_col, l_row, r_col, r_row) in &layout.equality {
-                let l_col = column_index(&cs, (*l_col).into());
-                let r_col = column_index(&cs, (*r_col).into());
+            for (l, r) in &layout.equality {
+                let l_col = column_index(&cs, l.column);
+                let r_col = column_index(&cs, r.column);
 
                 // Deduplicate cells.
-                cells.insert((l_col, *l_row));
-                cells.insert((r_col, *r_row));
+                cells.insert((l_col, l.row));
+                cells.insert((r_col, r.row));
             }
 
             for (col, row) in cells {
@@ -274,11 +274,11 @@ impl CircuitLayout {
 
         // Draw lines between equality-constrained cells.
         if self.show_equality_constraints {
-            for (l_col, l_row, r_col, r_row) in &layout.equality {
-                let l_col = column_index(&cs, (*l_col).into());
-                let r_col = column_index(&cs, (*r_col).into());
+            for (l, r) in &layout.equality {
+                let l_col = column_index(&cs, l.column);
+                let r_col = column_index(&cs, r.column);
                 root.draw(&PathElement::new(
-                    [(l_col, *l_row), (r_col, *r_row)],
+                    [(l_col, l.row), (r_col, r.row)],
                     ShapeStyle::from(&RED),
                 ))?;
             }