[fix] soundness bug in `BasicDynLookupConfig::assign_virtual_table_to…

…_raw` (#224)

* fix: `BasicDynLookupConfig::assign_virtual_table_to_raw`

* feat: add safety check on all assigned_advice HashMap insertions

halo2-base/src/gates/flex_gate/threads/single_phase.rs

@@ -219,9 +219,15 @@ pub fn assign_with_constraints<F: ScalarField, const ROTATIONS: usize>(
                 .assign_advice(|| "", column, row_offset, || value.map(|v| *v))
                 .unwrap()
                 .cell();
-            copy_manager
+            if let Some(old_cell) = copy_manager
                 .assigned_advices
-                .insert(ContextCell::new(ctx.type_id, ctx.context_id, i), cell);
+                .insert(ContextCell::new(ctx.type_id, ctx.context_id, i), cell)
+            {
+                assert!(
+                    old_cell.row_offset == cell.row_offset && old_cell.column == cell.column,
+                    "Trying to overwrite virtual cell with a different raw cell"
+                );
+            }
 
             // If selector enabled and row_offset is valid add break point, account for break point overlap, and enforce equality constraint for gate outputs.
             // ⚠️ This assumes overlap is of form: gate enabled at `i - delta` and `i`, where `delta = ROTATIONS - 1`. We currently do not support `delta < ROTATIONS - 1`.

halo2-base/src/utils/halo2.rs

@@ -1,9 +1,11 @@
+use std::collections::hash_map::Entry;
+
 use crate::ff::Field;
 use crate::halo2_proofs::{
     circuit::{AssignedCell, Cell, Region, Value},
     plonk::{Advice, Assigned, Column, Fixed},
 };
-use crate::virtual_region::copy_constraints::{CopyConstraintManager, SharedCopyConstraintManager};
+use crate::virtual_region::copy_constraints::{CopyConstraintManager, EXTERNAL_CELL_TYPE_ID};
 use crate::AssignedValue;
 
 /// Raw (physical) assigned cell in Plonkish arithmetization.
@@ -74,30 +76,11 @@ pub fn raw_constrain_equal<F: Field>(region: &mut Region<F>, left: Cell, right:
     region.constrain_equal(left, right).unwrap();
 }
 
-/// Assign virtual cell to raw halo2 cell in column `column` at row offset `offset` within the `region`.
-/// Stores the mapping between `virtual_cell` and the raw assigned cell in `copy_manager`, if provided.
-///
-/// `copy_manager` **must** be provided unless you are only doing witness generation
-/// without constraints.
-pub fn assign_virtual_to_raw<'v, F: Field + Ord>(
-    region: &mut Region<F>,
-    column: Column<Advice>,
-    offset: usize,
-    virtual_cell: AssignedValue<F>,
-    copy_manager: Option<&SharedCopyConstraintManager<F>>,
-) -> Halo2AssignedCell<'v, F> {
-    let raw = raw_assign_advice(region, column, offset, Value::known(virtual_cell.value));
-    if let Some(copy_manager) = copy_manager {
-        let mut copy_manager = copy_manager.lock().unwrap();
-        let cell = virtual_cell.cell.unwrap();
-        copy_manager.assigned_advices.insert(cell, raw.cell());
-        drop(copy_manager);
-    }
-    raw
-}
-
-/// Constrains that `virtual` is equal to `external`. The `virtual` cell must have
-/// **already** been raw assigned, with the raw assigned cell stored in `copy_manager`.
+/// Constrains that `virtual_cell` is equal to `external_cell`. The `virtual_cell` must have
+/// already been raw assigned with the raw assigned cell stored in `copy_manager`
+/// **unless** it is marked an external-only cell with type id [EXTERNAL_CELL_TYPE_ID].
+/// * When the virtual cell has already been assigned, the assigned cell is constrained to be equal to the external cell.
+/// * When the virtual cell has not been assigned **and** it is marked as an external cell, it is assigned to `external_cell` and the mapping is stored in `copy_manager`.
 ///
 /// This should only be called when `witness_gen_only` is false, otherwise it will panic.
 ///
@@ -107,9 +90,19 @@ pub fn constrain_virtual_equals_external<F: Field + Ord>(
     region: &mut Region<F>,
     virtual_cell: AssignedValue<F>,
     external_cell: Cell,
-    copy_manager: &CopyConstraintManager<F>,
+    copy_manager: &mut CopyConstraintManager<F>,
 ) {
     let ctx_cell = virtual_cell.cell.unwrap();
-    let acell = copy_manager.assigned_advices.get(&ctx_cell).expect("cell not assigned");
-    region.constrain_equal(*acell, external_cell);
+    match copy_manager.assigned_advices.entry(ctx_cell) {
+        Entry::Occupied(acell) => {
+            // The virtual cell has already been assigned, so we can constrain it to equal the external cell.
+            region.constrain_equal(*acell.get(), external_cell);
+        }
+        Entry::Vacant(assigned) => {
+            // The virtual cell **must** be an external cell
+            assert_eq!(ctx_cell.type_id, EXTERNAL_CELL_TYPE_ID);
+            // We map the virtual cell to point to the raw external cell in `copy_manager`
+            assigned.insert(external_cell);
+        }
+    }
 }

