halo2_base/virtual_region/
copy_constraints.rs

1use std::collections::{BTreeMap, HashMap};
2use std::ops::DerefMut;
3use std::sync::{Arc, Mutex, OnceLock};
4
5use itertools::Itertools;
6use rayon::slice::ParallelSliceMut;
7
8use crate::halo2_proofs::{
9    circuit::{Cell, Region},
10    plonk::{Assigned, Column, Fixed},
11};
12use crate::utils::halo2::{raw_assign_fixed, raw_constrain_equal, Halo2AssignedCell};
13use crate::AssignedValue;
14use crate::{ff::Field, ContextCell};
15
16use super::manager::VirtualRegionManager;
17
18/// Type ID to distinguish external raw Halo2 cells. **This Type ID must be unique.**
19pub const EXTERNAL_CELL_TYPE_ID: &str = "halo2-base:External Raw Halo2 Cell";
20
21/// Thread-safe shared global manager for all copy constraints.
22pub type SharedCopyConstraintManager<F> = Arc<Mutex<CopyConstraintManager<F>>>;
23
24/// Global manager for all copy constraints. Thread-safe.
25///
26/// This will only be accessed during key generation, not proof generation, so it does not need to be optimized.
27///
28/// Implements [VirtualRegionManager], which should be assigned only after all cells have been assigned
29/// by other managers.
30#[derive(Clone, Default, Debug)]
31pub struct CopyConstraintManager<F: Field + Ord> {
32    /// A [Vec] tracking equality constraints between pairs of virtual advice cells, tagged by [ContextCell].
33    /// These can be across different virtual regions.
34    pub advice_equalities: Vec<(ContextCell, ContextCell)>,
35
36    /// A [Vec] tracking equality constraints between virtual advice cell and fixed values.
37    /// Fixed values will only be added once globally.
38    pub constant_equalities: Vec<(F, ContextCell)>,
39
40    external_cell_count: usize,
41
42    // In circuit assignments
43    /// Advice assignments, mapping from virtual [ContextCell] to assigned physical [Cell]
44    pub assigned_advices: HashMap<ContextCell, Cell>,
45    /// Constant assignments, (key = constant, value = [Cell])
46    pub assigned_constants: BTreeMap<F, Cell>,
47    /// Flag for whether `assign_raw` has been called, for safety only.
48    assigned: OnceLock<()>,
49}
50
51impl<F: Field + Ord> CopyConstraintManager<F> {
52    /// Returns the number of distinct constants used.
53    pub fn num_distinct_constants(&self) -> usize {
54        self.constant_equalities.iter().map(|(x, _)| x).sorted().dedup().count()
55    }
56
57    /// Adds external raw [Halo2AssignedCell] to `self.assigned_advices` and returns a new virtual [AssignedValue]
58    /// that can be used in any virtual region. No copy constraint is imposed, as the virtual cell "points" to the
59    /// raw assigned cell. The returned [ContextCell] will have `type_id` the `TypeId::of::<Cell>()`.
60    pub fn load_external_assigned(
61        &mut self,
62        assigned_cell: Halo2AssignedCell<F>,
63    ) -> AssignedValue<F> {
64        let context_cell = self.load_external_cell(assigned_cell.cell());
65        let mut value = Assigned::Trivial(F::ZERO);
66        assigned_cell.value().map(|v| {
67            #[cfg(feature = "halo2-axiom")]
68            {
69                value = **v;
70            }
71            #[cfg(not(feature = "halo2-axiom"))]
72            {
73                value = *v;
74            }
75        });
76        AssignedValue { value, cell: Some(context_cell) }
77    }
78
79    /// Adds external raw Halo2 cell to `self.assigned_advices` and returns a new virtual cell that can be
80    /// used as a tag (but will not be re-assigned). The returned [ContextCell] will have `type_id` the `TypeId::of::<Cell>()`.
81    pub fn load_external_cell(&mut self, cell: Cell) -> ContextCell {
82        self.load_external_cell_impl(Some(cell))
83    }
84
85    /// Mock to load an external cell for base circuit simulation. If any mock external cell is loaded, calling `assign_raw` will panic.
86    pub fn mock_external_assigned(&mut self, v: F) -> AssignedValue<F> {
87        let context_cell = self.load_external_cell_impl(None);
88        AssignedValue { value: Assigned::Trivial(v), cell: Some(context_cell) }
89    }
90
91    fn load_external_cell_impl(&mut self, cell: Option<Cell>) -> ContextCell {
92        let context_cell = ContextCell::new(EXTERNAL_CELL_TYPE_ID, 0, self.external_cell_count);
93        self.external_cell_count += 1;
94        if let Some(cell) = cell {
95            if let Some(old_cell) = self.assigned_advices.insert(context_cell, cell) {
96                assert!(
97                    old_cell.row_offset == cell.row_offset && old_cell.column == cell.column,
98                    "External cell already assigned"
99                )
100            }
101        }
102        context_cell
103    }
104
105    /// Clears state
106    pub fn clear(&mut self) {
107        self.advice_equalities.clear();
108        self.constant_equalities.clear();
109        self.assigned_advices.clear();
110        self.assigned_constants.clear();
111        self.external_cell_count = 0;
112        self.assigned.take();
113    }
114}
115
116impl<F: Field + Ord> Drop for CopyConstraintManager<F> {
117    fn drop(&mut self) {
118        if self.assigned.get().is_some() {
119            return;
120        }
121        if !self.advice_equalities.is_empty() {
122            dbg!("WARNING: advice_equalities not empty");
123        }
124        if !self.constant_equalities.is_empty() {
125            dbg!("WARNING: constant_equalities not empty");
126        }
127    }
128}
129
130impl<F: Field + Ord> VirtualRegionManager<F> for SharedCopyConstraintManager<F> {
131    // The fixed columns
132    type Config = Vec<Column<Fixed>>;
133    type Assignment = ();
134
135    /// This should be the last manager to be assigned, after all other managers have assigned cells.
136    fn assign_raw(&self, config: &Self::Config, region: &mut Region<F>) -> Self::Assignment {
137        let mut guard = self.lock().unwrap();
138        let manager = guard.deref_mut();
139        // sort by constant so constant assignment order is deterministic
140        // this is necessary because constants can be assigned by multiple CPU threads
141        // We further sort by ContextCell because the backend implementation of `raw_constrain_equal` (permutation argument) seems to depend on the order you specify copy constraints...
142        manager
143            .constant_equalities
144            .par_sort_unstable_by(|(c1, cell1), (c2, cell2)| c1.cmp(c2).then(cell1.cmp(cell2)));
145        // Assign fixed cells, we go left to right, then top to bottom, to avoid needing to know number of rows here
146        let mut fixed_col = 0;
147        let mut fixed_offset = 0;
148        for (c, _) in manager.constant_equalities.iter() {
149            if !manager.assigned_constants.contains_key(c) {
150                // this will panic if you run out of rows
151                let cell = raw_assign_fixed(region, config[fixed_col], fixed_offset, *c);
152                manager.assigned_constants.insert(*c, cell);
153                fixed_col += 1;
154                if fixed_col >= config.len() {
155                    fixed_col = 0;
156                    fixed_offset += 1;
157                }
158            }
159        }
160
161        // Just in case: we sort by ContextCell because the backend implementation of `raw_constrain_equal` (permutation argument) seems to depend on the order you specify copy constraints...
162        manager.advice_equalities.par_sort_unstable();
163        // Impose equality constraints between assigned advice cells
164        // At this point we assume all cells have been assigned by other VirtualRegionManagers
165        for (left, right) in &manager.advice_equalities {
166            let left = manager.assigned_advices.get(left).expect("virtual cell not assigned");
167            let right = manager.assigned_advices.get(right).expect("virtual cell not assigned");
168            raw_constrain_equal(region, *left, *right);
169        }
170        for (left, right) in &manager.constant_equalities {
171            let left = manager.assigned_constants[left];
172            let right = manager.assigned_advices.get(right).expect("virtual cell not assigned");
173            raw_constrain_equal(region, left, *right);
174        }
175        // We can't clear advice_equalities and constant_equalities because keygen_vk and keygen_pk will call this function twice
176        let _ = manager.assigned.set(());
177        // When keygen_vk and keygen_pk are both run, you need to clear assigned constants
178        // so the second run still assigns constants in the pk
179        manager.assigned_constants.clear();
180    }
181}
halo2_base/virtual_region/copy_constraints.rs

halo2_base/virtual_region/
copy_constraints.rs
