halo2_base/virtual_region/
lookups.rs

1use std::collections::BTreeMap;
2use std::sync::{Arc, Mutex, OnceLock};
3
4use getset::{CopyGetters, Getters, Setters};
5
6use crate::ff::Field;
7use crate::halo2_proofs::{
8    circuit::{Region, Value},
9    plonk::{Advice, Column},
10};
11use crate::utils::halo2::{constrain_virtual_equals_external, raw_assign_advice};
12use crate::{AssignedValue, ContextTag};
13
14use super::copy_constraints::SharedCopyConstraintManager;
15use super::manager::VirtualRegionManager;
16
17/// Basic dynamic lookup table gadget.
18pub mod basic;
19
20/// A manager that can be used for any lookup argument. This manager automates
21/// the process of copying cells to designed advice columns with lookup enabled.
22/// It also manages how many such advice columns are necessary.
23///
24/// ## Detailed explanation
25/// If we have a lookup argument that uses `ADVICE_COLS` advice columns and `TABLE_COLS` table columns, where
26/// the table is either fixed or dynamic (advice), then we want to dynamically allocate chunks of `ADVICE_COLS` columns
27/// that have the lookup into the table **always on** so that:
28/// - every time we want to lookup [_; ADVICE_COLS] values, we copy them over to a row in the special
29///
30/// lookup-enabled advice columns.
31/// - note that just for assignment, we don't need to know anything about the table itself.
32///
33/// Note: the manager does not need to know the value of `TABLE_COLS`.
34///
35/// We want this manager to be CPU thread safe, while ensuring that the resulting circuit is
36/// deterministic -- the order in which the cells to lookup are added matters.
37/// The current solution is to tag the cells to lookup with the context id from the [`Context`](crate::Context) in which
38/// it was called, and add virtual cells sequentially to buckets labelled by id.
39/// The virtual cells will be assigned to physical cells sequentially by id.
40/// We use a `BTreeMap` for the buckets instead of sorting to cells, to ensure that the order of the cells
41/// within a bucket is deterministic.
42/// The assumption is that the [`Context`](crate::Context) is thread-local.
43///
44/// Cheap to clone across threads because everything is in [Arc].
45#[derive(Clone, Debug, Getters, CopyGetters, Setters)]
46pub struct LookupAnyManager<F: Field + Ord, const ADVICE_COLS: usize> {
47    /// Shared cells to lookup, tagged by (type id, context id).
48    #[allow(clippy::type_complexity)]
49    pub cells_to_lookup: Arc<Mutex<BTreeMap<ContextTag, Vec<[AssignedValue<F>; ADVICE_COLS]>>>>,
50    /// Global shared copy manager
51    #[getset(get = "pub", set = "pub")]
52    copy_manager: SharedCopyConstraintManager<F>,
53    /// Specify whether constraints should be imposed for additional safety.
54    #[getset(get_copy = "pub")]
55    witness_gen_only: bool,
56    /// Flag for whether `assign_raw` has been called, for safety only.
57    pub(crate) assigned: Arc<OnceLock<()>>,
58}
59
60impl<F: Field + Ord, const ADVICE_COLS: usize> LookupAnyManager<F, ADVICE_COLS> {
61    /// Creates a new [LookupAnyManager] with a given copy manager.
62    pub fn new(witness_gen_only: bool, copy_manager: SharedCopyConstraintManager<F>) -> Self {
63        Self {
64            witness_gen_only,
65            cells_to_lookup: Default::default(),
66            copy_manager,
67            assigned: Default::default(),
68        }
69    }
70
71    /// Add a lookup argument to the manager.
72    pub fn add_lookup(&self, tag: ContextTag, cells: [AssignedValue<F>; ADVICE_COLS]) {
73        self.cells_to_lookup
74            .lock()
75            .unwrap()
76            .entry(tag)
77            .and_modify(|thread| thread.push(cells))
78            .or_insert(vec![cells]);
79    }
80
81    /// The total number of virtual rows needed to special lookups
82    pub fn total_rows(&self) -> usize {
83        self.cells_to_lookup.lock().unwrap().iter().flat_map(|(_, advices)| advices).count()
84    }
85
86    /// The optimal number of `ADVICE_COLS` chunks of advice columns with lookup enabled for this
87    /// particular lookup argument that we should allocate.
88    pub fn num_advice_chunks(&self, usable_rows: usize) -> usize {
89        let total = self.total_rows();
90        total.div_ceil(usable_rows)
91    }
92
93    /// Clears state
94    pub fn clear(&mut self) {
95        self.cells_to_lookup.lock().unwrap().clear();
96        self.copy_manager.lock().unwrap().clear();
97        self.assigned = Arc::new(OnceLock::new());
98    }
99
100    /// Deep clone with the specified copy manager. Unsets `assigned`.
101    pub fn deep_clone(&self, copy_manager: SharedCopyConstraintManager<F>) -> Self {
102        Self {
103            witness_gen_only: self.witness_gen_only,
104            cells_to_lookup: Arc::new(Mutex::new(self.cells_to_lookup.lock().unwrap().clone())),
105            copy_manager,
106            assigned: Default::default(),
107        }
108    }
109}
110
111impl<F: Field + Ord, const ADVICE_COLS: usize> Drop for LookupAnyManager<F, ADVICE_COLS> {
112    /// Sanity checks whether the manager has assigned cells to lookup,
113    /// to prevent user error.
114    fn drop(&mut self) {
115        if Arc::strong_count(&self.cells_to_lookup) > 1 {
116            return;
117        }
118        if self.total_rows() > 0 && self.assigned.get().is_none() {
119            dbg!("WARNING: LookupAnyManager was not assigned!");
120        }
121    }
122}
123
124impl<F: Field + Ord, const ADVICE_COLS: usize> VirtualRegionManager<F>
125    for LookupAnyManager<F, ADVICE_COLS>
126{
127    type Config = Vec<[Column<Advice>; ADVICE_COLS]>;
128    type Assignment = ();
129
130    fn assign_raw(&self, config: &Self::Config, region: &mut Region<F>) {
131        let mut copy_manager =
132            (!self.witness_gen_only).then(|| self.copy_manager().lock().unwrap());
133        let cells_to_lookup = self.cells_to_lookup.lock().unwrap();
134        // Copy the cells to the config columns, going left to right, then top to bottom.
135        // Will panic if out of rows
136        let mut lookup_offset = 0;
137        let mut lookup_col = 0;
138        for advices in cells_to_lookup.iter().flat_map(|(_, advices)| advices) {
139            if lookup_col >= config.len() {
140                lookup_col = 0;
141                lookup_offset += 1;
142            }
143            for (advice, &column) in advices.iter().zip(config[lookup_col].iter()) {
144                let bcell =
145                    raw_assign_advice(region, column, lookup_offset, Value::known(advice.value));
146                if let Some(copy_manager) = copy_manager.as_mut() {
147                    constrain_virtual_equals_external(region, *advice, bcell.cell(), copy_manager);
148                }
149            }
150
151            lookup_col += 1;
152        }
153        // We cannot clear `cells_to_lookup` because keygen_vk and keygen_pk both call this function
154        let _ = self.assigned.set(());
155    }
156}
halo2_base/virtual_region/lookups.rs

halo2_base/virtual_region/
lookups.rs