halo2-base/src/virtual_region/copy_constraints.rs

@@ -15,6 +15,9 @@ use crate::{ff::Field, ContextCell};
 
 use super::manager::VirtualRegionManager;
 
+/// Type ID to distinguish external raw Halo2 cells. **This Type ID must be unique.**
+pub const EXTERNAL_CELL_TYPE_ID: &str = "halo2-base:External Raw Halo2 Cell";
+
 /// Thread-safe shared global manager for all copy constraints.
 pub type SharedCopyConstraintManager<F> = Arc<Mutex<CopyConstraintManager<F>>>;
 
@@ -86,11 +89,15 @@ impl<F: Field + Ord> CopyConstraintManager<F> {
     }
 
     fn load_external_cell_impl(&mut self, cell: Option<Cell>) -> ContextCell {
-        let context_cell =
-            ContextCell::new("halo2-base:External Raw Halo2 Cell", 0, self.external_cell_count);
+        let context_cell = ContextCell::new(EXTERNAL_CELL_TYPE_ID, 0, self.external_cell_count);
         self.external_cell_count += 1;
         if let Some(cell) = cell {
-            self.assigned_advices.insert(context_cell, cell);
+            if let Some(old_cell) = self.assigned_advices.insert(context_cell, cell) {
+                assert!(
+                    old_cell.row_offset == cell.row_offset && old_cell.column == cell.column,
+                    "External cell already assigned"
+                )
+            }
         }
         context_cell
     }

halo2-base/src/virtual_region/lookups.rs

@@ -125,7 +125,8 @@ impl<F: Field + Ord, const ADVICE_COLS: usize> VirtualRegionManager<F>
     type Config = Vec<[Column<Advice>; ADVICE_COLS]>;
 
     fn assign_raw(&self, config: &Self::Config, region: &mut Region<F>) {
-        let copy_manager = (!self.witness_gen_only).then(|| self.copy_manager().lock().unwrap());
+        let mut copy_manager =
+            (!self.witness_gen_only).then(|| self.copy_manager().lock().unwrap());
         let cells_to_lookup = self.cells_to_lookup.lock().unwrap();
         // Copy the cells to the config columns, going left to right, then top to bottom.
         // Will panic if out of rows
@@ -139,7 +140,7 @@ impl<F: Field + Ord, const ADVICE_COLS: usize> VirtualRegionManager<F>
             for (advice, &column) in advices.iter().zip(config[lookup_col].iter()) {
                 let bcell =
                     raw_assign_advice(region, column, lookup_offset, Value::known(advice.value));
-                if let Some(copy_manager) = copy_manager.as_ref() {
+                if let Some(copy_manager) = copy_manager.as_mut() {
                     constrain_virtual_equals_external(region, *advice, bcell.cell(), copy_manager);
                 }
             }

halo2-base/src/virtual_region/lookups/basic.rs

@@ -8,10 +8,7 @@ use crate::{
         poly::Rotation,
     },
     utils::{
-        halo2::{
-            assign_virtual_to_raw, constrain_virtual_equals_external, raw_assign_advice,
-            raw_assign_fixed,
-        },
+        halo2::{constrain_virtual_equals_external, raw_assign_advice, raw_assign_fixed},
         ScalarField,
     },
     virtual_region::copy_constraints::SharedCopyConstraintManager,
@@ -83,7 +80,7 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
         Self { table_is_enabled, table, to_lookup }
     }
 
-    /// Assign managed lookups
+    /// Assign managed lookups. The `keys` must have already been raw assigned beforehand.
     ///
     /// `copy_manager` **must** be provided unless you are only doing witness generation
     /// without constraints.
@@ -114,6 +111,8 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
             .unwrap();
     }
 
+    /// Assign managed lookups. The `keys` must have already been raw assigned beforehand.
+    ///
     /// `copy_manager` **must** be provided unless you are only doing witness generation
     /// without constraints.
     pub fn assign_virtual_to_lookup_to_raw_from_offset<F: ScalarField>(
@@ -123,7 +122,7 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
         mut offset: usize,
         copy_manager: Option<&SharedCopyConstraintManager<F>>,
     ) {
-        let copy_manager = copy_manager.map(|c| c.lock().unwrap());
+        let mut copy_manager = copy_manager.map(|c| c.lock().unwrap());
         // Copied from `LookupAnyManager::assign_raw` but modified to set `key_is_enabled` to 1.
         // Copy the cells to the config columns, going left to right, then top to bottom.
         // Will panic if out of rows
@@ -138,7 +137,7 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
             raw_assign_fixed(region, key_is_enabled_col, offset, F::ONE);
             for (advice, column) in zip(key, key_col) {
                 let bcell = raw_assign_advice(region, column, offset, Value::known(advice.value));
-                if let Some(copy_manager) = copy_manager.as_ref() {
+                if let Some(copy_manager) = copy_manager.as_mut() {
                     constrain_virtual_equals_external(region, advice, bcell.cell(), copy_manager);
                 }
             }
@@ -147,7 +146,7 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
         }
     }
 
