1use std::sync::Arc;
2
3use crate::{MemInput, MemModule, MemPreviousSegment};
4use fields::PrimeField64;
5use mem_common::{MEM_BYTES_BITS, SEGMENT_ADDR_MAX_RANGE};
6use pil_std_lib::Std;
7use proofman_common::{AirInstance, FromTrace, ProofmanResult};
8#[cfg(feature = "debug_mem")]
9use std::{
10 env,
11 fs::File,
12 io::{BufWriter, Write},
13};
14use zisk_common::SegmentId;
15use zisk_core::{ROM_ADDR, ROM_ADDR_MAX};
16use zisk_pil::RomDataAirValues;
17#[cfg(not(feature = "packed"))]
18use zisk_pil::RomDataTrace;
19#[cfg(feature = "packed")]
20use zisk_pil::RomDataTracePacked;
21
22#[cfg(feature = "packed")]
23type RomDataTraceType<F> = RomDataTracePacked<F>;
24
25#[cfg(not(feature = "packed"))]
26type RomDataTraceType<F> = RomDataTrace<F>;
27
28pub const ROM_DATA_W_ADDR_INIT: u32 = ROM_ADDR as u32 >> MEM_BYTES_BITS;
29pub const ROM_DATA_W_ADDR_END: u32 = ROM_ADDR_MAX as u32 >> MEM_BYTES_BITS;
30
31const _: () = {
32 assert!(ROM_ADDR_MAX <= 0xFFFF_FFFF, "ROM_DATA memory exceeds the 32-bit addressable range");
33 assert!(
34 (ROM_ADDR_MAX - ROM_ADDR) <= (128 << 20),
35 "ROM_DATA is too large. ROM size must be <= 128MB"
36 );
37};
38
39pub struct RomDataSM<F: PrimeField64> {
40 /// PIL2 standard library
41 std: Arc<Std<F>>,
42
43 range_id: usize,
44}
45
46#[allow(unused, unused_variables)]
47impl<F: PrimeField64> RomDataSM<F> {
48 pub fn new(std: Arc<Std<F>>) -> Arc<Self> {
49 let range_id = std
50 .get_range_id(0, SEGMENT_ADDR_MAX_RANGE as i64, None)
51 .expect("Failed to get range ID");
52
53 Arc::new(Self { std: std.clone(), range_id })
54 }
55 pub fn get_from_addr() -> u32 {
56 ROM_DATA_W_ADDR_INIT
57 }
58 fn get_u32_values(&self, value: u64) -> (u32, u32) {
59 (value as u32, (value >> 32) as u32)
60 }
61 pub fn get_to_addr() -> u32 {
62 ROM_DATA_W_ADDR_END
63 }
64 #[cfg(feature = "debug_mem")]
65 pub fn save_to_file(trace: &RomDataTrace<F>, file_name: &str) {
66 let file = File::create(file_name).unwrap();
67 let mut writer = BufWriter::new(file);
68 let num_rows = RomDataTrace::NUM_ROWS;
69
70 for i in 0..num_rows {
71 let addr = trace[i].get_addr() * 8;
72 let step = trace[i].get_step();
73 let sel = trace[i].get_sel();
74 // TODO: chunk_size * 4 = 20
75 writeln!(
76 writer,
77 "{:#010X} {} {:?} S:{sel} @{}",
78 addr,
79 step,
80 trace[i].value,
81 (step - 1) >> 20
82 )
83 .unwrap();
84 }
85 }
86}
87
88impl<F: PrimeField64> MemModule<F> for RomDataSM<F> {
89 fn get_addr_range(&self) -> (u32, u32) {
90 (ROM_DATA_W_ADDR_INIT, ROM_DATA_W_ADDR_END)
91 }
92 fn is_dual(&self) -> bool {
93 false
94 }
95
96 /// Finalizes the witness accumulation process and triggers the proof generation.
97 ///
98 /// This method is invoked by the executor when no further witness data remains to be added.
99 ///
100 /// # Parameters
101 ///
102 /// - `mem_inputs`: A slice of all `MemoryInput` inputs
103 fn compute_witness(
104 &self,
105 mem_ops: &[MemInput],
106 segment_id: SegmentId,
107 is_last_segment: bool,
108 previous_segment: &MemPreviousSegment,
109 trace_buffer: Vec<F>,
110 ) -> ProofmanResult<AirInstance<F>> {
111 let mut trace = RomDataTraceType::<F>::new_from_vec(trace_buffer)?;
112 let num_rows = RomDataTraceType::<F>::NUM_ROWS;
113 assert!(
114 !mem_ops.is_empty() && mem_ops.len() <= num_rows,
115 "RomDataSM: mem_ops.len()={} out of range {}",
116 mem_ops.len(),
117 num_rows
118 );
119
120 // range of instance
121 self.std.range_check(
122 self.range_id,
123 (previous_segment.addr - ROM_DATA_W_ADDR_INIT) as i64,
124 1,
125 );
126
127 let mut max_range_distance_count = 0;
128
129 // Fill the remaining rows
130 let mut last_addr: u32 = previous_segment.addr;
131 let mut last_step: u64 = previous_segment.step;
132 let mut last_value: u64 = previous_segment.value;
133
134 if segment_id == 0 {
135 // In the pil, in first row of first segment, we use previous_segment less 1, to
136 // allow to use ROM_DATA_W_ADDR_INIT as address, and active address change flag
137 // to free the value, if not
138 last_addr = ROM_DATA_W_ADDR_INIT - 1;
139 }
140 let mut i = 0;
141
142 for mem_op in mem_ops.iter() {
143 let distance = mem_op.addr - last_addr;
144 if i >= num_rows {
145 break;
146 }
147 if distance > SEGMENT_ADDR_MAX_RANGE as u32 {
148 let mut internal_reads = (distance - 1) / SEGMENT_ADDR_MAX_RANGE as u32;
149
150 #[cfg(feature = "debug_mem")]
151 println!(
152 "INTERNAL_READS[{},{}] {} 0x{:X},{} LAST:0x{:X}",
153 segment_id,
154 i,
155 internal_reads,
156 mem_op.addr * 8,
157 mem_op.step,
158 last_addr * 8
159 );
160
161 // check if has enough rows to complete the internal reads + regular memory
162 let incomplete = (i + internal_reads as usize) >= num_rows;
163 if incomplete {
164 internal_reads = (num_rows - i) as u32;
165 }
166
167 trace[i].set_addr_changes(true);
168 last_addr += SEGMENT_ADDR_MAX_RANGE as u32;
169 max_range_distance_count += 1;
170 trace[i].set_addr(last_addr);
171 trace[i].set_value(0, 0);
172 trace[i].set_value(1, 0);
173 trace[i].set_sel(false);
174 // the step, value of internal reads isn't relevant
175 trace[i].set_step(0);
176 i += 1;
177
178 for _j in 1..internal_reads {
179 trace[i] = trace[i - 1];
180 last_addr += SEGMENT_ADDR_MAX_RANGE as u32;
181 max_range_distance_count += 1;
182 trace[i].set_addr(last_addr);
183 i += 1;
184 }
185 if incomplete {
186 break;
187 }
188 }
189 trace[i].set_addr(mem_op.addr);
190 trace[i].set_step(mem_op.step);
191 trace[i].set_sel(true);
192
193 let (low_val, high_val) = self.get_u32_values(mem_op.value);
194 trace[i].set_value(0, low_val);
195 trace[i].set_value(1, high_val);
196
197 let addr_changes = last_addr != mem_op.addr;
198 if addr_changes || (i == 0 && segment_id == 0) {
199 trace[i].set_addr_changes(true);
200 self.std.range_check(self.range_id, (mem_op.addr - last_addr - 1) as i64, 1);
201 } else {
202 trace[i].set_addr_changes(false);
203 }
204
205 last_addr = mem_op.addr;
206 last_step = mem_op.step;
207 last_value = mem_op.value;
208 i += 1;
209 }
210 let count = i;
211 // STEP3. Add dummy rows to the output vector to fill the remaining rows
212 // PADDING: At end of memory fill with same addr, incrementing step, same value, sel = 0, rd
213 // = 1, wr = 0
214 let last_row_idx = count - 1;
215 if count < num_rows {
216 trace[count] = trace[last_row_idx];
217 trace[count].set_addr_changes(false);
218 trace[count].set_sel(false);
219
220 for i in count + 1..num_rows {
221 trace[i] = trace[i - 1];
222 }
223 // address doesn't change in padding rows, no range check is required
224 }
225
226 self.std.range_check(
227 self.range_id,
228 SEGMENT_ADDR_MAX_RANGE as i64,
229 max_range_distance_count,
230 );
231 self.std.range_check(self.range_id, (ROM_DATA_W_ADDR_END - last_addr) as i64, 1);
232
233 let mut air_values = RomDataAirValues::<F>::new();
234 air_values.segment_id = F::from_usize(segment_id.into());
235 air_values.is_first_segment = F::from_bool(segment_id == 0);
236 air_values.is_last_segment = F::from_bool(is_last_segment);
237 air_values.previous_segment_step = F::from_u64(previous_segment.step);
238 air_values.previous_segment_addr = F::from_u32(previous_segment.addr);
239 air_values.segment_last_addr = F::from_u32(last_addr);
240 air_values.segment_last_step = F::from_u64(last_step);
241
242 air_values.previous_segment_value[0] = F::from_u32(previous_segment.value as u32);
243 air_values.previous_segment_value[1] = F::from_u32((previous_segment.value >> 32) as u32);
244
245 air_values.segment_last_value[0] = F::from_u32(last_value as u32);
246 air_values.segment_last_value[1] = F::from_u32((last_value >> 32) as u32);
247
248 #[cfg(feature = "debug_mem")]
249 {
250 let path = env::var("MEM_TRACE_DIR").unwrap_or("tmp/mem_trace".to_string());
251 let filename = format!("{path}/rom_trace_{segment_id:04}.txt");
252 Self::save_to_file(&trace, &filename);
253 }
254
255 Ok(AirInstance::new_from_trace(FromTrace::new(&mut trace).with_air_values(&mut air_values)))
256 }
257}