[fix] BasicDynLookupConfig needs selector on advice table to prevent lookup poison

halo2-base/src/utils/halo2.rs

@@ -3,6 +3,8 @@ use crate::halo2_proofs::{
     circuit::{AssignedCell, Cell, Region, Value},
     plonk::{Advice, Assigned, Column, Fixed},
 };
+use crate::virtual_region::copy_constraints::{CopyConstraintManager, SharedCopyConstraintManager};
+use crate::AssignedValue;
 
 /// Raw (physical) assigned cell in Plonkish arithmetization.
 #[cfg(feature = "halo2-axiom")]
@@ -71,3 +73,43 @@ pub fn raw_constrain_equal<F: Field>(region: &mut Region<F>, left: Cell, right:
     #[cfg(not(feature = "halo2-axiom"))]
     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`.
+///
+/// This should only be called when `witness_gen_only` is false, otherwise it will panic.
+///
+/// ## Panics
+/// If witness generation only mode is true.
+pub fn constrain_virtual_equals_external<F: Field + Ord>(
+    region: &mut Region<F>,
+    virtual_cell: AssignedValue<F>,
+    external_cell: Cell,
+    copy_manager: &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);
+}

halo2-base/src/virtual_region/lookups.rs

@@ -8,7 +8,7 @@ use crate::halo2_proofs::{
     circuit::{Region, Value},
     plonk::{Advice, Column},
 };
-use crate::utils::halo2::raw_assign_advice;
+use crate::utils::halo2::{constrain_virtual_equals_external, raw_assign_advice};
 use crate::{AssignedValue, ContextTag};
 
 use super::copy_constraints::SharedCopyConstraintManager;
@@ -125,6 +125,7 @@ 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 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
@@ -138,12 +139,8 @@ 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 !self.witness_gen_only {
-                    let ctx_cell = advice.cell.unwrap();
-                    let copy_manager = self.copy_manager.lock().unwrap();
-                    let acell =
-                        copy_manager.assigned_advices.get(&ctx_cell).expect("cell not assigned");
-                    region.constrain_equal(*acell, bcell.cell());
+                if let Some(copy_manager) = copy_manager.as_ref() {
+                    constrain_virtual_equals_external(region, *advice, bcell.cell(), copy_manager);
                 }
             }
 

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

@@ -1,18 +1,20 @@
+use std::iter::zip;
+
 use crate::{
     halo2_proofs::{
         circuit::{Layouter, Region, Value},
         halo2curves::ff::Field,
-        plonk::{Advice, Column, ConstraintSystem, Phase},
+        plonk::{Advice, Column, ConstraintSystem, Fixed, Phase},
         poly::Rotation,
     },
     utils::{
-        halo2::{raw_assign_advice, Halo2AssignedCell},
+        halo2::{
+            assign_virtual_to_raw, constrain_virtual_equals_external, raw_assign_advice,
+            raw_assign_fixed,
+        },
         ScalarField,
     },
-    virtual_region::{
-        copy_constraints::SharedCopyConstraintManager, lookups::LookupAnyManager,
-        manager::VirtualRegionManager,
-    },
+    virtual_region::copy_constraints::SharedCopyConstraintManager,
     AssignedValue,
 };
 
@@ -24,12 +26,27 @@ use crate::{
 ///
 /// We can have multiple sets of dedicated columns to be looked up: these can be specified
 /// when calling `new`, but typically we just need 1 set.
+///
+/// The `table` consists of advice columns. Since this table may have poisoned rows (blinding factors),
+/// we use a fixed column `table_selector` which is default 0 and only 1 on enabled rows of the table.
+/// The dynamic lookup will check that for `(key, key_is_enabled)` in `to_lookup` we have `key` matches one of
+/// the rows in `table` where `table_selector == key_is_enabled`.
+/// Reminder: the Halo2 lookup argument will ignore the poisoned rows in `to_lookup`
+/// (see [https://zcash.github.io/halo2/design/proving-system/lookup.html#zero-knowledge-adjustment]), but it will
+/// not ignore the poisoned rows in `table`.
+///
+/// Part of this design consideration is to allow a key of `[F::ZERO; KEY_COL]` to still be used as a valid key
+/// in the lookup argument. By default, unfilled rows in `to_lookup` will be all zeros; we require
+/// at least one row in `table` where `table_is_enabled = 0` and the rest of the row in `table` are also 0s.
 #[derive(Clone, Debug)]
 pub struct BasicDynLookupConfig<const KEY_COL: usize> {
-    /// Columns for cells to be looked up.
-    pub to_lookup: Vec<[Column<Advice>; KEY_COL]>,
+    /// Columns for cells to be looked up. Consists of `(key, key_is_enabled)`.
+    pub to_lookup: Vec<([Column<Advice>; KEY_COL], Column<Fixed>)>,
     /// Table to look up against.
     pub table: [Column<Advice>; KEY_COL],
+    /// Selector to enable a row in `table` to actually be part of the lookup table. This is to prevent
+    /// blinding factors in `table` advice columns from being used in the lookup.
+    pub table_is_enabled: Column<Fixed>,
 }
 
 impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
@@ -42,44 +59,95 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
         num_lu_sets: usize,
     ) -> Self {
         let mut make_columns = || {
-            [(); KEY_COL].map(|_| {
+            let advices = [(); KEY_COL].map(|_| {
                 let advice = meta.advice_column_in(phase());
                 meta.enable_equality(advice);
                 advice
-            })
+            });
+            let is_enabled = meta.fixed_column();
+            (advices, is_enabled)
         };
-        let table = make_columns();
+        let (table, table_is_enabled) = make_columns();
         let to_lookup: Vec<_> = (0..num_lu_sets).map(|_| make_columns()).collect();
 
-        for to_lookup in &to_lookup {
+        for (key, key_is_enabled) in &to_lookup {
             meta.lookup_any("dynamic lookup table", |meta| {
                 let table = table.map(|c| meta.query_advice(c, Rotation::cur()));
-                let to_lu = to_lookup.map(|c| meta.query_advice(c, Rotation::cur()));
-                to_lu.into_iter().zip(table).collect()
+                let table_is_enabled = meta.query_fixed(table_is_enabled, Rotation::cur());
+                let key = key.map(|c| meta.query_advice(c, Rotation::cur()));
+                let key_is_enabled = meta.query_fixed(*key_is_enabled, Rotation::cur());
+                zip(key, table).chain([(key_is_enabled, table_is_enabled)]).collect()
             });
         }
 
-        Self { table, to_lookup }
+        Self { table_is_enabled, table, to_lookup }
     }
 
     /// Assign managed lookups
-    pub fn assign_managed_lookups<F: ScalarField>(
+    ///
+    /// `copy_manager` **must** be provided unless you are only doing witness generation
+    /// without constraints.
+    pub fn assign_virtual_to_lookup_to_raw<F: ScalarField>(
         &self,
         mut layouter: impl Layouter<F>,
-        lookup_manager: &LookupAnyManager<F, KEY_COL>,
+        keys: impl IntoIterator<Item = [AssignedValue<F>; KEY_COL]>,
+        copy_manager: Option<&SharedCopyConstraintManager<F>>,
     ) {
+        #[cfg(not(feature = "halo2-axiom"))]
+        let keys = keys.into_iter().collect::<Vec<_>>();
         layouter
             .assign_region(
-                || "Managed lookup advice",
+                || "[BasicDynLookupConfig] Advice cells to lookup",
                 |mut region| {
-                    lookup_manager.assign_raw(&self.to_lookup, &mut region);
+                    self.assign_virtual_to_lookup_to_raw_from_offset(
+                        &mut region,
+                        keys,
+                        0,
+                        copy_manager,
+                    );
                     Ok(())
                 },
             )
             .unwrap();
     }
 
-    /// Assign virtual table to raw
+    /// `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>(
+        &self,
+        region: &mut Region<F>,
+        keys: impl IntoIterator<Item = [AssignedValue<F>; KEY_COL]>,
+        mut offset: usize,
+        copy_manager: Option<&SharedCopyConstraintManager<F>>,
+    ) {
+        let 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
+        let mut lookup_col = 0;
+        for key in keys {
+            if lookup_col >= self.to_lookup.len() {
+                lookup_col = 0;
+                offset += 1;
+            }
+            let (key_col, key_is_enabled_col) = self.to_lookup[lookup_col];
+            // set key_is_enabled to 1
+            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() {
+                    constrain_virtual_equals_external(region, advice, bcell.cell(), copy_manager);
+                }
+            }
+
+            lookup_col += 1;
+        }
+    }
+
+    /// Assign virtual table to raw.
+    ///
+    /// `copy_manager` **must** be provided unless you are only doing witness generation
+    /// without constraints.
     pub fn assign_virtual_table_to_raw<F: ScalarField>(
         &self,
         mut layouter: impl Layouter<F>,
@@ -90,7 +158,7 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
         let rows = rows.into_iter().collect::<Vec<_>>();
         layouter
             .assign_region(
-                || "Dynamic Lookup Table",
+                || "[BasicDynLookupConfig] Dynamic Lookup Table",
                 |mut region| {
                     self.assign_virtual_table_to_raw_from_offset(
                         &mut region,
@@ -113,33 +181,21 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
         &self,
         region: &mut Region<F>,
         rows: impl IntoIterator<Item = [AssignedValue<F>; KEY_COL]>,
-        offset: usize,
+        mut offset: usize,
         copy_manager: Option<&SharedCopyConstraintManager<F>>,
     ) {
-        for (i, row) in rows.into_iter().enumerate() {
+        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 + i, virtual_cell, copy_manager);
+                assign_virtual_to_raw(region, col, offset, virtual_cell, copy_manager);
             }
+            offset += 1;
+        }
+        // always assign one disabled row with all 0s, so disabled to_lookup works for sure
+        raw_assign_fixed(region, self.table_is_enabled, offset, F::ZERO);
+        for col in self.table {
+            raw_assign_advice(region, col, offset, Value::known(F::ZERO));
         }
     }
 }
-
-/// Assign virtual cell to raw halo2 cell.
-/// `copy_manager` **must** be provided unless you are only doing witness generation
-/// without constraints.
-pub fn assign_virtual_to_raw<'v, F: ScalarField>(
-    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
-}