1use std::sync::Arc;
2
3use fields::PrimeField64;
4use pil_std_lib::Std;
5
6use proofman_common::{AirInstance, FromTrace, ProofmanResult};
7use proofman_util::{timer_start_trace, timer_stop_and_log_trace};
8
9#[cfg(not(feature = "packed"))]
10use zisk_pil::{KeccakfTrace, KeccakfTraceRow};
11#[cfg(feature = "packed")]
12use zisk_pil::{KeccakfTracePacked, KeccakfTraceRowPacked};
13
14#[cfg(feature = "packed")]
15type KeccakfTraceType<F> = KeccakfTracePacked<F>;
16#[cfg(feature = "packed")]
17type KeccakfTraceRowType<F> = KeccakfTraceRowPacked<F>;
18
19#[cfg(not(feature = "packed"))]
20type KeccakfTraceType<F> = KeccakfTrace<F>;
21#[cfg(not(feature = "packed"))]
22type KeccakfTraceRowType<F> = KeccakfTraceRow<F>;
23
24use precompiles_helpers::{keccak_f_rounds, keccakf_state_from_linear, keccakf_state_to_linear_1d};
25
26use crate::KeccakfInput;
27
28use super::{keccakf_constants::*, KeccakfTableSM};
29
30use rayon::prelude::*;
31
32/// The `KeccakfSM` struct encapsulates the logic of the Keccakf State Machine.
33pub struct KeccakfSM<F: PrimeField64> {
34 /// Number of available keccakfs in the trace.
35 pub num_available_keccakfs: usize,
36
37 /// Reference to the PIL2 standard library.
38 std: Arc<Std<F>>,
39
40 /// The table ID for the Keccakf Table State Machine
41 table_id: usize,
42}
43
44impl<F: PrimeField64> KeccakfSM<F> {
45 const NUM_REM: usize = WIDTH % TABLE_CHUNK_SIZE;
46 const NUM_REDUCED: usize = (WIDTH - Self::NUM_REM) / TABLE_CHUNK_SIZE;
47
48 /// Creates a new Keccakf State Machine instance.
49 ///
50 /// # Arguments
51 /// * `keccakf_table_sm` - An `Arc`-wrapped reference to the Keccakf Table State Machine.
52 ///
53 /// # Returns
54 /// A new `KeccakfSM` instance.
55 pub fn new(std: Arc<Std<F>>) -> Arc<Self> {
56 // Compute some useful values
57 let num_non_usable_rows = KeccakfTraceType::<F>::NUM_ROWS % ROWS_BY_KECCAKF;
58 let num_available_keccakfs = if num_non_usable_rows == 0 {
59 KeccakfTraceType::<F>::NUM_ROWS / ROWS_BY_KECCAKF
60 } else {
61 // Subtract 1 because we can't fit a complete cycle in the remaining rows
62 (KeccakfTraceType::<F>::NUM_ROWS - num_non_usable_rows) / ROWS_BY_KECCAKF - 1
63 };
64
65 // Get the table ID
66 let table_id = std
67 .get_virtual_table_id(KeccakfTableSM::TABLE_ID)
68 .expect("Failed to get Keccakf table ID");
69
70 Arc::new(Self { num_available_keccakfs, std, table_id })
71 }
72
73 /// Processes a slice of operation data, updating the trace and multiplicities.
74 ///
75 /// # Arguments
76 /// * `trace` - A mutable reference to the Keccakf trace.
77 /// * `input` - The operation data to process.
78 #[inline(always)]
79 #[allow(clippy::needless_range_loop)]
80 fn process_trace(
81 &self,
82 trace: &mut [KeccakfTraceRowType<F>],
83 initial_state: &[u64; 25],
84 addr: Option<u32>,
85 step: Option<u64>,
86 ) {
87 let lookup_active = addr.is_some() && step.is_some();
88
89 // Fill the states
90 // Convert input state to 5x5x64 representation
91 let initial_state = keccakf_state_from_linear(initial_state);
92 let round_states = keccak_f_rounds(initial_state);
93 for (state_3d, r) in round_states {
94 // Convert 3D state to 1D for processing
95 let state_1d = keccakf_state_to_linear_1d(&state_3d);
96
97 // Fill keccakf_state
98 for i in 0..1600 {
99 trace[r].set_state(i, (state_1d[i] % 2) == 1);
100 }
101
102 // Fill keccakf_reduced
103 for i in 0..Self::NUM_REDUCED {
104 let offset = i * TABLE_CHUNK_SIZE;
105 let mut acc = 0u32;
106 for j in 0..TABLE_CHUNK_SIZE {
107 let idx = offset + j;
108 let value = state_1d[idx] as u32;
109 acc += value * BASE.pow(j as u32);
110 }
111 if r > 0 {
112 trace[r - 1].set_chunk_acc(i, acc);
113 }
114 }
115
116 // Fill keccakf_rem
117 let offset = Self::NUM_REDUCED * TABLE_CHUNK_SIZE;
118 let mut acc = 0u8;
119 for j in 0..Self::NUM_REM {
120 let idx = offset + j;
121 let bit_value = state_1d[idx] as u8;
122 acc += bit_value * (BASE.pow(j as u32) as u8);
123 }
124 if r > 0 {
125 trace[r - 1].set_rem_acc(acc);
126 }
127 }
128
129 if !lookup_active {
130 return;
131 }
132
133 // Fill step and addr
134 trace[0].set_step_addr(step.unwrap_or(0));
135 trace[1].set_step_addr(addr.unwrap_or(0) as u64);
136
137 // Fill in_use_clk_0
138 trace[0].set_in_use_clk_0(true);
139
140 // Fill in_use
141 for i in 0..ROWS_BY_KECCAKF {
142 trace[i].set_in_use(true);
143 }
144 }
145
146 /// Computes the witness for a series of inputs and produces an `AirInstance`.
147 ///
148 /// # Arguments
149 /// * `inputs` - A slice of operations to process.
150 ///
151 /// # Returns
152 /// An `AirInstance` containing the computed witness data.
153 pub fn compute_witness(
154 &self,
155 inputs: &[Vec<KeccakfInput>],
156 trace_buffer: Vec<F>,
157 ) -> ProofmanResult<AirInstance<F>> {
158 let mut trace = KeccakfTraceType::new_from_vec_zeroes(trace_buffer)?;
159 let num_rows = trace.num_rows();
160
161 // Check that we can fit all the keccakfs in the trace
162 let num_available_keccakfs = self.num_available_keccakfs;
163 let num_inputs = inputs.iter().map(|v| v.len()).sum::<usize>();
164 let num_rows_needed = if num_inputs < num_available_keccakfs {
165 num_inputs * ROWS_BY_KECCAKF
166 } else if num_inputs == num_available_keccakfs {
167 num_rows
168 } else {
169 panic!(
170 "Exceeded available Keccakfs inputs: requested {}, but only {} are available.",
171 num_inputs, num_available_keccakfs
172 );
173 };
174
175 tracing::debug!(
176 "··· Creating Keccakf instance [{} / {} rows filled {:.2}%]",
177 num_rows_needed,
178 num_rows,
179 num_rows_needed as f64 / num_rows as f64 * 100.0
180 );
181
182 timer_start_trace!(KECCAKF_TRACE);
183
184 // 1] Fill the trace with the provided inputs
185 let mut trace_rows = &mut trace.buffer[..];
186 let mut par_traces = Vec::new();
187 let mut inputs_indexes = Vec::new();
188 for (i, inputs) in inputs.iter().enumerate() {
189 for (j, _) in inputs.iter().enumerate() {
190 let (head, tail) = trace_rows.split_at_mut(ROWS_BY_KECCAKF);
191 par_traces.push(head);
192 inputs_indexes.push((i, j));
193 trace_rows = tail;
194 }
195 }
196
197 par_traces.par_iter_mut().enumerate().for_each(|(index, trace)| {
198 let input_index = inputs_indexes[index];
199 let input = &inputs[input_index.0][input_index.1];
200 self.process_trace(trace, &input.state, Some(input.addr_main), Some(input.step_main));
201 });
202
203 // 2] Update lookup table
204 let mut table = vec![0u32; TABLE_SIZE as usize];
205 for trace in &par_traces {
206 for r in 1..ROWS_BY_KECCAKF {
207 for i in 0..Self::NUM_REDUCED {
208 let table_row =
209 KeccakfTableSM::calculate_table_row(trace[r - 1].get_chunk_acc(i));
210 table[table_row as usize] += 1;
211 }
212 let table_row =
213 KeccakfTableSM::calculate_table_row(trace[r - 1].get_rem_acc() as u32);
214 table[table_row as usize] += 1;
215 }
216 }
217 table.into_par_iter().enumerate().for_each(|(row, value)| {
218 if value > 0 {
219 self.std.inc_virtual_row(self.table_id, row as u64, value as u64);
220 }
221 });
222 timer_stop_and_log_trace!(KECCAKF_TRACE);
223
224 timer_start_trace!(KECCAKF_PADDING);
225
226 // 3] Fill the padding rows with Keccakf(0)
227 let padding_rows_start = num_rows_needed;
228 let padding_rows_end =
229 padding_rows_start + ((num_available_keccakfs - num_inputs) * ROWS_BY_KECCAKF);
230
231 // Split the padding trace into padding chunks
232 let padding_trace = &mut trace.buffer[padding_rows_start..padding_rows_end];
233 let mut padding_chunks: Vec<_> = padding_trace.chunks_mut(ROWS_BY_KECCAKF).collect();
234
235 // Process padding in parallel
236 if let Some((first, rest)) = padding_chunks.split_first_mut() {
237 self.process_trace(first, &[0u64; 25], None, None);
238
239 rest.par_iter_mut().for_each(|chunk| {
240 chunk.copy_from_slice(first);
241 });
242 }
243
244 // 4] The non-usable rows should be zeroes, which are already set at initialization
245
246 timer_stop_and_log_trace!(KECCAKF_PADDING);
247
248 Ok(AirInstance::new_from_trace(FromTrace::new(&mut trace)))
249 }
250}