1use fields::PrimeField64;
2use std::sync::Arc;
3
4use pil_std_lib::Std;
5use proofman_common::{AirInstance, FromTrace, ProofmanResult, SetupCtx};
6use proofman_util::{timer_start_trace, timer_stop_and_log_trace};
7#[cfg(not(feature = "packed"))]
8use zisk_pil::{ArithEqTrace, ArithEqTraceRow};
9#[cfg(feature = "packed")]
10use zisk_pil::{ArithEqTracePacked, ArithEqTraceRowPacked};
11
12#[cfg(feature = "packed")]
13type ArithEqTraceRowType<F> = ArithEqTraceRowPacked<F>;
14#[cfg(feature = "packed")]
15type ArithEqTraceType<F> = ArithEqTracePacked<F>;
16
17#[cfg(not(feature = "packed"))]
18type ArithEqTraceRowType<F> = ArithEqTraceRow<F>;
19#[cfg(not(feature = "packed"))]
20type ArithEqTraceType<F> = ArithEqTrace<F>;
21
22use crate::{
23 arith_eq_constants::*, executors, Arith256Input, Arith256ModInput, ArithEqInput,
24 ArithEqLtTableSM, Bn254ComplexAddInput, Bn254ComplexMulInput, Bn254ComplexSubInput,
25 Bn254CurveAddInput, Bn254CurveDblInput, Secp256k1AddInput, Secp256k1DblInput,
26 SECP256K1_PRIME_CHUNKS, SEL_OP_ARITH256, SEL_OP_ARITH256_MOD, SEL_OP_SECP256K1_ADD,
27 SEL_OP_SECP256K1_DBL,
28};
29use rayon::prelude::*;
30
31/// The `ArithEqSM` struct encapsulates the logic of the ArithEq State Machine.
32pub struct ArithEqSM<F: PrimeField64> {
33 /// Number of available arith256s in the trace.
34 pub num_available_ops: usize,
35
36 /// Reference to the PIL2 standard library.
37 pub std: Arc<Std<F>>,
38
39 /// The table ID for the Keccakf Table State Machine
40 table_id: usize,
41
42 pub q_hsc_range_id: usize,
43 pub chunk_range_id: usize,
44 pub carry_range_id: usize,
45}
46#[derive(Debug, Default)]
47struct ArithEqStepAddr {
48 main_step: u64,
49 addr_op: u32,
50 addr_x1: u32,
51 addr_y1: u32,
52 addr_x2: u32,
53 addr_y2: u32,
54 addr_x3: u32,
55 addr_y3: u32,
56 addr_ind: [u32; 5],
57}
58
59impl<F: PrimeField64> ArithEqSM<F> {
60 /// Creates a new ArithEq State Machine instance.
61 ///
62 /// # Returns
63 /// A new `ArithEqSM` instance.
64 pub fn new(std: Arc<Std<F>>) -> Arc<Self> {
65 // Compute some useful values
66 let num_available_ops = ArithEqTraceType::<F>::NUM_ROWS / ARITH_EQ_ROWS_BY_OP;
67 let p2_22 = 1 << 22;
68 let q_hsc_range_id = std.get_range_id(0, p2_22 - 1, None).expect("Failed to get range ID");
69 let chunk_range_id = std.get_range_id(0, 0xFFFF, None).expect("Failed to get range ID");
70 let carry_range_id =
71 std.get_range_id(-(p2_22 - 1), p2_22, None).expect("Failed to get range ID");
72
73 // Get the table ID
74 let table_id =
75 std.get_virtual_table_id(ArithEqLtTableSM::TABLE_ID).expect("Failed to get table ID");
76
77 Arc::new(Self {
78 std,
79 num_available_ops,
80 q_hsc_range_id,
81 chunk_range_id,
82 carry_range_id,
83 table_id,
84 })
85 }
86 fn expand_addr_step_on_trace(data: &ArithEqStepAddr, trace: &mut [ArithEqTraceRowType<F>]) {
87 trace[0].set_step_addr(data.main_step);
88 trace[1].set_step_addr(data.addr_op as u64);
89 trace[2].set_step_addr(data.addr_x1 as u64);
90 trace[3].set_step_addr(data.addr_y1 as u64);
91 trace[4].set_step_addr(data.addr_x2 as u64);
92 trace[5].set_step_addr(data.addr_y2 as u64);
93 trace[6].set_step_addr(data.addr_x3 as u64);
94 trace[7].set_step_addr(data.addr_y3 as u64);
95 for (i, addr_ind) in data.addr_ind.iter().enumerate() {
96 trace[i + 8].set_step_addr(*addr_ind as u64);
97 }
98 for i in 0..(ARITH_EQ_ROWS_BY_OP - 8 - data.addr_ind.len()) {
99 trace[i + 8 + data.addr_ind.len()].set_step_addr(0);
100 }
101 }
102
103 fn process_arith256(&self, input: &Arith256Input, trace: &mut [ArithEqTraceRowType<F>]) {
104 let data = executors::Arith256::execute(&input.a, &input.b, &input.c);
105 self.expand_data_on_trace(&data, trace, SEL_OP_ARITH256);
106 Self::expand_addr_step_on_trace(
107 &ArithEqStepAddr {
108 main_step: input.step,
109 addr_op: input.addr,
110 addr_x1: input.a_addr,
111 addr_y1: input.b_addr,
112 addr_x2: input.c_addr,
113 addr_y2: 0,
114 addr_x3: input.dl_addr,
115 addr_y3: input.dh_addr,
116 addr_ind: [input.a_addr, input.b_addr, input.c_addr, input.dl_addr, input.dh_addr],
117 },
118 trace,
119 );
120 }
121
122 fn process_arith256_mod(&self, input: &Arith256ModInput, trace: &mut [ArithEqTraceRowType<F>]) {
123 let data = executors::Arith256Mod::execute(&input.a, &input.b, &input.c, &input.module);
124 self.expand_data_on_trace(&data, trace, SEL_OP_ARITH256_MOD);
125 Self::expand_addr_step_on_trace(
126 &ArithEqStepAddr {
127 main_step: input.step,
128 addr_op: input.addr,
129 addr_x1: input.a_addr,
130 addr_y1: input.b_addr,
131 addr_x2: input.c_addr,
132 addr_y2: input.module_addr,
133 addr_x3: input.d_addr,
134 addr_y3: 0,
135 addr_ind: [
136 input.a_addr,
137 input.b_addr,
138 input.c_addr,
139 input.module_addr,
140 input.d_addr,
141 ],
142 },
143 trace,
144 );
145 }
146 fn process_secp256k1_add(
147 &self,
148 input: &Secp256k1AddInput,
149 trace: &mut [ArithEqTraceRowType<F>],
150 ) {
151 let data = executors::Secp256k1::execute_add(&input.p1, &input.p2);
152 self.expand_data_on_trace(&data, trace, SEL_OP_SECP256K1_ADD);
153 Self::expand_addr_step_on_trace(
154 &ArithEqStepAddr {
155 main_step: input.step,
156 addr_op: input.addr,
157 addr_x1: input.p1_addr,
158 addr_y1: input.p1_addr + 32,
159 addr_x2: input.p2_addr,
160 addr_y2: input.p2_addr + 32,
161 addr_x3: input.p1_addr,
162 addr_y3: input.p1_addr + 32,
163 addr_ind: [input.p1_addr, input.p2_addr, 0, 0, 0],
164 },
165 trace,
166 );
167 }
168 fn process_secp256k1_dbl(
169 &self,
170 input: &Secp256k1DblInput,
171 trace: &mut [ArithEqTraceRowType<F>],
172 ) {
173 let data = executors::Secp256k1::execute_dbl(&input.p1);
174 self.expand_data_on_trace(&data, trace, SEL_OP_SECP256K1_DBL);
175 Self::expand_addr_step_on_trace(
176 &ArithEqStepAddr {
177 main_step: input.step,
178 addr_op: input.addr,
179 addr_x1: input.addr,
180 addr_y1: input.addr + 32,
181 addr_x2: input.addr,
182 addr_y2: input.addr + 32,
183 addr_x3: input.addr,
184 addr_y3: input.addr + 32,
185 addr_ind: [0, 0, 0, 0, 0],
186 },
187 trace,
188 );
189 }
190
191 fn process_bn254_curve_add(
192 &self,
193 input: &Bn254CurveAddInput,
194 trace: &mut [ArithEqTraceRowType<F>],
195 ) {
196 let data = executors::Bn254Curve::execute_add(&input.p1, &input.p2);
197 self.expand_data_on_trace(&data, trace, SEL_OP_BN254_CURVE_ADD);
198 Self::expand_addr_step_on_trace(
199 &ArithEqStepAddr {
200 main_step: input.step,
201 addr_op: input.addr,
202 addr_x1: input.p1_addr,
203 addr_y1: input.p1_addr + 32,
204 addr_x2: input.p2_addr,
205 addr_y2: input.p2_addr + 32,
206 addr_x3: input.p1_addr,
207 addr_y3: input.p1_addr + 32,
208 addr_ind: [input.p1_addr, input.p2_addr, 0, 0, 0],
209 },
210 trace,
211 );
212 }
213
214 fn process_bn254_curve_dbl(
215 &self,
216 input: &Bn254CurveDblInput,
217 trace: &mut [ArithEqTraceRowType<F>],
218 ) {
219 let data = executors::Bn254Curve::execute_dbl(&input.p1);
220 self.expand_data_on_trace(&data, trace, SEL_OP_BN254_CURVE_DBL);
221 Self::expand_addr_step_on_trace(
222 &ArithEqStepAddr {
223 main_step: input.step,
224 addr_op: input.addr,
225 addr_x1: input.addr,
226 addr_y1: input.addr + 32,
227 addr_x2: input.addr,
228 addr_y2: input.addr + 32,
229 addr_x3: input.addr,
230 addr_y3: input.addr + 32,
231 addr_ind: [0, 0, 0, 0, 0],
232 },
233 trace,
234 );
235 }
236
237 fn process_bn254_complex_add(
238 &self,
239 input: &Bn254ComplexAddInput,
240 trace: &mut [ArithEqTraceRowType<F>],
241 ) {
242 let data = executors::Bn254Complex::execute_add(&input.f1, &input.f2);
243 self.expand_data_on_trace(&data, trace, SEL_OP_BN254_COMPLEX_ADD);
244 Self::expand_addr_step_on_trace(
245 &ArithEqStepAddr {
246 main_step: input.step,
247 addr_op: input.addr,
248 addr_x1: input.f1_addr,
249 addr_y1: input.f1_addr + 32,
250 addr_x2: input.f2_addr,
251 addr_y2: input.f2_addr + 32,
252 addr_x3: input.f1_addr,
253 addr_y3: input.f1_addr + 32,
254 addr_ind: [input.f1_addr, input.f2_addr, 0, 0, 0],
255 },
256 trace,
257 );
258 }
259
260 fn process_bn254_complex_sub(
261 &self,
262 input: &Bn254ComplexSubInput,
263 trace: &mut [ArithEqTraceRowType<F>],
264 ) {
265 let data = executors::Bn254Complex::execute_sub(&input.f1, &input.f2);
266 self.expand_data_on_trace(&data, trace, SEL_OP_BN254_COMPLEX_SUB);
267 Self::expand_addr_step_on_trace(
268 &ArithEqStepAddr {
269 main_step: input.step,
270 addr_op: input.addr,
271 addr_x1: input.f1_addr,
272 addr_y1: input.f1_addr + 32,
273 addr_x2: input.f2_addr,
274 addr_y2: input.f2_addr + 32,
275 addr_x3: input.f1_addr,
276 addr_y3: input.f1_addr + 32,
277 addr_ind: [input.f1_addr, input.f2_addr, 0, 0, 0],
278 },
279 trace,
280 );
281 }
282
283 fn process_bn254_complex_mul(
284 &self,
285 input: &Bn254ComplexMulInput,
286 trace: &mut [ArithEqTraceRowType<F>],
287 ) {
288 let data = executors::Bn254Complex::execute_mul(&input.f1, &input.f2);
289 self.expand_data_on_trace(&data, trace, SEL_OP_BN254_COMPLEX_MUL);
290 Self::expand_addr_step_on_trace(
291 &ArithEqStepAddr {
292 main_step: input.step,
293 addr_op: input.addr,
294 addr_x1: input.f1_addr,
295 addr_y1: input.f1_addr + 32,
296 addr_x2: input.f2_addr,
297 addr_y2: input.f2_addr + 32,
298 addr_x3: input.f1_addr,
299 addr_y3: input.f1_addr + 32,
300 addr_ind: [input.f1_addr, input.f2_addr, 0, 0, 0],
301 },
302 trace,
303 );
304 }
305
306 #[inline(always)]
307 fn to_ranged_field(&self, value: i64, range_id: usize) -> u64 {
308 self.std.range_check(range_id, value, 1);
309 if value >= 0 {
310 value as u64
311 } else {
312 (F::ORDER_U64 as i64 + value) as u64
313 }
314 }
315
316 fn expand_data_on_trace(
317 &self,
318 data: &executors::ArithEqData,
319 trace: &mut [ArithEqTraceRowType<F>],
320 sel_op: usize,
321 ) {
322 let mut x1_x2_different = false;
323 let mut prev_x3_lt = false;
324 let mut prev_y3_lt = false;
325
326 #[allow(clippy::needless_range_loop)]
327 for i in 0..ARITH_EQ_ROWS_BY_OP {
328 for j in 0..3 {
329 // first position without carry
330 let carry_0 = if i == 0 { 0 } else { data.cout[i * 2 - 1][j] };
331 trace[i].set_carry(j, 0, self.to_ranged_field(carry_0, self.carry_range_id));
332 trace[i].set_carry(
333 j,
334 1,
335 self.to_ranged_field(data.cout[i * 2][j], self.carry_range_id),
336 );
337 }
338 let q_range_id = if i == ARITH_EQ_ROWS_BY_OP - 1 {
339 self.q_hsc_range_id
340 } else {
341 self.chunk_range_id
342 };
343 trace[i].set_x1(self.to_ranged_field(data.x1[i], self.chunk_range_id) as u16);
344 trace[i].set_y1(self.to_ranged_field(data.y1[i], self.chunk_range_id) as u16);
345 trace[i].set_x2(self.to_ranged_field(data.x2[i], self.chunk_range_id) as u16);
346 trace[i].set_y2(self.to_ranged_field(data.y2[i], self.chunk_range_id) as u16);
347 trace[i].set_x3(self.to_ranged_field(data.x3[i], self.chunk_range_id) as u16);
348 trace[i].set_y3(self.to_ranged_field(data.y3[i], self.chunk_range_id) as u16);
349 trace[i].set_q0(self.to_ranged_field(data.q0[i], q_range_id) as u32);
350 trace[i].set_q1(self.to_ranged_field(data.q1[i], q_range_id) as u32);
351 trace[i].set_q2(self.to_ranged_field(data.q2[i], q_range_id) as u32);
352 trace[i].set_s(self.to_ranged_field(data.s[i], self.chunk_range_id) as u32);
353
354 // TODO Range check
355 for j in 0..ARITH_EQ_OP_NUM {
356 let selected = j == sel_op;
357 trace[i].set_sel_op(j, selected);
358 if i == 0 {
359 trace[i].set_sel_op_clk0(j, selected);
360 } else {
361 trace[i].set_sel_op_clk0(j, false);
362 }
363 }
364 match sel_op {
365 SEL_OP_ARITH256_MOD => {
366 let x3_lt = data.x3[i] < data.y2[i] || (data.x3[i] == data.y2[i] && prev_x3_lt);
367 trace[i].set_x3_lt(x3_lt);
368 let row = ArithEqLtTableSM::calculate_table_row(
369 prev_x3_lt,
370 x3_lt,
371 data.x3[i] - data.y2[i],
372 );
373 self.std.inc_virtual_row(self.table_id, row as u64, 1);
374 prev_x3_lt = x3_lt;
375
376 trace[i].set_y3_lt(false);
377 }
378 SEL_OP_SECP256K1_ADD | SEL_OP_SECP256K1_DBL => {
379 let x3_lt = data.x3[i] < SECP256K1_PRIME_CHUNKS[i]
380 || (data.x3[i] == SECP256K1_PRIME_CHUNKS[i] && prev_x3_lt);
381 trace[i].set_x3_lt(x3_lt);
382 let row = ArithEqLtTableSM::calculate_table_row(
383 prev_x3_lt,
384 x3_lt,
385 data.x3[i] - SECP256K1_PRIME_CHUNKS[i],
386 );
387 self.std.inc_virtual_row(self.table_id, row as u64, 1);
388 prev_x3_lt = x3_lt;
389
390 let y3_lt = data.y3[i] < SECP256K1_PRIME_CHUNKS[i]
391 || (data.y3[i] == SECP256K1_PRIME_CHUNKS[i] && prev_y3_lt);
392 trace[i].set_y3_lt(y3_lt);
393 let row = ArithEqLtTableSM::calculate_table_row(
394 prev_y3_lt,
395 y3_lt,
396 data.y3[i] - SECP256K1_PRIME_CHUNKS[i],
397 );
398 self.std.inc_virtual_row(self.table_id, row as u64, 1);
399 prev_y3_lt = y3_lt;
400 }
401 SEL_OP_BN254_CURVE_ADD
402 | SEL_OP_BN254_CURVE_DBL
403 | SEL_OP_BN254_COMPLEX_ADD
404 | SEL_OP_BN254_COMPLEX_SUB
405 | SEL_OP_BN254_COMPLEX_MUL => {
406 let x3_lt = data.x3[i] < BN254_PRIME_CHUNKS[i]
407 || (data.x3[i] == BN254_PRIME_CHUNKS[i] && prev_x3_lt);
408 trace[i].set_x3_lt(x3_lt);
409 let row = ArithEqLtTableSM::calculate_table_row(
410 prev_x3_lt,
411 x3_lt,
412 data.x3[i] - BN254_PRIME_CHUNKS[i],
413 );
414 self.std.inc_virtual_row(self.table_id, row as u64, 1);
415 prev_x3_lt = x3_lt;
416
417 let y3_lt = data.y3[i] < BN254_PRIME_CHUNKS[i]
418 || (data.y3[i] == BN254_PRIME_CHUNKS[i] && prev_y3_lt);
419 trace[i].set_y3_lt(y3_lt);
420 let row = ArithEqLtTableSM::calculate_table_row(
421 prev_y3_lt,
422 y3_lt,
423 data.y3[i] - BN254_PRIME_CHUNKS[i],
424 );
425 self.std.inc_virtual_row(self.table_id, row as u64, 1);
426 prev_y3_lt = y3_lt;
427 }
428 _ => {
429 trace[i].set_x3_lt(false);
430 trace[i].set_y3_lt(false);
431 }
432 }
433 if (sel_op == SEL_OP_SECP256K1_ADD) || (sel_op == SEL_OP_BN254_CURVE_ADD) {
434 if x1_x2_different {
435 trace[i].set_x_are_different(true);
436 trace[i].set_x_delta_chunk_inv(0);
437 } else if data.x1[i] != data.x2[i] {
438 x1_x2_different = true;
439 trace[i].set_x_are_different(true);
440 trace[i].set_x_delta_chunk_inv(
441 F::inverse(&F::from_i64(data.x2[i] - data.x1[i])).as_canonical_u64(),
442 );
443 } else {
444 trace[i].set_x_are_different(false);
445 trace[i].set_x_delta_chunk_inv(0);
446 }
447 } else {
448 trace[i].set_x_are_different(false);
449 trace[i].set_x_delta_chunk_inv(0);
450 }
451 }
452 }
453
454 /// Computes the witness for a series of inputs and produces an `AirInstance`.
455 ///
456 /// # Arguments
457 /// * `inputs` - A slice of operations to process.
458 ///
459 /// # Returns
460 /// An `AirInstance` containing the computed witness data.
461 pub fn compute_witness(
462 &self,
463 _sctx: &SetupCtx<F>,
464 inputs: &[Vec<ArithEqInput>],
465 trace_buffer: Vec<F>,
466 ) -> ProofmanResult<AirInstance<F>> {
467 let mut trace = ArithEqTraceType::new_from_vec(trace_buffer)?;
468 let num_rows = trace.num_rows();
469 let total_inputs: usize = inputs.iter().map(|x| x.len()).sum();
470 let num_rows_needed = total_inputs * ARITH_EQ_ROWS_BY_OP;
471
472 tracing::debug!(
473 "··· Creating ArithEq instance [{} / {} rows filled {:.2}%]",
474 num_rows_needed,
475 num_rows,
476 num_rows_needed as f64 / num_rows as f64 * 100.0
477 );
478
479 timer_start_trace!(ARITH_EQ_TRACE);
480
481 let mut trace_rows = &mut trace.buffer[..];
482 let mut par_traces = Vec::new();
483 let mut inputs_indexes = Vec::new();
484 for (i, inputs) in inputs.iter().enumerate() {
485 for (j, _) in inputs.iter().enumerate() {
486 let (head, tail) = trace_rows.split_at_mut(ARITH_EQ_ROWS_BY_OP);
487 par_traces.push(head);
488 inputs_indexes.push((i, j));
489 trace_rows = tail;
490 }
491 }
492 let index = par_traces.len();
493
494 par_traces.into_par_iter().enumerate().for_each(|(index, trace)| {
495 let input_index = inputs_indexes[index];
496 let input = &inputs[input_index.0][input_index.1];
497 match input {
498 ArithEqInput::Arith256(idata) => self.process_arith256(idata, trace),
499 ArithEqInput::Arith256Mod(idata) => self.process_arith256_mod(idata, trace),
500 ArithEqInput::Secp256k1Add(idata) => self.process_secp256k1_add(idata, trace),
501 ArithEqInput::Secp256k1Dbl(idata) => self.process_secp256k1_dbl(idata, trace),
502 ArithEqInput::Bn254CurveAdd(idata) => self.process_bn254_curve_add(idata, trace),
503 ArithEqInput::Bn254CurveDbl(idata) => self.process_bn254_curve_dbl(idata, trace),
504 ArithEqInput::Bn254ComplexAdd(idata) => {
505 self.process_bn254_complex_add(idata, trace);
506 }
507 ArithEqInput::Bn254ComplexSub(idata) => {
508 self.process_bn254_complex_sub(idata, trace);
509 }
510 ArithEqInput::Bn254ComplexMul(idata) => {
511 self.process_bn254_complex_mul(idata, trace);
512 }
513 }
514 });
515
516 let padding_ops = (self.num_available_ops - index) as u64;
517 self.std.range_check(self.q_hsc_range_id, 0, 3 * padding_ops);
518 self.std.range_check(self.chunk_range_id, 0, 157 * padding_ops);
519 self.std.range_check(self.carry_range_id, 0, 96 * padding_ops);
520
521 let padding_row = ArithEqTraceRowType::default();
522
523 trace.buffer[num_rows_needed..num_rows].par_iter_mut().for_each(|slot| *slot = padding_row);
524
525 timer_stop_and_log_trace!(ARITH_EQ_TRACE);
526
527 Ok(AirInstance::new_from_trace(FromTrace::new(&mut trace)))
528 }
529}