-    /// Assign virtual table to raw.
+    /// Assign virtual table to raw. The `rows` must have already been raw assigned beforehand.
     ///
     /// `copy_manager` **must** be provided unless you are only doing witness generation
     /// without constraints.
@@ -178,6 +177,8 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
             .unwrap();
     }
 
+    /// Assign virtual table to raw. The `rows` must have already been raw assigned beforehand.
+    ///
     /// `copy_manager` **must** be provided unless you are only doing witness generation
     /// without constraints.
     pub fn assign_virtual_table_to_raw_from_offset<F: ScalarField>(
@@ -187,11 +188,15 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
         mut offset: usize,
         copy_manager: Option<&SharedCopyConstraintManager<F>>,
     ) {
+        let mut copy_manager = copy_manager.map(|c| c.lock().unwrap());
         for row in rows {
             // Enable this row in the table
             raw_assign_fixed(region, self.table_is_enabled, offset, F::ONE);
-            for (col, virtual_cell) in self.table.into_iter().zip(row) {
-                assign_virtual_to_raw(region, col, offset, virtual_cell, copy_manager);
+            for (advice, column) in zip(row, self.table) {
+                let bcell = raw_assign_advice(region, column, offset, Value::known(advice.value));
+                if let Some(copy_manager) = copy_manager.as_mut() {
+                    constrain_virtual_equals_external(region, advice, bcell.cell(), copy_manager);
+                }
             }
             offset += 1;
         }

halo2-base/src/virtual_region/tests/lookups/memory.rs

@@ -1,5 +1,3 @@
-use std::any::TypeId;
-
 use crate::{
     halo2_proofs::{
         arithmetic::Field,
@@ -8,7 +6,9 @@ use crate::{
         halo2curves::bn256::Fr,
         plonk::{keygen_pk, keygen_vk, Assigned, Circuit, ConstraintSystem, Error},
     },
-    virtual_region::lookups::basic::BasicDynLookupConfig,
+    virtual_region::{
+        copy_constraints::EXTERNAL_CELL_TYPE_ID, lookups::basic::BasicDynLookupConfig,
+    },
     AssignedValue, ContextCell,
 };
 use halo2_proofs_axiom::plonk::FirstPhase;
@@ -109,34 +109,38 @@ impl<F: ScalarField, const CYCLES: usize> Circuit<F> for RAMCircuit<F, CYCLES> {
         config: Self::Config,
         mut layouter: impl Layouter<F>,
     ) -> Result<(), Error> {
+        layouter.assign_region(
+            || "cpu",
+            |mut region| {
+                self.cpu.assign_raw(
+                    &(config.cpu.basic_gates[0].clone(), config.cpu.max_rows),
+                    &mut region,
+                );
+                Ok(())
+            },
+        )?;
+
+        let copy_manager = (!self.cpu.witness_gen_only()).then_some(&self.cpu.copy_manager);
+
         // Make purely virtual cells so we can raw assign them
         let memory = self.memory.iter().enumerate().map(|(i, value)| {
             let idx = Assigned::Trivial(F::from(i as u64));
-            let idx =
-                AssignedValue { value: idx, cell: Some(ContextCell::new("RAM Config", 0, i)) };
+            let idx = AssignedValue {
+                value: idx,
+                cell: Some(ContextCell::new(EXTERNAL_CELL_TYPE_ID, 0, i)),
+            };
             let value = Assigned::Trivial(*value);
-            let value = AssignedValue { value, cell: Some(ContextCell::new("RAM Config", 1, i)) };
+            let value =
+                AssignedValue { value, cell: Some(ContextCell::new(EXTERNAL_CELL_TYPE_ID, 1, i)) };
             [idx, value]
         });
 
-        let copy_manager = (!self.cpu.witness_gen_only()).then_some(&self.cpu.copy_manager);
-
         config.memory.assign_virtual_table_to_raw(
             layouter.namespace(|| "memory"),
             memory,
             copy_manager,
         );
 
-        layouter.assign_region(
-            || "cpu",
-            |mut region| {
-                self.cpu.assign_raw(
-                    &(config.cpu.basic_gates[0].clone(), config.cpu.max_rows),
-                    &mut region,
-                );
-                Ok(())
-            },
-        )?;
         config.memory.assign_virtual_to_lookup_to_raw(
             layouter.namespace(|| "memory accesses"),
             self.mem_access.clone(),