From 0efc52e7f6e29cbf6af9813106bd7e6df389e526 Mon Sep 17 00:00:00 2001
From: Brechtpd <Brechtp.Devos@gmail.com>
Date: Sat, 7 May 2022 19:59:14 +0200
Subject: [PATCH] Use a hash map to track cell assignment

---
 halo2_proofs/src/dev.rs | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/halo2_proofs/src/dev.rs b/halo2_proofs/src/dev.rs
index a458d3211b..ce8327974a 100644
--- a/halo2_proofs/src/dev.rs
+++ b/halo2_proofs/src/dev.rs
@@ -249,9 +249,9 @@ struct Region {
     /// The selectors that have been enabled in this region. All other selectors are by
     /// construction not enabled.
     enabled_selectors: HashMap<Selector, Vec<usize>>,
-    /// The cells assigned in this region. We store this as a `Vec` so that if any cells
-    /// are double-assigned, they will be visibly darker.
-    cells: Vec<(Column<Any>, usize)>,
+    /// The cells assigned in this region. We store this as a `HashMap` with count
+    /// so that if any cells are double-assigned, they will be visibly darker.
+    cells: HashMap<(Column<Any>, usize), usize>,
 }
 
 impl Region {
@@ -270,6 +270,16 @@ impl Region {
         }
         self.rows = Some((start, end));
     }
+
+    fn track_cell(&mut self, column: Column<Any>, row: usize) {
+        // Keep track of how many times this cell has been assigned to.
+        let count = *self.cells.get(&(column, row)).unwrap_or(&0);
+        self.cells.insert((column, row), count + 1);
+    }
+
+    fn is_assigned(&self, column: Column<Any>, row: usize) -> bool {
+        self.cells.contains_key(&(column, row))
+    }
 }
 
 /// The value of a particular cell within the circuit.
@@ -509,7 +519,7 @@ impl<F: Field + Group> Assignment<F> for MockProver<F> {
             columns: HashSet::default(),
             rows: None,
             enabled_selectors: HashMap::default(),
-            cells: vec![],
+            cells: HashMap::default(),
         });
     }
 
@@ -572,7 +582,7 @@ impl<F: Field + Group> Assignment<F> for MockProver<F> {
 
         if let Some(region) = self.current_region.as_mut() {
             region.update_extent(column.into(), row);
-            region.cells.push((column.into(), row));
+            region.track_cell(column.into(), row);
         }
 
         *self
@@ -603,7 +613,7 @@ impl<F: Field + Group> Assignment<F> for MockProver<F> {
 
         if let Some(region) = self.current_region.as_mut() {
             region.update_extent(column.into(), row);
-            region.cells.push((column.into(), row));
+            region.track_cell(column.into(), row);
         }
 
         *self
@@ -799,7 +809,7 @@ impl<F: FieldExt> MockProver<F> {
                                 let cell_row = ((gate_row + n + cell.rotation.0) % n) as usize;
 
                                 // Check that it was assigned!
-                                if r.cells.contains(&(cell.column, cell_row)) {
+                                if r.is_assigned(cell.column, cell_row) {
                                     None
                                 } else {
                                     Some(VerifyFailure::CellNotAssigned {