folding/
quadraticization.rs

1//! A library to reduce constraints into degree 2.
2// TODO: more documentation
3
4use crate::{
5    columns::ExtendedFoldingColumn,
6    error_term::{eval_sided, ExtendedEnv, Side},
7    expressions::{Degree, FoldingExp},
8    FoldingConfig,
9};
10use std::collections::{BTreeMap, HashMap, VecDeque};
11
12pub struct Quadraticized<C: FoldingConfig> {
13    pub original_constraints: Vec<FoldingExp<C>>,
14    pub extra_constraints: Vec<FoldingExp<C>>,
15    pub extended_witness_generator: ExtendedWitnessGenerator<C>,
16}
17
18/// Returns the constraints converted into degree 2 or less and the extra
19/// constraints added in the process
20pub fn quadraticize<C: FoldingConfig>(
21    constraints: Vec<FoldingExp<C>>,
22) -> (Quadraticized<C>, usize) {
23    let mut recorder = ExpRecorder::new();
24    let original_constraints = constraints
25        .into_iter()
26        .map(|exp| lower_degree_to_2(exp, &mut recorder))
27        .collect();
28    let added_columns = recorder.next;
29    let (extra_constraints, exprs) = recorder.into_constraints();
30    (
31        Quadraticized {
32            original_constraints,
33            extra_constraints,
34            extended_witness_generator: ExtendedWitnessGenerator { exprs },
35        },
36        added_columns,
37    )
38}
39
40/// Records expressions that have been extracted into an extra column
41struct ExpRecorder<C: FoldingConfig> {
42    recorded_exprs: HashMap<FoldingExp<C>, usize>,
43    pub(crate) next: usize,
44}
45
46impl<C: FoldingConfig> ExpRecorder<C> {
47    fn new() -> Self {
48        Self {
49            recorded_exprs: Default::default(),
50            next: 0,
51        }
52    }
53
54    fn get_id(&mut self, exp: FoldingExp<C>) -> usize {
55        *self.recorded_exprs.entry(exp).or_insert_with(|| {
56            let id = self.next;
57            self.next += 1;
58            id
59        })
60    }
61
62    #[allow(clippy::type_complexity)]
63    fn into_constraints(self) -> (Vec<FoldingExp<C>>, VecDeque<(usize, FoldingExp<C>)>) {
64        let ExpRecorder { recorded_exprs, .. } = self;
65        let mut witness_generator = VecDeque::with_capacity(recorded_exprs.len());
66        let mut new_constraints = BTreeMap::new();
67        for (exp, id) in recorded_exprs.into_iter() {
68            let left = FoldingExp::Atom(ExtendedFoldingColumn::WitnessExtended(id));
69            let constraint = FoldingExp::Sub(Box::new(left), Box::new(exp.clone()));
70            new_constraints.insert(id, constraint);
71            witness_generator.push_front((id, exp));
72        }
73        (new_constraints.into_values().collect(), witness_generator)
74    }
75}
76
77impl<C: FoldingConfig> FoldingExp<C> {
78    fn degree(&self) -> usize {
79        match self {
80            e @ FoldingExp::Atom(_) => match e.folding_degree() {
81                Degree::Zero => 0,
82                Degree::One => 1,
83                Degree::Two => 2,
84            },
85            FoldingExp::Double(exp) => exp.degree(),
86            FoldingExp::Square(exp) => exp.degree() * 2,
87            FoldingExp::Add(e1, e2) | FoldingExp::Sub(e1, e2) => {
88                core::cmp::max(e1.degree(), e2.degree())
89            }
90            FoldingExp::Mul(e1, e2) => e1.degree() + e2.degree(),
91            FoldingExp::Pow(e, i) => e.degree() * (*i as usize),
92        }
93    }
94}
95
96fn lower_degree_to_1<C: FoldingConfig>(
97    exp: FoldingExp<C>,
98    rec: &mut ExpRecorder<C>,
99) -> FoldingExp<C> {
100    let degree = exp.degree();
101    match degree {
102        0 => exp,
103        1 => exp,
104        _ => match exp {
105            FoldingExp::Add(e1, e2) => FoldingExp::Add(
106                Box::new(lower_degree_to_1(*e1, rec)),
107                Box::new(lower_degree_to_1(*e2, rec)),
108            ),
109            FoldingExp::Sub(e1, e2) => FoldingExp::Sub(
110                Box::new(lower_degree_to_1(*e1, rec)),
111                Box::new(lower_degree_to_1(*e2, rec)),
112            ),
113            e @ FoldingExp::Square(_) | e @ FoldingExp::Mul(_, _) => {
114                let exp = lower_degree_to_2(e, rec);
115                let id = rec.get_id(exp);
116                FoldingExp::Atom(ExtendedFoldingColumn::WitnessExtended(id))
117            }
118            FoldingExp::Double(exp) => FoldingExp::Double(Box::new(lower_degree_to_1(*exp, rec))),
119            _ => todo!(),
120        },
121    }
122}
123
124fn lower_degree_to_2<C: FoldingConfig>(
125    exp: FoldingExp<C>,
126    rec: &mut ExpRecorder<C>,
127) -> FoldingExp<C> {
128    use FoldingExp::*;
129    let degree = exp.degree();
130    if degree <= 2 {
131        return exp;
132    }
133
134    match exp {
135        FoldingExp::Atom(_) => panic!("a column shouldn't be above degree 1"),
136        FoldingExp::Double(exp) => Double(Box::new(lower_degree_to_2(*exp, rec))),
137        FoldingExp::Square(exp) => Square(Box::new(lower_degree_to_1(*exp, rec))),
138        FoldingExp::Add(e1, e2) => {
139            let e1 = lower_degree_to_2(*e1, rec);
140            let e2 = lower_degree_to_2(*e2, rec);
141            Add(Box::new(e1), Box::new(e2))
142        }
143        FoldingExp::Sub(e1, e2) => {
144            let e1 = lower_degree_to_2(*e1, rec);
145            let e2 = lower_degree_to_2(*e2, rec);
146            Sub(Box::new(e1), Box::new(e2))
147        }
148        FoldingExp::Mul(e1, e2) => {
149            let d1 = e1.degree();
150            let d2 = e2.degree();
151            assert_eq!(degree, d1 + d2);
152            let (e1, e2) = (*e1, *e2);
153            let (e1, e2) = match (d1, d2) {
154                (0, _) => (e1, lower_degree_to_2(e2, rec)),
155                (_, 0) => (lower_degree_to_2(e1, rec), e2),
156                (_, _) => (lower_degree_to_1(e1, rec), lower_degree_to_1(e2, rec)),
157            };
158            Mul(Box::new(e1), Box::new(e2))
159        }
160        exp @ FoldingExp::Pow(_, 0) => exp,
161        FoldingExp::Pow(e, 1) => lower_degree_to_2(*e, rec),
162        FoldingExp::Pow(e, 2) => FoldingExp::Pow(Box::new(lower_degree_to_1(*e, rec)), 2),
163        FoldingExp::Pow(e, i) => {
164            let e = lower_degree_to_1(*e, rec);
165            let mut acc = e.clone();
166            for _ in 1..i - 1 {
167                acc = lower_degree_to_1(FoldingExp::Mul(Box::new(e.clone()), Box::new(acc)), rec);
168            }
169            FoldingExp::Mul(Box::new(e.clone()), Box::new(acc))
170        }
171    }
172}
173
174pub struct ExtendedWitnessGenerator<C: FoldingConfig> {
175    exprs: VecDeque<(usize, FoldingExp<C>)>,
176}
177
178impl<C: FoldingConfig> ExtendedWitnessGenerator<C> {
179    pub(crate) fn compute_extended_witness(
180        &self,
181        mut env: ExtendedEnv<C>,
182        side: Side,
183    ) -> ExtendedEnv<C> {
184        if env.needs_extension(side) {
185            let mut pending = self.exprs.clone();
186
187            while let Some((i, exp)) = pending.pop_front() {
188                if check_evaluable(&exp, &env, side) {
189                    let evals = eval_sided(&exp, &env, side).unwrap();
190                    env.add_witness_evals(i, evals, side);
191                } else {
192                    pending.push_back((i, exp))
193                }
194            }
195        }
196
197        env
198    }
199}
200
201/// Checks if the expression can be evaluated in the current environment
202fn check_evaluable<C: FoldingConfig>(
203    exp: &FoldingExp<C>,
204    env: &ExtendedEnv<C>,
205    side: Side,
206) -> bool {
207    match exp {
208        FoldingExp::Atom(col) => env.col_try(col, side),
209        FoldingExp::Double(e) | FoldingExp::Square(e) => check_evaluable(e, env, side),
210        FoldingExp::Add(e1, e2) | FoldingExp::Sub(e1, e2) | FoldingExp::Mul(e1, e2) => {
211            check_evaluable(e1, env, side) && check_evaluable(e2, env, side)
212        }
213        FoldingExp::Pow(e, _) => check_evaluable(e, env, side),
214    }
215}
folding/quadraticization.rs

folding/
quadraticization.rs
