1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use ark_ff::PrimeField;
use kimchi::circuits::{
    berkeley_columns::BerkeleyChallengeTerm,
    expr::{ConstantExpr, ConstantTerm, Expr, ExprInner, Variable},
    gate::CurrOrNext,
};
use std::collections::BTreeMap;

use crate::{
    circuit_design::capabilities::{ColAccessCap, HybridCopyCap, LookupCap},
    columns::{Column, ColumnIndexer},
    expr::E,
    logup::{constraint_lookups, LookupTableID},
};

pub struct ConstraintBuilderEnv<F: PrimeField, LT: LookupTableID> {
    /// An indexed set of constraints.
    pub constraints: Vec<Expr<ConstantExpr<F, BerkeleyChallengeTerm>, Column>>,
    /// Aggregated lookups or "reads".
    pub lookup_reads: BTreeMap<LT, Vec<Vec<E<F>>>>,
    /// Aggregated "write" lookups, for runtime tables.
    pub lookup_writes: BTreeMap<LT, Vec<Vec<E<F>>>>,
    /// The function that maps the argument of `assert_zero`.
    pub assert_mapper: Box<dyn Fn(E<F>) -> E<F>>,
}

impl<F: PrimeField, LT: LookupTableID> ConstraintBuilderEnv<F, LT> {
    pub fn create() -> Self {
        Self {
            constraints: vec![],
            lookup_reads: BTreeMap::new(),
            lookup_writes: BTreeMap::new(),
            assert_mapper: Box::new(|x| x),
        }
    }
}

impl<F: PrimeField, CIx: ColumnIndexer, LT: LookupTableID> ColAccessCap<F, CIx>
    for ConstraintBuilderEnv<F, LT>
{
    type Variable = E<F>;

    fn assert_zero(&mut self, cst: Self::Variable) {
        self.constraints.push((self.assert_mapper)(cst));
    }

    fn set_assert_mapper(&mut self, mapper: Box<dyn Fn(Self::Variable) -> Self::Variable>) {
        self.assert_mapper = mapper;
    }

    fn read_column(&self, position: CIx) -> Self::Variable {
        Expr::Atom(ExprInner::Cell(Variable {
            col: position.to_column(),
            row: CurrOrNext::Curr,
        }))
    }

    fn constant(value: F) -> Self::Variable {
        let cst_expr_inner = ConstantExpr::from(ConstantTerm::Literal(value));
        Expr::Atom(ExprInner::Constant(cst_expr_inner))
    }
}

impl<F: PrimeField, CIx: ColumnIndexer, LT: LookupTableID> HybridCopyCap<F, CIx>
    for ConstraintBuilderEnv<F, LT>
{
    fn hcopy(&mut self, x: &Self::Variable, position: CIx) -> Self::Variable {
        let y = Expr::Atom(ExprInner::Cell(Variable {
            col: position.to_column(),
            row: CurrOrNext::Curr,
        }));
        <ConstraintBuilderEnv<F, LT> as ColAccessCap<F, CIx>>::assert_zero(
            self,
            y.clone() - x.clone(),
        );
        y
    }
}

impl<F: PrimeField, CIx: ColumnIndexer, LT: LookupTableID> LookupCap<F, CIx, LT>
    for ConstraintBuilderEnv<F, LT>
{
    fn lookup(&mut self, table_id: LT, value: Vec<<Self as ColAccessCap<F, CIx>>::Variable>) {
        self.lookup_reads
            .entry(table_id)
            .or_default()
            .push(value.clone());
    }

    fn lookup_runtime_write(&mut self, table_id: LT, value: Vec<Self::Variable>) {
        assert!(
            !table_id.is_fixed(),
            "lookup_runtime_write must be called on non-fixed tables only"
        );
        self.lookup_writes
            .entry(table_id)
            .or_default()
            .push(value.clone());
    }
}

impl<F: PrimeField, LT: LookupTableID> ConstraintBuilderEnv<F, LT> {
    /// Get constraints related to the application logic itself.
    pub fn get_relation_constraints(&self) -> Vec<E<F>> {
        self.constraints.clone()
    }

    /// Get constraints related to the lookup argument.
    pub fn get_lookup_constraints(&self) -> Vec<E<F>> {
        constraint_lookups(&self.lookup_reads, &self.lookup_writes)
    }

    /// Get all relevant constraints generated by the constraint builder.
    pub fn get_constraints(&self) -> Vec<E<F>> {
        let mut constraints: Vec<E<F>> = vec![];
        constraints.extend(self.get_relation_constraints());
        if !self.lookup_reads.is_empty() {
            constraints.extend(self.get_lookup_constraints());
        }
        constraints
    }
}