Use a hash map to track cell assignment (#71)

privacy-scaling-explorations · May 9, 2022 · bffb83e · bffb83e
1 parent 9992785
commit bffb83e
Showing 1 changed file with 17 additions and 7 deletions.
diff --git 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 {