1//! The `RomSM` module implements the ROM State Machine,
2//! directly managing the ROM execution process, generating traces, and computing custom traces.
3//!
4//! Key components of this module include:
5//! - The `RomSM` struct, which represents the ROM State Machine and encapsulates ROM-related
6//! operations.
7//! - Methods for proving instances and computing traces from the ROM data.
8//! - `ComponentBuilder` trait implementations for creating counters, planners, and input
9//! collectors.
10
11use std::{
12 path::PathBuf,
13 sync::{
14 atomic::{AtomicBool, AtomicU32},
15 Arc, Mutex,
16 },
17 thread::JoinHandle,
18};
19
20use crate::{RomInstance, RomPlanner};
21use asm_runner::{AsmRHData, AsmRunnerRH};
22use fields::PrimeField64;
23use itertools::Itertools;
24use proofman_common::{AirInstance, FromTrace, ProofmanResult};
25use zisk_common::{
26 create_atomic_vec, BusDeviceMetrics, ComponentBuilder, CounterStats, Instance, InstanceCtx,
27 Planner,
28};
29use zisk_core::{
30 zisk_ops::ZiskOp, Riscv2zisk, ZiskRom, ROM_ADDR, ROM_ADDR_MAX, ROM_ENTRY, ROM_EXIT, SRC_IMM,
31};
32use zisk_pil::{MainTrace, RomRomTrace, RomRomTraceRow, RomTrace};
33
34/// The `RomSM` struct represents the ROM State Machine
35pub struct RomSM {
36 /// Zisk Rom
37 zisk_rom: Arc<ZiskRom>,
38
39 /// Shared biod instruction counter for monitoring ROM operations.
40 bios_inst_count: Arc<Vec<AtomicU32>>,
41
42 /// Shared program instruction counter for monitoring ROM operations.
43 prog_inst_count: Arc<Vec<AtomicU32>>,
44
45 asm_runner_handler: Mutex<Option<JoinHandle<AsmRunnerRH>>>,
46}
47
48impl RomSM {
49 /// Creates a new instance of the `RomSM` state machine.
50 ///
51 /// # Arguments
52 /// * `zisk_rom` - The Zisk ROM representation.
53 ///
54 /// # Returns
55 /// An `Arc`-wrapped instance of `RomSM`.
56 pub fn new(zisk_rom: Arc<ZiskRom>, asm_rom_path: Option<PathBuf>) -> Arc<Self> {
57 let (bios_inst_count, prog_inst_count) = if asm_rom_path.is_some() {
58 (vec![], vec![])
59 } else {
60 (
61 create_atomic_vec(((ROM_ADDR - ROM_ENTRY) as usize) >> 2), // No atomics, we can divide by 4
62 create_atomic_vec((ROM_ADDR_MAX - ROM_ADDR) as usize), // Cannot be dividede by 4
63 )
64 };
65
66 Arc::new(Self {
67 zisk_rom,
68 bios_inst_count: Arc::new(bios_inst_count),
69 prog_inst_count: Arc::new(prog_inst_count),
70 asm_runner_handler: Mutex::new(None),
71 })
72 }
73
74 pub fn set_asm_runner_handler(&self, handler: JoinHandle<AsmRunnerRH>) {
75 *self.asm_runner_handler.lock().unwrap() = Some(handler);
76 }
77
78 /// Computes the witness for the provided plan using the given ROM.
79 ///
80 /// # Arguments
81 /// * `rom` - Reference to the Zisk ROM.
82 /// * `plan` - The execution plan for computing the witness.
83 ///
84 /// # Returns
85 /// An `AirInstance` containing the computed witness trace data.
86 pub fn compute_witness<F: PrimeField64>(
87 rom: &ZiskRom,
88 counter_stats: &CounterStats,
89 calculated: &AtomicBool,
90 trace_buffer: Vec<F>,
91 ) -> ProofmanResult<AirInstance<F>> {
92 let mut rom_trace = RomTrace::new_from_vec_zeroes(trace_buffer)?;
93
94 let main_trace_len = MainTrace::<F>::NUM_ROWS as u64;
95
96 tracing::debug!("··· Creating Rom instance [{} rows]", RomTrace::<F>::NUM_ROWS);
97
98 // For every instruction in the rom, fill its corresponding ROM trace
99 for (i, key) in rom.insts.keys().sorted().enumerate() {
100 // Get the Zisk instruction
101 let inst = &rom.insts[key].i;
102
103 // Calculate the multiplicity, i.e. the number of times this pc is used in this
104 // execution
105 let mut multiplicity: u64;
106 if inst.paddr < ROM_ADDR {
107 if counter_stats.bios_inst_count.is_empty() {
108 multiplicity = 1; // If the histogram is empty, we use 1 for all pc's
109 } else {
110 match calculated.load(std::sync::atomic::Ordering::Relaxed) {
111 true => {
112 multiplicity = counter_stats.bios_inst_count
113 [((inst.paddr - ROM_ENTRY) as usize) >> 2]
114 .swap(0, std::sync::atomic::Ordering::Relaxed)
115 as u64;
116 }
117 false => {
118 multiplicity = counter_stats.bios_inst_count
119 [((inst.paddr - ROM_ENTRY) as usize) >> 2]
120 .load(std::sync::atomic::Ordering::Relaxed)
121 as u64;
122 }
123 }
124
125 if multiplicity == 0 {
126 continue;
127 }
128 if inst.paddr == counter_stats.end_pc {
129 multiplicity += main_trace_len - counter_stats.steps % main_trace_len;
130 }
131 }
132 } else {
133 match calculated.load(std::sync::atomic::Ordering::Relaxed) {
134 true => {
135 multiplicity = counter_stats.prog_inst_count
136 [(inst.paddr - ROM_ADDR) as usize]
137 .swap(0, std::sync::atomic::Ordering::Relaxed)
138 as u64
139 }
140 false => {
141 multiplicity = counter_stats.prog_inst_count
142 [(inst.paddr - ROM_ADDR) as usize]
143 .load(std::sync::atomic::Ordering::Relaxed)
144 as u64
145 }
146 }
147 if multiplicity == 0 {
148 continue;
149 }
150 if inst.paddr == counter_stats.end_pc {
151 multiplicity += main_trace_len - counter_stats.steps % main_trace_len;
152 }
153 }
154 rom_trace[i].multiplicity = F::from_u64(multiplicity);
155 }
156
157 Ok(AirInstance::new_from_trace(FromTrace::new(&mut rom_trace)))
158 }
159
160 pub fn compute_witness_from_asm<F: PrimeField64>(
161 rom: &ZiskRom,
162 asm_romh: &AsmRHData,
163 trace_buffer: Vec<F>,
164 ) -> ProofmanResult<AirInstance<F>> {
165 let mut rom_trace = RomTrace::new_from_vec_zeroes(trace_buffer)?;
166
167 tracing::debug!("··· Creating Rom instance [{} rows]", RomTrace::<F>::NUM_ROWS);
168
169 let main_trace_len = MainTrace::<F>::NUM_ROWS as u64;
170
171 for (i, key) in rom.insts.keys().sorted().enumerate() {
172 // Get the Zisk instruction
173 let inst = &rom.insts[key].i;
174
175 // Calculate the multiplicity, i.e. the number of times this pc is used in this
176 // execution
177 let mut multiplicity: u64;
178 if inst.paddr < ROM_ADDR {
179 if asm_romh.bios_inst_count.is_empty() {
180 multiplicity = 1; // If the histogram is empty, we use 1 for all pc's
181 } else {
182 let idx = ((inst.paddr - ROM_ENTRY) as usize) >> 2;
183
184 multiplicity = asm_romh.bios_inst_count[idx];
185
186 if multiplicity == 0 {
187 continue;
188 }
189
190 if inst.paddr == ROM_EXIT {
191 multiplicity += main_trace_len - asm_romh.steps % main_trace_len;
192 }
193 }
194 } else {
195 let idx = (inst.paddr - ROM_ADDR) as usize;
196 multiplicity = asm_romh.prog_inst_count[idx];
197
198 if multiplicity == 0 {
199 continue;
200 }
201 }
202
203 rom_trace[i].multiplicity = F::from_u64(multiplicity);
204 }
205
206 Ok(AirInstance::new_from_trace(FromTrace::new(&mut rom_trace)))
207 }
208
209 /// Computes the ROM trace based on the ROM instructions.
210 ///
211 /// # Arguments
212 /// * `rom` - Reference to the Zisk ROM.
213 /// * `rom_custom_trace` - Reference to the custom ROM trace.
214 fn compute_trace_rom<F: PrimeField64>(rom: &ZiskRom, rom_custom_trace: &mut RomRomTrace<F>) {
215 // For every instruction in the rom, fill its corresponding ROM trace
216 for (i, key) in rom.insts.keys().sorted().enumerate() {
217 // Get the Zisk instruction
218 let inst = &rom.insts[key].i;
219
220 // Convert the i64 offsets to F
221 let jmp_offset1 = if inst.jmp_offset1 >= 0 {
222 F::from_u64(inst.jmp_offset1 as u64)
223 } else {
224 F::neg(F::from_u64((-inst.jmp_offset1) as u64))
225 };
226 let jmp_offset2 = if inst.jmp_offset2 >= 0 {
227 F::from_u64(inst.jmp_offset2 as u64)
228 } else {
229 F::neg(F::from_u64((-inst.jmp_offset2) as u64))
230 };
231 let store_offset = if inst.store_offset >= 0 {
232 F::from_u64(inst.store_offset as u64)
233 } else {
234 F::neg(F::from_u64((-inst.store_offset) as u64))
235 };
236 let a_offset_imm0 = if inst.a_offset_imm0 as i64 >= 0 {
237 F::from_u64(inst.a_offset_imm0)
238 } else {
239 F::neg(F::from_u64((-(inst.a_offset_imm0 as i64)) as u64))
240 };
241 let b_offset_imm0 = if inst.b_offset_imm0 as i64 >= 0 {
242 F::from_u64(inst.b_offset_imm0)
243 } else {
244 F::neg(F::from_u64((-(inst.b_offset_imm0 as i64)) as u64))
245 };
246
247 // Fill the rom trace row fields
248 rom_custom_trace[i].line = F::from_u64(inst.paddr); // TODO: unify names: pc, paddr, line
249 rom_custom_trace[i].a_offset_imm0 = a_offset_imm0;
250 rom_custom_trace[i].a_imm1 =
251 F::from_u64(if inst.a_src == SRC_IMM { inst.a_use_sp_imm1 } else { 0 });
252 rom_custom_trace[i].b_offset_imm0 = b_offset_imm0;
253 rom_custom_trace[i].b_imm1 =
254 F::from_u64(if inst.b_src == SRC_IMM { inst.b_use_sp_imm1 } else { 0 });
255 rom_custom_trace[i].ind_width = F::from_u64(inst.ind_width);
256 // IMPORTANT: the opcodes fcall, fcall_get, and fcall_param are really a variant
257 // of the copyb, use to get free-input information
258 rom_custom_trace[i].op = if inst.op == ZiskOp::Fcall.code()
259 || inst.op == ZiskOp::FcallGet.code()
260 || inst.op == ZiskOp::FcallParam.code()
261 {
262 F::from_u8(ZiskOp::CopyB.code())
263 } else {
264 F::from_u8(inst.op)
265 };
266 rom_custom_trace[i].store_offset = store_offset;
267 rom_custom_trace[i].jmp_offset1 = jmp_offset1;
268 rom_custom_trace[i].jmp_offset2 = jmp_offset2;
269 rom_custom_trace[i].flags = F::from_u64(inst.get_flags());
270 }
271
272 // Padd with zeroes
273 let num_rows: usize = RomRomTrace::<F>::NUM_ROWS;
274 for i in rom.insts.len()..num_rows {
275 rom_custom_trace[i] = RomRomTraceRow::default();
276 }
277 }
278
279 /// Computes a custom trace ROM from the given ELF file.
280 ///
281 /// # Arguments
282 /// * `rom_path` - The path to the ELF file.
283 /// * `rom_custom_trace` - Reference to the custom ROM trace.
284 pub fn compute_custom_trace_rom<F: PrimeField64>(
285 rom_path: PathBuf,
286 rom_custom_trace: &mut RomRomTrace<F>,
287 ) {
288 // Get the ELF file path as a string
289 let elf_filename: String = rom_path.to_str().unwrap().into();
290 tracing::info!("Computing custom trace ROM");
291
292 // Load and parse the ELF file, and transpile it into a ZisK ROM using Riscv2zisk
293
294 // Create an instance of the RISCV -> ZisK program converter
295 let riscv2zisk = Riscv2zisk::new(elf_filename);
296
297 // Convert program to rom
298 let rom = riscv2zisk.run().expect("RomSM::prover() failed converting elf to rom");
299
300 let rom_len = rom.insts.len();
301 let air_rom_len = RomTrace::<F>::NUM_ROWS;
302 if rom_len > air_rom_len {
303 panic!(
304 "Error: The generated ROM has {} instructions, which exceeds the maximum supported by the Zisk PIL ROM trace ({} instructions). Please review zisk.pil and increase the ROM trace size accordingly.",
305 rom_len, air_rom_len
306 );
307 }
308
309 Self::compute_trace_rom(&rom, rom_custom_trace);
310 }
311}
312
313impl<F: PrimeField64> ComponentBuilder<F> for RomSM {
314 /// Builds and returns a new counter for monitoring ROM operations.
315 ///
316 /// # Returns
317 /// A boxed implementation of `RomCounter`.
318 fn build_counter(&self) -> Option<Box<dyn BusDeviceMetrics>> {
319 None
320 }
321
322 /// Builds a planner for ROM-related instances.
323 ///
324 /// # Returns
325 /// A boxed implementation of `RomPlanner`.
326 fn build_planner(&self) -> Box<dyn Planner> {
327 Box::new(RomPlanner)
328 }
329
330 /// Builds an instance of the ROM state machine.
331 ///
332 /// # Arguments
333 /// * `ictx` - The context of the instance, containing the plan and its associated
334 ///
335 /// # Returns
336 /// A boxed implementation of `RomInstance`.
337 fn build_instance(&self, ictx: InstanceCtx) -> Box<dyn Instance<F>> {
338 let mut handle_rh_guard = self.asm_runner_handler.lock().unwrap();
339 let handle_rh = handle_rh_guard.take();
340
341 Box::new(RomInstance::new(
342 self.zisk_rom.clone(),
343 ictx,
344 self.bios_inst_count.clone(),
345 self.prog_inst_count.clone(),
346 handle_rh,
347 ))
348 }
349}