1use crate::{
6 columns::ExtendedFoldingColumn,
7 decomposable_folding::check_selector,
8 eval_leaf::EvalLeaf,
9 expressions::{Degree, FoldingExp, IntegratedFoldingExpr, Sign},
10 quadraticization::ExtendedWitnessGenerator,
11 FoldingConfig, FoldingEnv, Instance, RelaxedInstance, RelaxedWitness, ScalarField,
12};
13use ark_ec::AdditiveGroup;
14use ark_ff::{Field, One, Zero};
15use ark_poly::{Evaluations, Radix2EvaluationDomain};
16use kimchi::circuits::expr::Variable;
17use poly_commitment::{PolyComm, SRS};
18
19#[derive(Clone, Copy)]
25pub enum Side {
26 Left = 0,
27 Right = 1,
28}
29
30impl Side {
31 pub fn other(self) -> Self {
32 match self {
33 Side::Left => Side::Right,
34 Side::Right => Side::Left,
35 }
36 }
37}
38
39pub(crate) fn eval_sided<'a, C: FoldingConfig>(
41 exp: &FoldingExp<C>,
42 env: &'a ExtendedEnv<C>,
43 side: Side,
44) -> EvalLeaf<'a, ScalarField<C>> {
45 use FoldingExp::*;
46
47 match exp {
48 Atom(col) => env.col(col, side),
49 Double(e) => {
50 let col = eval_sided(e, env, side);
51 col.map(AdditiveGroup::double, |f| {
52 AdditiveGroup::double_in_place(f);
53 })
54 }
55 Square(e) => {
56 let col = eval_sided(e, env, side);
57 col.map(Field::square, |f| {
58 Field::square_in_place(f);
59 })
60 }
61 Add(e1, e2) => eval_sided(e1, env, side) + eval_sided(e2, env, side),
62 Sub(e1, e2) => eval_sided(e1, env, side) - eval_sided(e2, env, side),
63 Mul(e1, e2) => {
64 let selector = check_selector(e1)
66 .or(check_selector(e2))
67 .zip(env.enabled_selector())
68 .map(|(s1, s2)| s1 == s2);
69 match selector {
70 Some(false) => {
71 let zero_vec = vec![ScalarField::<C>::zero(); env.domain.size as usize];
72 EvalLeaf::Result(zero_vec)
73 }
74 Some(true) | None => {
75 let d1 = e1.folding_degree();
76 let d2 = e2.folding_degree();
77 let e1 = match d1 {
78 Degree::Two => eval_sided(e1, env, side),
79 _ => eval_exp_error(e1, env, side),
80 };
81 let e2 = match d2 {
82 Degree::Two => eval_sided(e2, env, side),
83 _ => eval_exp_error(e2, env, side),
84 };
85 e1 * e2
86 }
87 }
88 }
89 Pow(e, i) => match i {
90 0 => EvalLeaf::Const(ScalarField::<C>::one()),
91 1 => eval_sided(e, env, side),
92 i => {
93 let err = eval_sided(e, env, side);
94 let mut acc = err.clone();
95 for _ in 1..*i {
96 acc = acc * err.clone()
97 }
98 acc
99 }
100 },
101 }
102}
103
104pub(crate) fn eval_exp_error<'a, C: FoldingConfig>(
105 exp: &FoldingExp<C>,
106 env: &'a ExtendedEnv<C>,
107 side: Side,
108) -> EvalLeaf<'a, ScalarField<C>> {
109 use FoldingExp::*;
110
111 match exp {
112 Atom(col) => env.col(col, side),
113 Double(e) => {
114 let col = eval_exp_error(e, env, side);
115 col.map(AdditiveGroup::double, |f| {
116 AdditiveGroup::double_in_place(f);
117 })
118 }
119 Square(e) => match exp.folding_degree() {
120 Degree::Two => {
121 let cross = eval_exp_error(e, env, side) * eval_exp_error(e, env, side.other());
122 cross.map(AdditiveGroup::double, |f| {
123 AdditiveGroup::double_in_place(f);
124 })
125 }
126 _ => {
127 let e = eval_exp_error(e, env, side);
128 e.map(Field::square, |f| {
129 Field::square_in_place(f);
130 })
131 }
132 },
133 Add(e1, e2) => eval_exp_error(e1, env, side) + eval_exp_error(e2, env, side),
134 Sub(e1, e2) => eval_exp_error(e1, env, side) - eval_exp_error(e2, env, side),
135 Mul(e1, e2) => {
136 let selector = check_selector(e1)
138 .or(check_selector(e2))
139 .zip(env.enabled_selector())
140 .map(|(s1, s2)| s1 == s2);
141 match selector {
142 Some(false) => {
143 let zero_vec = vec![ScalarField::<C>::zero(); env.domain.size as usize];
144 EvalLeaf::Result(zero_vec)
145 }
146 Some(true) | None => match (exp.folding_degree(), e1.folding_degree()) {
147 (Degree::Two, Degree::One) => {
148 let first =
149 eval_exp_error(e1, env, side) * eval_exp_error(e2, env, side.other());
150 let second =
151 eval_exp_error(e1, env, side.other()) * eval_exp_error(e2, env, side);
152 first + second
153 }
154 _ => eval_exp_error(e1, env, side) * eval_exp_error(e2, env, side),
155 },
156 }
157 }
158 Pow(_, 0) => EvalLeaf::Const(ScalarField::<C>::one()),
159 Pow(e, 1) => eval_exp_error(e, env, side),
160 Pow(e, 2) => match (exp.folding_degree(), e.folding_degree()) {
161 (Degree::Two, Degree::One) => {
162 let first = eval_exp_error(e, env, side) * eval_exp_error(e, env, side.other());
163 let second = eval_exp_error(e, env, side.other()) * eval_exp_error(e, env, side);
164 first + second
165 }
166 _ => {
167 let err = eval_exp_error(e, env, side);
168 err.clone() * err
169 }
170 },
171 Pow(e, i) => match exp.folding_degree() {
172 Degree::Zero => {
173 let e = eval_exp_error(e, env, side);
174 let mut acc = e.clone();
176 for _ in 1..*i {
177 acc = acc * e.clone();
178 }
179 acc
180 }
181 _ => panic!("degree over 2"),
182 },
183 }
184}
185
186pub(crate) fn compute_error<C: FoldingConfig>(
193 exp: &IntegratedFoldingExpr<C>,
194 env: &ExtendedEnv<C>,
195 u: (ScalarField<C>, ScalarField<C>),
196) -> [Vec<ScalarField<C>>; 2] {
197 let (ul, ur) = (u.0, u.1);
204 let u_cross = ul * ur;
205 let zero_vec = vec![ScalarField::<C>::zero(); env.domain.size as usize];
206 let zero = || EvalLeaf::Result(zero_vec.clone());
207
208 let alphas_l = env
209 .get_relaxed_instance(Side::Left)
210 .extended_instance
211 .instance
212 .get_alphas();
213 let alphas_r = env
214 .get_relaxed_instance(Side::Right)
215 .extended_instance
216 .instance
217 .get_alphas();
218
219 let t_0 = {
220 let t_0 = (zero(), zero());
221 let (l, r) = exp.degree_0.iter().fold(t_0, |(l, r), (exp, sign, alpha)| {
222 let exp = eval_exp_error(exp, env, Side::Left);
224 let alpha_l = alphas_l.get(*alpha).expect("alpha not present");
225 let alpha_r = alphas_r.get(*alpha).expect("alpha not present");
226 let left = exp.clone() * alpha_l;
227 let right = exp * alpha_r;
228 match sign {
229 Sign::Pos => (l + left, r + right),
230 Sign::Neg => (l - left, r - right),
231 }
232 });
233 let cross2 = u_cross.double();
234 let e0 = l.clone() * cross2 + r.clone() * ul.square();
235 let e1 = r * cross2 + l * ur.square();
236 (e0, e1)
237 };
238
239 let t_1 = {
240 let t_1 = (zero(), zero(), zero());
241 let (l, cross, r) = exp
242 .degree_1
243 .iter()
244 .fold(t_1, |(l, cross, r), (exp, sign, alpha)| {
245 let expl = eval_exp_error(exp, env, Side::Left);
246 let expr = eval_exp_error(exp, env, Side::Right);
247 let alpha_l = alphas_l.get(*alpha).expect("alpha not present");
248 let alpha_r = alphas_r.get(*alpha).expect("alpha not present");
249 let expr_cross = expl.clone() * alpha_r + expr.clone() * alpha_l;
250 let left = expl * alpha_l;
251 let right = expr * alpha_r;
252 match sign {
253 Sign::Pos => (l + left, cross + expr_cross, r + right),
254 Sign::Neg => (l - left, cross - expr_cross, r - right),
255 }
256 });
257 let e0 = cross.clone() * ul + l * ur;
258 let e1 = cross.clone() * ur + r * ul;
259 (e0, e1)
260 };
261 let t_2 = (zero(), zero());
262 let t_2 = exp.degree_2.iter().fold(t_2, |(l, r), (exp, sign, alpha)| {
263 let expl = eval_sided(exp, env, Side::Left);
264 let expr = eval_sided(exp, env, Side::Right);
265 let cross = eval_exp_error(exp, env, Side::Left);
267 let alpha_l = alphas_l.get(*alpha).expect("alpha not present");
268 let alpha_r = alphas_r.get(*alpha).expect("alpha not present");
269 let left = expl * alpha_r + cross.clone() * alpha_l;
270 let right = expr * alpha_l + cross * alpha_r;
271 match sign {
272 Sign::Pos => (l + left, r + right),
273 Sign::Neg => (l - left, r - right),
274 }
275 });
276 let t = [t_1, t_2]
277 .into_iter()
278 .fold(t_0, |(tl, tr), (txl, txr)| (tl + txl, tr + txr));
279
280 match t {
281 (EvalLeaf::Result(l), EvalLeaf::Result(r)) => [l, r],
282 _ => unreachable!(),
283 }
284}
285
286pub(crate) struct ExtendedEnv<CF: FoldingConfig> {
292 inner: CF::Env,
293 instances: [RelaxedInstance<CF::Curve, CF::Instance>; 2],
294 witnesses: [RelaxedWitness<CF::Curve, CF::Witness>; 2],
295 domain: Radix2EvaluationDomain<ScalarField<CF>>,
296 selector: Option<CF::Selector>,
297}
298
299impl<CF: FoldingConfig> ExtendedEnv<CF> {
300 pub fn new(
301 structure: &CF::Structure,
302 instances: [RelaxedInstance<CF::Curve, CF::Instance>; 2],
304 witnesses: [RelaxedWitness<CF::Curve, CF::Witness>; 2],
305 domain: Radix2EvaluationDomain<ScalarField<CF>>,
306 selector: Option<CF::Selector>,
307 ) -> Self {
308 let inner_instances = [
309 &instances[0].extended_instance.instance,
310 &instances[1].extended_instance.instance,
311 ];
312 let inner_witnesses = [
313 &witnesses[0].extended_witness.witness,
314 &witnesses[1].extended_witness.witness,
315 ];
316 let inner = <CF::Env>::new(structure, inner_instances, inner_witnesses);
317 Self {
318 inner,
319 instances,
320 witnesses,
321 domain,
322 selector,
323 }
324 }
325
326 pub fn enabled_selector(&self) -> Option<&CF::Selector> {
327 self.selector.as_ref()
328 }
329
330 #[allow(clippy::type_complexity)]
331 pub fn unwrap(
332 self,
333 ) -> (
334 [RelaxedInstance<CF::Curve, CF::Instance>; 2],
335 [RelaxedWitness<CF::Curve, CF::Witness>; 2],
336 ) {
337 let Self {
338 instances,
339 witnesses,
340 ..
341 } = self;
342 (instances, witnesses)
343 }
344
345 pub fn get_relaxed_instance(&self, side: Side) -> &RelaxedInstance<CF::Curve, CF::Instance> {
346 &self.instances[side as usize]
347 }
348
349 pub fn get_relaxed_witness(&self, side: Side) -> &RelaxedWitness<CF::Curve, CF::Witness> {
350 &self.witnesses[side as usize]
351 }
352
353 pub fn col(&self, col: &ExtendedFoldingColumn<CF>, side: Side) -> EvalLeaf<ScalarField<CF>> {
354 use EvalLeaf::Col;
355 use ExtendedFoldingColumn::*;
356 let relaxed_instance = self.get_relaxed_instance(side);
357 let relaxed_witness = self.get_relaxed_witness(side);
358 let alphas = relaxed_instance.extended_instance.instance.get_alphas();
359 match col {
360 Inner(Variable { col, row }) => Col(self.inner.col(*col, *row, side)),
361 WitnessExtended(i) => Col(&relaxed_witness
362 .extended_witness
363 .extended
364 .get(i)
365 .expect("extended column not present")
366 .evals),
367 Error => panic!("shouldn't happen"),
368 Constant(c) => EvalLeaf::Const(*c),
369 Challenge(chall) => EvalLeaf::Const(self.inner.challenge(*chall, side)),
370 Alpha(i) => {
371 let alpha = alphas.get(*i).expect("alpha not present");
372 EvalLeaf::Const(alpha)
373 }
374 Selector(s) => Col(self.inner.selector(s, side)),
375 }
376 }
377
378 pub fn col_try(&self, col: &ExtendedFoldingColumn<CF>, side: Side) -> bool {
379 use ExtendedFoldingColumn::*;
380 let relaxed_witness = self.get_relaxed_witness(side);
381 match col {
382 WitnessExtended(i) => relaxed_witness.extended_witness.extended.contains_key(i),
383 Error => panic!("shouldn't happen"),
384 Inner(_) | Constant(_) | Challenge(_) | Alpha(_) | Selector(_) => true,
385 }
386 }
387
388 pub fn add_witness_evals(&mut self, i: usize, evals: Vec<ScalarField<CF>>, side: Side) {
389 let (_instance, relaxed_witness) = match side {
390 Side::Left => (&self.instances[0], &mut self.witnesses[0]),
391 Side::Right => (&self.instances[1], &mut self.witnesses[1]),
392 };
393 let evals = Evaluations::from_vec_and_domain(evals, self.domain);
394 relaxed_witness.extended_witness.add_witness_evals(i, evals);
395 }
396
397 pub fn needs_extension(&self, side: Side) -> bool {
398 !match side {
399 Side::Left => self.witnesses[0].extended_witness.is_extended(),
400 Side::Right => self.witnesses[1].extended_witness.is_extended(),
401 }
402 }
403
404 pub fn compute_extension(
407 self,
408 witness_generator: &ExtendedWitnessGenerator<CF>,
409 srs: &CF::Srs,
410 ) -> Self {
411 let env = self;
412 let env = witness_generator.compute_extended_witness(env, Side::Left);
413 let env = witness_generator.compute_extended_witness(env, Side::Right);
414 let env = env.compute_extended_commitments(srs, Side::Left);
415 env.compute_extended_commitments(srs, Side::Right)
416 }
417
418 fn compute_extended_commitments(mut self, srs: &CF::Srs, side: Side) -> Self {
430 let (relaxed_instance, relaxed_witness) = match side {
431 Side::Left => (&mut self.instances[0], &self.witnesses[0]),
432 Side::Right => (&mut self.instances[1], &self.witnesses[1]),
433 };
434
435 let blinder = PolyComm::new(vec![relaxed_instance.blinder]);
437 for (expected_i, (i, wit)) in relaxed_witness.extended_witness.extended.iter().enumerate() {
438 assert_eq!(*i, expected_i);
440 let commit = srs
443 .commit_evaluations_custom(self.domain, wit, &blinder)
444 .unwrap()
445 .commitment;
446 relaxed_instance.extended_instance.extended.push(commit)
447 }
448 self
451 }
452
453 pub(crate) fn to_absorb(
456 &self,
457 t0: &CF::Curve,
458 t1: &CF::Curve,
459 ) -> (Vec<ScalarField<CF>>, Vec<CF::Curve>) {
460 let mut left = self.instances[0].to_absorb();
461 let right = self.instances[1].to_absorb();
462
463 left.0.extend(right.0);
464 left.1.extend(right.1);
465 left.1.extend([t0, t1]);
466 left
467 }
468}