1//! The `BinaryExtensionSM` module defines the Binary Extension State Machine.
2//!
3//! This state machine handles binary extension-related operations, computes traces, and manages
4//! range checks and multiplicities for table rows based on the operations provided.
5
6use std::sync::Arc;
7
8use crate::{
9 binary_constants::*, BinaryExtensionFrops, BinaryExtensionTableOp, BinaryExtensionTableSM,
10 BinaryInput,
11};
12
13use fields::PrimeField64;
14use pil_std_lib::Std;
15use proofman_common::{AirInstance, FromTrace, ProofmanResult};
16use rayon::prelude::*;
17use zisk_core::zisk_ops::ZiskOp;
18use zisk_pil::BinaryExtensionAirValues;
19#[cfg(not(feature = "packed"))]
20use zisk_pil::{BinaryExtensionTrace, BinaryExtensionTraceRow};
21#[cfg(feature = "packed")]
22use zisk_pil::{BinaryExtensionTracePacked, BinaryExtensionTraceRowPacked};
23
24#[cfg(feature = "packed")]
25type BinaryExtensionTraceRowType<F> = BinaryExtensionTraceRowPacked<F>;
26#[cfg(feature = "packed")]
27type BinaryExtensionTraceType<F> = BinaryExtensionTracePacked<F>;
28
29#[cfg(not(feature = "packed"))]
30type BinaryExtensionTraceRowType<F> = BinaryExtensionTraceRow<F>;
31#[cfg(not(feature = "packed"))]
32type BinaryExtensionTraceType<F> = BinaryExtensionTrace<F>;
33
34// Constants for bit masks and operations.
35const MASK_32: u64 = 0xFFFFFFFF;
36const MASK_64: u64 = 0xFFFFFFFFFFFFFFFF;
37
38const SE_MASK_32: u64 = 0xFFFFFFFF00000000;
39const SE_MASK_16: u64 = 0xFFFFFFFFFFFF0000;
40const SE_MASK_8: u64 = 0xFFFFFFFFFFFFFF00;
41
42const SIGN_32_BIT: u64 = 0x80000000;
43const SIGN_BYTE: u64 = 0x80;
44
45const LS_5_BITS: u64 = 0x1F;
46const LS_6_BITS: u64 = 0x3F;
47
48/// The `BinaryExtensionSM` struct defines the Binary Extension State Machine.
49///
50/// It processes binary extension-related operations and generates necessary traces and multiplicity
51/// tables for the operations. It also manages range checks through the PIL2 standard library.
52pub struct BinaryExtensionSM<F: PrimeField64> {
53 /// Reference to the PIL2 standard library.
54 std: Arc<Std<F>>,
55
56 /// The range check ID
57 range_id: usize,
58
59 /// The table ID for the Binary Basic State Machine
60 table_id: usize,
61
62 /// The table ID for the Binary Extension FROPS
63 frops_table_id: usize,
64}
65
66impl<F: PrimeField64> BinaryExtensionSM<F> {
67 /// Creates a new instance of the `BinaryExtensionSM`.
68 ///
69 /// # Arguments
70 /// * `std` - An `Arc`-wrapped reference to the PIL2 standard library.
71 ///
72 /// # Returns
73 /// An `Arc`-wrapped instance of `BinaryExtensionSM`.
74 pub fn new(std: Arc<Std<F>>) -> Arc<Self> {
75 // Get the range check ID
76 let range_id = std.get_range_id(0, 0xFFFFFF, None).expect("Failed to get range ID");
77
78 // Get the table ID
79 let table_id = std
80 .get_virtual_table_id(BinaryExtensionTableSM::TABLE_ID)
81 .expect("Failed to get table ID");
82
83 // Get the FROPS table ID
84 let frops_table_id = std
85 .get_virtual_table_id(BinaryExtensionFrops::TABLE_ID)
86 .expect("Failed to get FROPS table ID");
87
88 Arc::new(Self { std, range_id, table_id, frops_table_id })
89 }
90
91 /// Determines if the given opcode represents a shift operation.
92 fn opcode_is_shift(opcode: ZiskOp) -> bool {
93 match opcode {
94 ZiskOp::Sll
95 | ZiskOp::Srl
96 | ZiskOp::Sra
97 | ZiskOp::SllW
98 | ZiskOp::SrlW
99 | ZiskOp::SraW => true,
100
101 ZiskOp::SignExtendB | ZiskOp::SignExtendH | ZiskOp::SignExtendW => false,
102
103 _ => panic!("BinaryExtensionSM::opcode_is_shift() got invalid opcode={opcode:?}"),
104 }
105 }
106
107 /// Determines if the given opcode represents a shift word operation.
108 fn opcode_is_shift_word(opcode: ZiskOp) -> bool {
109 match opcode {
110 ZiskOp::SllW | ZiskOp::SrlW | ZiskOp::SraW => true,
111
112 ZiskOp::Sll
113 | ZiskOp::Srl
114 | ZiskOp::Sra
115 | ZiskOp::SignExtendB
116 | ZiskOp::SignExtendH
117 | ZiskOp::SignExtendW => false,
118
119 _ => panic!("BinaryExtensionSM::opcode_is_shift() got invalid opcode={opcode:?}"),
120 }
121 }
122
123 /// Processes a single operation and generates the corresponding trace row.
124 ///
125 /// # Arguments
126 /// * `operation` - The operation to process.
127 /// * `multiplicity` - A mutable reference to the multiplicity table to update.
128 /// * `range_check` - A mutable reference to the range check table to update.
129 ///
130 /// # Returns
131 /// A `BinaryExtensionTraceRow` representing the processed trace.
132 pub fn process_slice(&self, input: &BinaryInput) -> BinaryExtensionTraceRowType<F> {
133 // Get a ZiskOp from the code
134 let opcode = ZiskOp::try_from_code(input.op).expect("Invalid ZiskOp opcode");
135
136 // Create an empty trace
137 let mut row = BinaryExtensionTraceRowType::default();
138 row.set_op(input.op);
139
140 // Set if the opcode is a shift operation
141 let op_is_shift = Self::opcode_is_shift(opcode);
142 row.set_op_is_shift(op_is_shift);
143
144 // Set if the opcode is a shift word operation
145 let op_is_shift_word = Self::opcode_is_shift_word(opcode);
146
147 // Detect if this is a sign extend operation
148 let a_val = if op_is_shift { input.a } else { input.b };
149 let b_val = if op_is_shift { input.b } else { input.a };
150
151 // Split a in bytes and store them in in1
152 let a_bytes: [u8; 8] = a_val.to_le_bytes();
153 for (i, value) in a_bytes.iter().enumerate() {
154 row.set_free_in_a(i, *value);
155 }
156
157 // Store b low part into in2_low
158 let in2_low: u64 = if op_is_shift { b_val & 0xFF } else { 0 };
159 row.set_free_in_b(in2_low as u8);
160
161 // Store b lower bits when shifting, depending on operation size
162 let b_low = if op_is_shift_word { b_val & LS_5_BITS } else { b_val & LS_6_BITS };
163
164 // Store b into in2
165 let in2_0: u32 = if op_is_shift {
166 ((b_val >> 8) & 0xFFFFFF) as u32
167 } else {
168 (b_val & 0xFFFFFFFF) as u32
169 };
170 let in2_1: u32 = ((b_val >> 32) & 0xFFFFFFFF) as u32;
171
172 row.set_b(0, in2_0);
173 row.set_b(1, in2_1);
174
175 // Calculate the trace output
176 let mut t_out: [[u64; 2]; 8] = [[0; 2]; 8];
177
178 // Calculate output based on opcode
179 let binary_extension_table_op: BinaryExtensionTableOp;
180 match opcode {
181 ZiskOp::Sll => {
182 binary_extension_table_op = BinaryExtensionTableOp::Sll;
183 for j in 0..8 {
184 let bits_to_shift = b_low + 8 * j as u64;
185 let out =
186 if bits_to_shift < 64 { (a_bytes[j] as u64) << bits_to_shift } else { 0 };
187 t_out[j][0] = out & 0xffffffff;
188 t_out[j][1] = (out >> 32) & 0xffffffff;
189 }
190 }
191 ZiskOp::Srl => {
192 binary_extension_table_op = BinaryExtensionTableOp::Srl;
193 for j in 0..8 {
194 let out = ((a_bytes[j] as u64) << (8 * j as u64)) >> b_low;
195 t_out[j][0] = out & 0xffffffff;
196 t_out[j][1] = (out >> 32) & 0xffffffff;
197 }
198 }
199 ZiskOp::Sra => {
200 binary_extension_table_op = BinaryExtensionTableOp::Sra;
201 for j in 0..8 {
202 let mut out = ((a_bytes[j] as u64) << (8 * j as u64)) >> b_low;
203 if j == 7 {
204 // most significant bit of most significant byte define if negative or not
205 // if negative then add b bits one on the left
206 if ((a_bytes[j] as u64) & SIGN_BYTE) != 0 && (b_low != 0) {
207 out |= MASK_64 << (64 - b_low);
208 }
209 }
210 t_out[j][0] = out & 0xffffffff;
211 t_out[j][1] = (out >> 32) & 0xffffffff;
212 }
213 }
214 ZiskOp::SllW => {
215 binary_extension_table_op = BinaryExtensionTableOp::SllW;
216 for j in 0..8 {
217 let mut out: u64;
218 if j >= 4 {
219 out = 0;
220 } else {
221 out = (((a_bytes[j] as u64) << b_low) << (8 * j as u64)) & MASK_32;
222 if (out & SIGN_32_BIT) != 0 {
223 out |= SE_MASK_32;
224 }
225 }
226 t_out[j][0] = out & 0xffffffff;
227 t_out[j][1] = (out >> 32) & 0xffffffff;
228 }
229 }
230 ZiskOp::SrlW => {
231 binary_extension_table_op = BinaryExtensionTableOp::SrlW;
232 for j in 0..8 {
233 let mut out: u64;
234 if j >= 4 {
235 out = 0;
236 } else {
237 out = (((a_bytes[j] as u64) << (8 * j as u64)) >> b_low) & MASK_32;
238 if (out & SIGN_32_BIT) != 0 {
239 out |= SE_MASK_32;
240 }
241 }
242 t_out[j][0] = out & 0xffffffff;
243 t_out[j][1] = (out >> 32) & 0xffffffff;
244 }
245 }
246 ZiskOp::SraW => {
247 binary_extension_table_op = BinaryExtensionTableOp::SraW;
248 for j in 0..8 {
249 let mut out: u64;
250 if j >= 4 {
251 out = 0;
252 } else {
253 out = ((a_bytes[j] as u64) << (8 * j as u64)) >> b_low;
254 if j == 3 && ((a_bytes[j] as u64) & SIGN_BYTE) != 0 {
255 out |= MASK_64 << (32 - b_low);
256 }
257 }
258 t_out[j][0] = out & 0xffffffff;
259 t_out[j][1] = (out >> 32) & 0xffffffff;
260 }
261 }
262 ZiskOp::SignExtendB => {
263 binary_extension_table_op = BinaryExtensionTableOp::SextB;
264 for j in 0..8 {
265 let out: u64;
266 if j == 0 {
267 if ((a_bytes[j] as u64) & SIGN_BYTE) != 0 {
268 out = (a_bytes[j] as u64) | SE_MASK_8;
269 } else {
270 out = a_bytes[j] as u64;
271 }
272 } else {
273 out = 0;
274 }
275 t_out[j][0] = out & 0xffffffff;
276 t_out[j][1] = (out >> 32) & 0xffffffff;
277 }
278 }
279 ZiskOp::SignExtendH => {
280 binary_extension_table_op = BinaryExtensionTableOp::SextH;
281 for j in 0..8 {
282 let out: u64;
283 if j == 0 {
284 out = a_bytes[j] as u64;
285 } else if j == 1 {
286 if ((a_bytes[j] as u64) & SIGN_BYTE) != 0 {
287 out = ((a_bytes[j] as u64) << 8) | SE_MASK_16;
288 } else {
289 out = (a_bytes[j] as u64) << 8;
290 }
291 } else {
292 out = 0;
293 }
294 t_out[j][0] = out & 0xffffffff;
295 t_out[j][1] = (out >> 32) & 0xffffffff;
296 }
297 }
298 ZiskOp::SignExtendW => {
299 binary_extension_table_op = BinaryExtensionTableOp::SextW;
300 for j in 0..4 {
301 let mut out = (a_bytes[j] as u64) << (8 * j as u64);
302 if j == 3 && ((a_bytes[j] as u64) & SIGN_BYTE) != 0 {
303 out |= SE_MASK_32;
304 }
305
306 t_out[j][0] = out & 0xffffffff;
307 t_out[j][1] = (out >> 32) & 0xffffffff;
308 }
309 }
310 _ => panic!("BinaryExtensionSM::process_slice() found invalid opcode={}", input.op),
311 }
312
313 // Convert the trace output to field elements
314 for (j, out) in t_out.iter().enumerate() {
315 row.set_free_in_c(j, 0, out[0] as u32);
316 row.set_free_in_c(j, 1, out[1] as u32);
317 }
318
319 for (i, a_byte) in a_bytes.iter().enumerate() {
320 let row = BinaryExtensionTableSM::calculate_table_row(
321 binary_extension_table_op,
322 i as u64,
323 *a_byte as u64,
324 in2_low,
325 );
326 self.std.inc_virtual_row(self.table_id, row, 1);
327 }
328
329 row
330 }
331
332 /// Computes the witness for the given set of operations.
333 ///
334 /// # Arguments
335 /// * `operations` - The list of operations to process.
336 ///
337 /// # Returns
338 /// An `AirInstance` representing the computed witness.
339 pub fn compute_witness(
340 &self,
341 inputs: &[Vec<BinaryInput>],
342 trace_buffer: Vec<F>,
343 ) -> ProofmanResult<AirInstance<F>> {
344 let mut binary_e_trace = BinaryExtensionTraceType::new_from_vec(trace_buffer)?;
345
346 let num_rows = binary_e_trace.num_rows();
347
348 let total_inputs: usize = inputs.iter().map(|c| c.len()).sum();
349 assert!(
350 total_inputs <= num_rows,
351 "{} <= {} ({})",
352 total_inputs,
353 num_rows,
354 BinaryExtensionTraceType::<F>::NUM_ROWS
355 );
356
357 tracing::debug!(
358 "··· Creating Binary Extension instance [{} / {} rows filled {:.2}%]",
359 total_inputs,
360 num_rows,
361 total_inputs as f64 / num_rows as f64 * 100.0
362 );
363
364 // Split the binary_e_trace.buffer into slices matching each inner vector’s length.
365 let sizes: Vec<usize> = inputs.iter().map(|v| v.len()).collect();
366 let mut slices = Vec::with_capacity(inputs.len());
367 let mut rest = &mut binary_e_trace.buffer[..];
368 for size in sizes {
369 let (head, tail) = rest.split_at_mut(size);
370 slices.push(head);
371 rest = tail;
372 }
373
374 // Process each slice in parallel, and use the corresponding inner input from `inputs`.
375 slices.into_par_iter().enumerate().for_each(|(i, slice)| {
376 slice.iter_mut().enumerate().for_each(|(j, trace_row)| {
377 *trace_row = self.process_slice(&inputs[i][j]);
378 });
379 });
380
381 // Iterate over all inputs and check opcode
382 // to update multiplicity for the corresponding table row.
383 for row in inputs.iter() {
384 for input in row.iter() {
385 let opcode = ZiskOp::try_from_code(input.op).expect("Invalid ZiskOp opcode");
386 let op_is_shift = Self::opcode_is_shift(opcode);
387 if op_is_shift {
388 let row = (input.b >> 8) & 0xFFFFFF;
389 self.std.range_check(self.range_id, row as i64, 1);
390 }
391 }
392 }
393
394 // Set SEXT_B(0) as the padding row
395 let mut padding_row = BinaryExtensionTraceRowType::default();
396 padding_row.set_op(SEXT_B_OP);
397
398 binary_e_trace.buffer[total_inputs..num_rows]
399 .par_iter_mut()
400 .for_each(|slot| *slot = padding_row);
401
402 let padding_size = num_rows - total_inputs;
403 for i in 0..8 {
404 let multiplicity = padding_size as u64;
405 let row =
406 BinaryExtensionTableSM::calculate_table_row(BinaryExtensionTableOp::SextB, i, 0, 0);
407 self.std.inc_virtual_row(self.table_id, row, multiplicity);
408 }
409
410 let mut air_values = BinaryExtensionAirValues::<F>::new();
411 air_values.padding_size = F::from_usize(padding_size);
412 Ok(AirInstance::new_from_trace(
413 FromTrace::new(&mut binary_e_trace).with_air_values(&mut air_values),
414 ))
415 }
416 pub fn compute_frops(&self, frops_inputs: &Vec<u32>) {
417 for row in frops_inputs {
418 self.std.inc_virtual_row(self.frops_table_id, *row as u64, 1);
419 }
420 }
421}