Source code for zisk/state-machines/binary/src/binary_add.rs
1//! The `BinaryAddSM` module implements the logic for the Binary Add State Machine.
2//!
3//! This state machine processes binary-related operations.
4
5use crate::BinaryBasicFrops;
6use fields::PrimeField64;
7use pil_std_lib::Std;
8use proofman_common::{AirInstance, FromTrace, ProofmanResult};
9use rayon::prelude::*;
10use std::sync::Arc;
11use zisk_pil::BinaryAddAirValues;
12#[cfg(not(feature = "packed"))]
13use zisk_pil::{BinaryAddTrace, BinaryAddTraceRow};
14#[cfg(feature = "packed")]
15use zisk_pil::{BinaryAddTracePacked, BinaryAddTraceRowPacked};
16
17#[cfg(feature = "packed")]
18type BinaryAddTraceRowType<F> = BinaryAddTraceRowPacked<F>;
19#[cfg(feature = "packed")]
20type BinaryAddTraceType<F> = BinaryAddTracePacked<F>;
21
22#[cfg(not(feature = "packed"))]
23type BinaryAddTraceRowType<F> = BinaryAddTraceRow<F>;
24#[cfg(not(feature = "packed"))]
25type BinaryAddTraceType<F> = BinaryAddTrace<F>;
26
27const MASK_U32: u64 = 0x0000_0000_FFFF_FFFF;
28
29/// The `BinaryAddSM` struct encapsulates the logic of the Binary Add State Machine.
30pub struct BinaryAddSM<F: PrimeField64> {
31 /// Reference to the PIL2 standard library.
32 std: Arc<Std<F>>,
33 range_id: usize,
34
35 /// The table ID for the FROPS
36 frops_table_id: usize,
37}
38
39impl<F: PrimeField64> BinaryAddSM<F> {
40 /// Creates a new BinaryAdd State Machine instance.
41 ///
42 /// # Arguments/// * `std` - An `Arc`-wrapped reference to the PIL2 standard library.
43 /// Machine.
44 ///
45 /// # Returns
46 /// A new `BinaryAddSM` instance.
47 pub fn new(std: Arc<Std<F>>) -> Arc<Self> {
48 let range_id = std.get_range_id(0, 0xFFFF, None).expect("Failed to get range ID");
49
50 // Get the Arithmetic FROPS table ID
51 let frops_table_id = std
52 .get_virtual_table_id(BinaryBasicFrops::TABLE_ID)
53 .expect("Failed to get FROPS table ID");
54 // Create the BinaryAdd state machine
55 Arc::new(Self { std, range_id, frops_table_id })
56 }
57
58 /// Processes a slice of operation data, generating a trace row and updating multiplicities.
59 ///
60 /// # Arguments
61 /// * `operation` - The operation data to process.
62 /// * `multiplicity` - A mutable slice to update with multiplicities for the operation.
63 ///
64 /// # Returns
65 /// A `BinaryAddTraceRow` representing the operation's result.
66 #[inline(always)]
67 pub fn process_slice(&self, input: &[u64; 2]) -> (BinaryAddTraceRowType<F>, [u64; 4]) {
68 // Create an empty trace
69 let mut row: BinaryAddTraceRowType<F> = Default::default();
70
71 // Execute the opcode
72 let mut a = input[0];
73 let mut b = input[1];
74 let mut cin = 0;
75
76 let mut range_checks = [0u64; 4];
77 for i in 0..2 {
78 let _a = a & 0xFFFF_FFFF;
79 let _b = b & 0xFFFF_FFFF;
80 let c = _a + _b + cin;
81 let _c = c & 0xFFFF_FFFF;
82 row.set_a(i, _a as u32);
83 row.set_b(i, _b as u32);
84 let c_chunks = [_c & 0xFFFF, _c >> 16];
85 row.set_c_chunks(i * 2, c_chunks[0] as u16);
86 row.set_c_chunks(i * 2 + 1, c_chunks[1] as u16);
87 if c > MASK_U32 {
88 row.set_cout(i, true);
89 cin = 1
90 } else {
91 row.set_cout(i, false);
92 cin = 0
93 };
94 range_checks[i * 2] = c_chunks[0];
95 range_checks[i * 2 + 1] = c_chunks[1];
96 a >>= 32;
97 b >>= 32;
98 }
99
100 // Return
101 (row, range_checks)
102 }
103
104 /// Computes the witness for a series of inputs and produces an `AirInstance`.
105 ///
106 /// # Arguments
107 /// * `operations` - A slice of operations to process.
108 ///
109 /// # Returns
110 /// An `AirInstance` containing the computed witness data.
111 pub fn compute_witness(
112 &self,
113 inputs: &[Vec<[u64; 2]>],
114 trace_buffer: Vec<F>,
115 ) -> ProofmanResult<AirInstance<F>> {
116 let mut add_trace = BinaryAddTraceType::new_from_vec(trace_buffer)?;
117
118 let num_rows = add_trace.num_rows();
119
120 let total_inputs: usize = inputs.iter().map(|c| c.len()).sum();
121 assert!(total_inputs <= num_rows);
122
123 tracing::debug!(
124 "··· Creating BinaryAdd instance [{} / {} rows filled {:.2}%]",
125 total_inputs,
126 num_rows,
127 total_inputs as f64 / num_rows as f64 * 100.0
128 );
129
130 // Split the add_e_trace.buffer into slices matching each inner vector’s length.
131 let flat_inputs: Vec<_> = inputs.iter().flatten().collect();
132 let mut range_checks: Vec<[u64; 4]> = vec![[0u64; 4]; flat_inputs.len()];
133
134 // Process each slice in parallel, and use the corresponding inner input from `inputs`.
135 flat_inputs
136 .into_par_iter()
137 .zip(add_trace.buffer.par_iter_mut())
138 .zip(range_checks.par_iter_mut())
139 .for_each(|((input, trace_row), range_check)| {
140 let (row, checks) = self.process_slice(input);
141 *trace_row = row;
142 *range_check = checks;
143 });
144
145 let mut multiplicities = vec![0u32; 0xFFFF + 1];
146 for range_check in range_checks {
147 multiplicities[range_check[0] as usize] += 1;
148 multiplicities[range_check[1] as usize] += 1;
149 multiplicities[range_check[2] as usize] += 1;
150 multiplicities[range_check[3] as usize] += 1;
151 }
152 multiplicities[0] += 4 * (num_rows - total_inputs) as u32;
153
154 self.std.range_checks(self.range_id, multiplicities);
155
156 // Set 0 + 0 as the padding row
157 let padding_size = num_rows - total_inputs;
158 if padding_size > 0 {
159 let padding_row = BinaryAddTraceRowType::<F>::default();
160 add_trace.buffer[total_inputs..num_rows]
161 .par_iter_mut()
162 .for_each(|slot| *slot = padding_row);
163 }
164
165 let mut air_values = BinaryAddAirValues::<F>::new();
166 air_values.padding_size = F::from_usize(padding_size);
167 Ok(AirInstance::new_from_trace(
168 FromTrace::new(&mut add_trace).with_air_values(&mut air_values),
169 ))
170 }
171
172 pub fn compute_frops(&self, frops_inputs: &Vec<u32>) {
173 for row in frops_inputs {
174 self.std.inc_virtual_row(self.frops_table_id, *row as u64, 1);
175 }
176 }
177}