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