kimchi_msm/circuit_design/
constraints.rs

1use ark_ff::PrimeField;
2use kimchi::circuits::{
3    berkeley_columns::BerkeleyChallengeTerm,
4    expr::{ConstantExpr, ConstantTerm, Expr, ExprInner, Variable},
5    gate::CurrOrNext,
6};
7use std::collections::BTreeMap;
8
9use crate::{
10    circuit_design::capabilities::{ColAccessCap, HybridCopyCap, LookupCap},
11    columns::{Column, ColumnIndexer},
12    expr::E,
13    logup::{constraint_lookups, LookupTableID},
14};
15
16pub struct ConstraintBuilderEnv<F: PrimeField, LT: LookupTableID> {
17    /// An indexed set of constraints.
18    pub constraints: Vec<Expr<ConstantExpr<F, BerkeleyChallengeTerm>, Column<usize>>>,
19    /// Aggregated lookups or "reads".
20    pub lookup_reads: BTreeMap<LT, Vec<Vec<E<F>>>>,
21    /// Aggregated "write" lookups, for runtime tables.
22    pub lookup_writes: BTreeMap<LT, Vec<Vec<E<F>>>>,
23    /// The function that maps the argument of `assert_zero`.
24    pub assert_mapper: Box<dyn Fn(E<F>) -> E<F>>,
25}
26
27impl<F: PrimeField, LT: LookupTableID> ConstraintBuilderEnv<F, LT> {
28    pub fn create() -> Self {
29        Self {
30            constraints: vec![],
31            lookup_reads: BTreeMap::new(),
32            lookup_writes: BTreeMap::new(),
33            assert_mapper: Box::new(|x| x),
34        }
35    }
36}
37
38impl<F: PrimeField, CIx: ColumnIndexer<usize>, LT: LookupTableID> ColAccessCap<F, CIx>
39    for ConstraintBuilderEnv<F, LT>
40{
41    type Variable = E<F>;
42
43    fn assert_zero(&mut self, cst: Self::Variable) {
44        self.constraints.push((self.assert_mapper)(cst));
45    }
46
47    fn set_assert_mapper(&mut self, mapper: Box<dyn Fn(Self::Variable) -> Self::Variable>) {
48        self.assert_mapper = mapper;
49    }
50
51    fn read_column(&self, position: CIx) -> Self::Variable {
52        Expr::Atom(ExprInner::Cell(Variable {
53            col: position.to_column(),
54            row: CurrOrNext::Curr,
55        }))
56    }
57
58    fn constant(value: F) -> Self::Variable {
59        let cst_expr_inner = ConstantExpr::from(ConstantTerm::Literal(value));
60        Expr::Atom(ExprInner::Constant(cst_expr_inner))
61    }
62}
63
64impl<F: PrimeField, CIx: ColumnIndexer<usize>, LT: LookupTableID> HybridCopyCap<F, CIx>
65    for ConstraintBuilderEnv<F, LT>
66{
67    fn hcopy(&mut self, x: &Self::Variable, position: CIx) -> Self::Variable {
68        let y = Expr::Atom(ExprInner::Cell(Variable {
69            col: position.to_column(),
70            row: CurrOrNext::Curr,
71        }));
72        <ConstraintBuilderEnv<F, LT> as ColAccessCap<F, CIx>>::assert_zero(
73            self,
74            y.clone() - x.clone(),
75        );
76        y
77    }
78}
79
80impl<F: PrimeField, CIx: ColumnIndexer<usize>, LT: LookupTableID> LookupCap<F, CIx, LT>
81    for ConstraintBuilderEnv<F, LT>
82{
83    fn lookup(&mut self, table_id: LT, value: Vec<<Self as ColAccessCap<F, CIx>>::Variable>) {
84        self.lookup_reads
85            .entry(table_id)
86            .or_default()
87            .push(value.clone());
88    }
89
90    fn lookup_runtime_write(&mut self, table_id: LT, value: Vec<Self::Variable>) {
91        assert!(
92            !table_id.is_fixed(),
93            "lookup_runtime_write must be called on non-fixed tables only"
94        );
95        self.lookup_writes
96            .entry(table_id)
97            .or_default()
98            .push(value.clone());
99    }
100}
101
102impl<F: PrimeField, LT: LookupTableID> ConstraintBuilderEnv<F, LT> {
103    /// Get constraints related to the application logic itself.
104    pub fn get_relation_constraints(&self) -> Vec<E<F>> {
105        self.constraints.clone()
106    }
107
108    /// Get constraints related to the lookup argument.
109    pub fn get_lookup_constraints(&self) -> Vec<E<F>> {
110        constraint_lookups(&self.lookup_reads, &self.lookup_writes)
111    }
112
113    /// Get all relevant constraints generated by the constraint builder.
114    pub fn get_constraints(&self) -> Vec<E<F>> {
115        let mut constraints: Vec<E<F>> = vec![];
116        constraints.extend(self.get_relation_constraints());
117        if !self.lookup_reads.is_empty() {
118            constraints.extend(self.get_lookup_constraints());
119        }
120        constraints
121    }
122}