1//! The `MainSM` module implements the Main State Machine,
2//! responsible for computing witness main state machine.
3//!
4//! Key components of this module include:
5//! - The `MainSM` struct, which handles the main execution trace computation.
6//! - The `MainInstance` struct, representing the execution context of a specific main trace
7//! segment.
8//! - Methods for computing the witness and setting up trace rows.
9
10use std::sync::Arc;
11
12use crate::MainCounter;
13use fields::PrimeField64;
14use mem_common::{MemHelpers, MEM_REGS_MAX_DIFF, MEM_STEPS_BY_MAIN_STEP};
15use pil_std_lib::Std;
16use proofman_common::{AirInstance, FromTrace, ProofCtx, ProofmanResult, SetupCtx};
17use rayon::prelude::*;
18use zisk_common::{BusDeviceMetrics, EmuTrace, InstanceCtx, SegmentId};
19use zisk_core::{ZiskRom, REGS_IN_MAIN, REGS_IN_MAIN_FROM, REGS_IN_MAIN_TO};
20use zisk_pil::MainAirValues;
21use ziskemu::{Emu, EmuRegTrace};
22
23#[cfg(not(feature = "packed"))]
24use zisk_pil::{MainTrace, MainTraceRow};
25#[cfg(feature = "packed")]
26use zisk_pil::{MainTracePacked, MainTraceRowPacked};
27
28#[cfg(feature = "packed")]
29type MainTraceRowType<F> = MainTraceRowPacked<F>;
30#[cfg(feature = "packed")]
31type MainTraceType<F> = MainTracePacked<F>;
32
33#[cfg(not(feature = "packed"))]
34type MainTraceRowType<F> = MainTraceRow<F>;
35#[cfg(not(feature = "packed"))]
36type MainTraceType<F> = MainTrace<F>;
37
38/// Represents an instance of the main state machine,
39/// containing context for managing a specific segment of the main trace.
40pub struct MainInstance<F: PrimeField64> {
41 /// Instance Context
42 pub ictx: InstanceCtx,
43
44 pub std: Arc<Std<F>>,
45}
46
47impl<F: PrimeField64> MainInstance<F> {
48 const MAX_SEGMENT_ID: usize = ((1 << 32) / MainTraceType::<F>::NUM_ROWS) - 1;
49
50 /// Creates a new `MainInstance`.
51 ///
52 /// # Arguments
53 /// * `ictx` - The instance context for this main instance.
54 ///
55 /// # Returns
56 /// A new `MainInstance`.
57 pub fn new(ictx: InstanceCtx, std: Arc<Std<F>>) -> Self {
58 Self { ictx, std }
59 }
60
61 /// Computes the main witness trace for a given segment based on the provided proof context,
62 /// ROM, and emulation traces.
63 ///
64 /// # Arguments
65 /// * `zisk_rom` - Reference to the Zisk ROM used for execution.
66 /// * `min_traces` - A vector of the minimal traces, each segment has num_within minimal traces
67 /// inside.
68 /// * `chunk_size` - The size of the minimal traces.
69 /// * `main_instance` - Reference to the `MainInstance` representing the current segment.
70 ///
71 /// The computed trace is added to the proof context's air instance repository.
72 pub fn compute_witness(
73 &self,
74 zisk_rom: &ZiskRom,
75 min_traces: &[EmuTrace],
76 chunk_size: u64,
77 main_instance: &MainInstance<F>,
78 trace_buffer: Vec<F>,
79 ) -> ProofmanResult<AirInstance<F>> {
80 // Create the main trace buffer
81 let mut main_trace = MainTraceType::new_from_vec(trace_buffer)?;
82
83 let segment_id = main_instance.ictx.plan.segment_id.unwrap();
84
85 let is_last_segment = main_instance
86 .ictx
87 .plan
88 .meta
89 .as_ref()
90 .and_then(|m| m.downcast_ref::<bool>())
91 .unwrap_or_else(|| {
92 panic!("create_main_instance: Invalid metadata format, expected bool")
93 });
94
95 // Determine the number of minimal traces per segment
96 let num_rows = main_trace.num_rows();
97 let num_within = num_rows / chunk_size as usize;
98
99 // Determine trace slice for the current segment
100 let start_idx = segment_id.as_usize() * num_within;
101 let end_idx = (start_idx + num_within).min(min_traces.len());
102 let segment_min_traces = &min_traces[start_idx..end_idx];
103
104 // Calculate total filled rows
105 let filled_rows: usize =
106 segment_min_traces.iter().map(|min_trace| min_trace.steps as usize).sum();
107
108 tracing::debug!(
109 "··· Creating Main segment #{} [{} / {} rows filled {:.2}%]",
110 segment_id,
111 filled_rows,
112 num_rows,
113 filled_rows as f64 / num_rows as f64 * 100.0
114 );
115
116 // Calculate final_step of instance, last mem slot of last row. The initial_step is 0 for the
117 // first segment, for the rest is the final_step of the previous segment
118
119 let last_row_previous_segment =
120 if segment_id == 0 { 0 } else { (segment_id.as_usize() * num_rows) as u64 - 1 };
121
122 let initial_step = MemHelpers::main_step_to_special_mem_step(last_row_previous_segment);
123
124 let final_step = MemHelpers::main_step_to_special_mem_step(
125 ((segment_id.as_usize() + 1) * num_rows) as u64 - 1,
126 );
127
128 // To reduce memory used, only take memory for the maximum range of mem_step inside the
129 // minimal trace.
130 let max_range = chunk_size * MEM_STEPS_BY_MAIN_STEP;
131
132 // We know each register's previous step, but only by instance. We don't have this
133 // information by chunk, so we need to store in the EmuRegTrace the location of the
134 // first mem_step register is used in the chunk and information about the last step
135 // where the register is used. The register's last steps of one chunk are the initial
136 // steps of the next chunk. In the end, we need to update with the correct values.
137
138 let fill_trace_outputs = main_trace
139 .par_iter_mut_chunks(num_within)
140 .enumerate()
141 .take(segment_min_traces.len())
142 .map(|(chunk_id, chunk)| {
143 let mut step_range_check = vec![0; max_range as usize];
144 let init_chunk_step = if chunk_id == 0 { initial_step } else { 0 };
145 let mut reg_trace = EmuRegTrace::from_init_step(init_chunk_step, chunk_id == 0);
146 let (pc, regs) = Self::fill_partial_trace(
147 zisk_rom,
148 chunk,
149 &segment_min_traces[chunk_id],
150 &mut reg_trace,
151 &mut step_range_check,
152 chunk_id == (end_idx - start_idx - 1),
153 );
154 (pc, regs, reg_trace, step_range_check)
155 })
156 .collect::<Vec<(u64, Vec<u64>, EmuRegTrace, Vec<u32>)>>();
157 let last_result = fill_trace_outputs.last().unwrap();
158 let next_pc = last_result.0;
159
160 let mut step_range_check: Vec<u32> = (0..max_range as usize)
161 .into_par_iter()
162 .map(|i| fill_trace_outputs.iter().map(|(_, _, _, local)| local[i]).sum())
163 .collect();
164
165 // In the range checks are values too large to store in steps_range_check, but there
166 // are only a few values that exceed this limit, for this reason, are stored in a vector
167
168 let mut reg_steps = [initial_step; REGS_IN_MAIN];
169 let mut large_range_checks = Self::complete_trace_with_initial_reg_steps_per_chunk(
170 num_rows,
171 &fill_trace_outputs,
172 &mut main_trace,
173 &mut step_range_check,
174 &mut reg_steps,
175 );
176
177 Self::update_reg_steps_with_last_chunk(&last_result.2, &mut reg_steps);
178
179 // Pad remaining rows with the last valid row
180 // In padding row must be clear of registers access, if not need to calculate previous
181 // register step and range check conntribution
182 let last_row = main_trace.buffer[filled_rows - 1];
183 main_trace.buffer[filled_rows..num_rows].par_iter_mut().for_each(|row| *row = last_row);
184
185 // Determine the last row of the previous segment
186 let prev_segment_last_c = if start_idx > 0 {
187 Emu::intermediate_value(min_traces[start_idx - 1].last_c)
188 } else {
189 [F::ZERO, F::ZERO]
190 };
191
192 // Prepare main AIR values
193 let mut air_values = MainAirValues::<F>::new();
194
195 air_values.main_segment = F::from_usize(segment_id.into());
196 air_values.main_last_segment = F::from_bool(*is_last_segment);
197 air_values.segment_initial_pc = F::from_u32(main_trace[0].get_pc());
198 air_values.segment_next_pc = F::from_u64(next_pc);
199 air_values.segment_previous_c = prev_segment_last_c;
200 air_values.segment_last_c[0] = F::from_u32(last_row.get_c(0));
201 air_values.segment_last_c[1] = F::from_u32(last_row.get_c(1));
202
203 Self::update_reg_airvalues(
204 &mut air_values,
205 final_step,
206 &last_result.1,
207 ®_steps,
208 &mut step_range_check,
209 &mut large_range_checks,
210 );
211 self.update_std_range_checks(segment_id, step_range_check, &large_range_checks);
212
213 // Generate and add the AIR instance
214 let from_trace = FromTrace::new(&mut main_trace).with_air_values(&mut air_values);
215 Ok(AirInstance::new_from_trace(from_trace))
216 }
217
218 /// Fills a partial trace in the main trace buffer based on the minimal trace.
219 /// This method processes the minimal trace in batches to improve performance.
220 ///
221 /// # Arguments
222 /// * `zisk_rom` - Reference to the Zisk ROM used for execution.
223 /// * `main_trace` - Reference to the main trace buffer to fill.
224 /// * `min_trace` - Reference to the minimal trace to process.
225 ///
226 /// # Returns
227 /// The next program counter value after processing the minimal trace.
228 fn fill_partial_trace(
229 zisk_rom: &ZiskRom,
230 main_trace: &mut [MainTraceRowType<F>],
231 min_trace: &EmuTrace,
232 reg_trace: &mut EmuRegTrace,
233 step_range_check: &mut [u32],
234 last_reg_values: bool,
235 ) -> (u64, Vec<u64>) {
236 // Initialize the emulator with the start state of the emu trace
237 let mut emu = Emu::from_emu_trace_start(zisk_rom, &min_trace.start_state);
238 let mut mem_reads_index: usize = 0;
239
240 for trace in main_trace {
241 *trace = emu.step_slice_full_trace(
242 &min_trace.mem_reads,
243 &mut mem_reads_index,
244 reg_trace,
245 Some(step_range_check),
246 );
247 }
248
249 (
250 emu.ctx.inst_ctx.pc,
251 if last_reg_values {
252 emu.ctx.inst_ctx.regs[REGS_IN_MAIN_FROM..=REGS_IN_MAIN_TO].to_vec()
253 } else {
254 vec![]
255 },
256 )
257 }
258
259 fn complete_trace_with_initial_reg_steps_per_chunk(
260 num_rows: usize,
261 fill_trace_outputs: &[(u64, Vec<u64>, EmuRegTrace, Vec<u32>)],
262 main_trace: &mut MainTraceType<F>,
263 step_range_check: &mut [u32],
264 reg_steps: &mut [u64; REGS_IN_MAIN],
265 ) -> Vec<u32> {
266 let mut large_range_checks: Vec<u32> = vec![];
267 let max_range = step_range_check.len() as u64;
268 for (index, (_, _, reg_trace, _)) in fill_trace_outputs.iter().enumerate().skip(1) {
269 #[allow(clippy::needless_range_loop)]
270 for reg_index in 0..REGS_IN_MAIN {
271 let reg_prev_mem_step = if fill_trace_outputs[index - 1].2.reg_steps[reg_index] == 0
272 {
273 reg_steps[reg_index]
274 } else {
275 fill_trace_outputs[index - 1].2.reg_steps[reg_index]
276 };
277 reg_steps[reg_index] = reg_prev_mem_step;
278 if reg_trace.first_step_uses[reg_index].is_some() {
279 let mem_step = reg_trace.first_step_uses[reg_index].unwrap();
280 let slot = MemHelpers::mem_step_to_slot(mem_step);
281 let row = MemHelpers::mem_step_to_row(mem_step) % num_rows;
282 let range = mem_step - reg_prev_mem_step - 1;
283 if range >= max_range {
284 large_range_checks.push(range as u32);
285 } else {
286 step_range_check[range as usize] += 1;
287 }
288 match slot {
289 0 => {
290 main_trace.buffer[row].set_a_reg_prev_mem_step(reg_prev_mem_step);
291 }
292 1 => {
293 main_trace.buffer[row].set_b_reg_prev_mem_step(reg_prev_mem_step);
294 }
295 2 => {
296 main_trace.buffer[row].set_store_reg_prev_mem_step(reg_prev_mem_step);
297 }
298 _ => panic!("Invalid slot {slot}"),
299 }
300 // TODO: range_check mem_step - reg_prev_mem_step
301 }
302 }
303 }
304 large_range_checks
305 }
306
307 fn update_reg_steps_with_last_chunk(
308 last_emu_reg_trace: &EmuRegTrace,
309 reg_steps: &mut [u64; REGS_IN_MAIN],
310 ) {
311 #[allow(clippy::needless_range_loop)]
312 for reg_index in 0..REGS_IN_MAIN {
313 let reg_prev_mem_step = if last_emu_reg_trace.reg_steps[reg_index] == 0 {
314 reg_steps[reg_index]
315 } else {
316 last_emu_reg_trace.reg_steps[reg_index]
317 };
318 reg_steps[reg_index] = reg_prev_mem_step;
319 }
320 }
321 fn update_reg_airvalues(
322 air_values: &mut MainAirValues<'_, F>,
323 final_step: u64,
324 last_reg_values: &[u64],
325 reg_steps: &[u64; REGS_IN_MAIN],
326 step_range_check: &mut [u32],
327 large_range_checks: &mut Vec<u32>,
328 ) {
329 let max_range = step_range_check.len() as u64;
330 for ireg in 0..REGS_IN_MAIN {
331 let reg_value = last_reg_values[ireg];
332 let values = [F::from_u32(reg_value as u32), F::from_u32((reg_value >> 32) as u32)];
333 air_values.last_reg_value[ireg] = values;
334 air_values.last_reg_mem_step[ireg] = F::from_u64(reg_steps[ireg]);
335 let range = (final_step - reg_steps[ireg] - 1) as usize;
336 if range >= max_range as usize {
337 large_range_checks.push(range as u32);
338 } else {
339 step_range_check[range] += 1;
340 }
341 }
342 }
343 fn update_std_range_checks(
344 &self,
345 segment_id: SegmentId,
346 step_range_check: Vec<u32>,
347 large_range_checks: &[u32],
348 ) {
349 let range_id = self
350 .std
351 .get_range_id(0, MEM_REGS_MAX_DIFF as i64, None)
352 .expect("Failed to get range ID");
353 self.std.range_checks(range_id, step_range_check);
354
355 for range in large_range_checks {
356 self.std.range_check(range_id, *range as i64, 1);
357 }
358 let range_id = self
359 .std
360 .get_range_id(0, Self::MAX_SEGMENT_ID as i64, None)
361 .expect("Failed to get range ID");
362 self.std.range_check(range_id, segment_id.as_usize() as i64, 1);
363 }
364}
365
366/// The `MainSM` struct represents the Main State Machine,
367/// responsible for generating the main witness.
368pub struct MainSM {}
369
370impl MainSM {
371 /// Debug method for the main state machine.
372 pub fn debug<F: PrimeField64>(_pctx: &ProofCtx<F>, _sctx: &SetupCtx<F>) {
373 // No debug information to display
374 }
375
376 pub fn build_counter() -> Box<dyn BusDeviceMetrics> {
377 Box::new(MainCounter::new())
378 }
379}