1use std::sync::Arc;
2
3use fields::PrimeField64;
4use rayon::prelude::*;
5
6use pil_std_lib::Std;
7use proofman_common::{AirInstance, FromTrace, ProofmanResult};
8use proofman_util::{timer_start_trace, timer_stop_and_log_trace};
9
10#[cfg(not(feature = "packed"))]
11use zisk_pil::{Add256Trace, Add256TraceRow};
12#[cfg(feature = "packed")]
13use zisk_pil::{Add256TracePacked, Add256TraceRowPacked};
14
15#[cfg(not(feature = "packed"))]
16type Add256TraceRowType<F> = Add256TraceRow<F>;
17#[cfg(feature = "packed")]
18type Add256TraceRowType<F> = Add256TraceRowPacked<F>;
19
20#[cfg(not(feature = "packed"))]
21type Add256TraceType<F> = Add256Trace<F>;
22#[cfg(feature = "packed")]
23type Add256TraceType<F> = Add256TracePacked<F>;
24
25use super::Add256Input;
26
27/// The `Add256SM` struct encapsulates the logic of the Add256 State Machine.
28pub struct Add256SM<F: PrimeField64> {
29 /// Reference to the PIL2 standard library.
30 pub std: Arc<Std<F>>,
31
32 /// Number of available add256s in the trace.
33 pub num_availables: usize,
34
35 /// Range checks ID's
36 range_id: usize,
37}
38
39impl<F: PrimeField64> Add256SM<F> {
40 /// Creates a new Add256 State Machine instance.
41 ///
42 /// # Returns
43 /// A new `Add256SM` instance.
44 pub fn new(std: Arc<Std<F>>) -> Arc<Self> {
45 // Compute some useful values
46 let num_availables = Add256TraceType::<F>::NUM_ROWS;
47
48 let range_id = std.get_range_id(0, (1 << 16) - 1, None).unwrap();
49
50 Arc::new(Self { std, num_availables, range_id })
51 }
52
53 /// Processes a slice of operation data, updating the trace.
54 ///
55 /// # Arguments
56 /// * `trace` - A mutable reference to the Add256 trace.
57 /// * `input` - The operation data to process.
58 #[inline(always)]
59 pub fn process_slice(
60 &self,
61 input: &Add256Input,
62 trace: &mut Add256TraceRowType<F>,
63 multiplicities: &mut [u32],
64 ) {
65 trace.set_cin(input.cin != 0);
66 let mut cout_2 = input.cin as u32;
67
68 for i in 0..4 {
69 let al = input.a[i] as u32;
70 let ah = (input.a[i] >> 32) as u32;
71
72 let bl = input.b[i] as u32;
73 let bh = (input.b[i] >> 32) as u32;
74
75 trace.set_a(i, 0, al);
76 trace.set_a(i, 1, ah);
77 trace.set_b(i, 0, bl);
78 trace.set_b(i, 1, bh);
79 let cl = al as u64 + bl as u64 + cout_2 as u64;
80 let cout_1 = cl >> 32;
81 let ch = ah as u64 + bh as u64 + cout_1;
82 cout_2 = (ch >> 32) as u32;
83
84 let cll = cl as u16;
85 let clh = (cl >> 16) as u16;
86 let chl = ch as u16;
87 let chh = (ch >> 16) as u16;
88
89 trace.set_c_chunks(i, 0, cll);
90 trace.set_c_chunks(i, 1, clh);
91 trace.set_c_chunks(i, 2, chl);
92 trace.set_c_chunks(i, 3, chh);
93
94 trace.set_cout(i, 0, cout_1 != 0);
95 trace.set_cout(i, 1, cout_2 != 0);
96
97 multiplicities[cll as usize] += 1;
98 multiplicities[clh as usize] += 1;
99 multiplicities[chl as usize] += 1;
100 multiplicities[chh as usize] += 1;
101 }
102 trace.set_addr_params(input.addr_main);
103 trace.set_addr_a(input.addr_a);
104 trace.set_addr_b(input.addr_b);
105 trace.set_addr_c(input.addr_c);
106 trace.set_step(input.step_main);
107 trace.set_sel(true);
108 }
109
110 /// Computes the witness for a series of inputs and produces an `AirInstance`.
111 ///
112 /// # Arguments
113 /// * `sctx` - The setup context containing the setup data.
114 /// * `inputs` - A slice of operations to process.
115 ///
116 /// # Returns
117 /// An `AirInstance` containing the computed witness data.
118 pub fn compute_witness(
119 &self,
120 inputs: &[Vec<Add256Input>],
121 trace_buffer: Vec<F>,
122 ) -> ProofmanResult<AirInstance<F>> {
123 let mut trace = Add256TraceType::<F>::new_from_vec(trace_buffer)?;
124
125 let num_rows = trace.num_rows();
126
127 let total_inputs: usize = inputs.iter().map(|c| c.len()).sum();
128 assert!(total_inputs <= num_rows);
129
130 tracing::debug!(
131 "··· Creating Add256 instance [{} / {} rows filled {:.2}%]",
132 total_inputs,
133 num_rows,
134 total_inputs as f64 / num_rows as f64 * 100.0
135 );
136
137 timer_start_trace!(ADD256_TRACE);
138
139 // Split the add256_trace.buffer into slices matching each inner vector’s length.
140 let flat_inputs: Vec<_> = inputs.iter().flatten().collect();
141 let trace_rows = trace.buffer.as_mut_slice();
142
143 // Determinar tamaño óptimo de chunks
144 let num_threads = rayon::current_num_threads();
145 let chunk_size = std::cmp::max(1, flat_inputs.len() / num_threads);
146
147 // Procesar en chunks para compartir arrays locales de multiplicities
148 let local_multiplicities_vec: Vec<Vec<u32>> = flat_inputs
149 .par_chunks(chunk_size)
150 .zip(trace_rows.par_chunks_mut(chunk_size))
151 .map(|(input_chunk, trace_chunk)| {
152 // Array local compartido por este chunk
153 let mut local_multiplicities = vec![0u32; 1 << 16];
154
155 // Procesar todos los inputs del chunk
156 for (input, trace_row) in input_chunk.iter().zip(trace_chunk.iter_mut()) {
157 self.process_slice(input, trace_row, &mut local_multiplicities);
158 }
159
160 local_multiplicities
161 })
162 .collect();
163
164 // Sumar todos los arrays locales en uno global
165 let mut global_multiplicities = vec![0u32; 1 << 16];
166 for local_multiplicities in local_multiplicities_vec {
167 for (i, count) in local_multiplicities.iter().enumerate() {
168 global_multiplicities[i] += count;
169 }
170 }
171
172 // Enviar el resultado final al std
173 self.std.range_checks(self.range_id, global_multiplicities);
174
175 timer_stop_and_log_trace!(ADD256_TRACE);
176
177 let padding_row = Add256TraceRowType::<F>::default();
178 trace.buffer[total_inputs..num_rows].par_iter_mut().for_each(|slot| *slot = padding_row);
179
180 Ok(AirInstance::<F>::new_from_trace(FromTrace::new(&mut trace)))
181 }
182}