1use std::sync::Arc;
2
3#[cfg(feature = "debug_mem_align")]
4use std::sync::Mutex;
5
6use fields::PrimeField64;
7use pil_std_lib::Std;
8use rayon::prelude::*;
9
10use crate::MemAlignInput;
11use proofman_common::{AirInstance, FromTrace, ProofmanResult};
12use zisk_pil::{MemAlignByteAirValues, MemAlignReadByteAirValues, MemAlignWriteByteAirValues};
13
14#[cfg(not(feature = "packed"))]
15use zisk_pil::{
16 MemAlignByteTrace, MemAlignByteTraceRow, MemAlignReadByteTrace, MemAlignReadByteTraceRow,
17 MemAlignWriteByteTrace, MemAlignWriteByteTraceRow,
18};
19#[cfg(feature = "packed")]
20use zisk_pil::{
21 MemAlignByteTracePacked, MemAlignByteTraceRowPacked, MemAlignReadByteTracePacked,
22 MemAlignReadByteTraceRowPacked, MemAlignWriteByteTracePacked, MemAlignWriteByteTraceRowPacked,
23};
24
25#[cfg(feature = "packed")]
26pub type MemAlignByteTraceRowType<F> = MemAlignByteTraceRowPacked<F>;
27#[cfg(feature = "packed")]
28pub type MemAlignByteTraceType<F> = MemAlignByteTracePacked<F>;
29#[cfg(feature = "packed")]
30pub type MemAlignReadByteTraceRowType<F> = MemAlignReadByteTraceRowPacked<F>;
31#[cfg(feature = "packed")]
32pub type MemAlignReadByteTraceType<F> = MemAlignReadByteTracePacked<F>;
33#[cfg(feature = "packed")]
34pub type MemAlignWriteByteTraceRowType<F> = MemAlignWriteByteTraceRowPacked<F>;
35#[cfg(feature = "packed")]
36pub type MemAlignWriteByteTraceType<F> = MemAlignWriteByteTracePacked<F>;
37
38#[cfg(not(feature = "packed"))]
39pub type MemAlignByteTraceRowType<F> = MemAlignByteTraceRow<F>;
40#[cfg(not(feature = "packed"))]
41pub type MemAlignByteTraceType<F> = MemAlignByteTrace<F>;
42#[cfg(not(feature = "packed"))]
43pub type MemAlignReadByteTraceRowType<F> = MemAlignReadByteTraceRow<F>;
44#[cfg(not(feature = "packed"))]
45pub type MemAlignReadByteTraceType<F> = MemAlignReadByteTrace<F>;
46#[cfg(not(feature = "packed"))]
47pub type MemAlignWriteByteTraceRowType<F> = MemAlignWriteByteTraceRow<F>;
48#[cfg(not(feature = "packed"))]
49pub type MemAlignWriteByteTraceType<F> = MemAlignWriteByteTrace<F>;
50
51pub trait MemAlignByteRow<F: PrimeField64, T> {
52 #[allow(clippy::too_many_arguments)]
53 fn set_common_fields(
54 &mut self,
55 sel_high_4b: bool,
56 sel_high_2b: bool,
57 sel_high_b: bool,
58 direct_value: u32,
59 composed_value: u32,
60 value_16b: u16,
61 value_8b: u8,
62 byte_value: u8,
63 addr_w: u32,
64 step: u64,
65 );
66 fn set_write_fields(
67 &mut self,
68 is_write: bool,
69 written_composed_value: u32,
70 written_byte_value: u8,
71 mem_write_values: [u32; 2],
72 );
73 fn valid_for_read() -> bool;
74 fn valid_for_write() -> bool;
75 fn create_trace(trace_buffer: Vec<F>) -> ProofmanResult<T>;
76 fn get_num_rows(trace: &T) -> usize;
77 fn name() -> &'static str;
78 fn create_instance_from_trace(trace: &mut T, padding_row: usize) -> AirInstance<F>;
79 fn get_row_mut(trace: &mut T, index: usize) -> &mut Self;
80}
81
82// Helper function to avoid code duplication in create_instance_from_trace
83// fn create_instance_from_trace_helper<F: PrimeField64, T, V>(
84// trace: &mut T,
85// padding_row: usize,
86// mut air_values: V,
87// set_padding_size: impl FnOnce(&mut V, F),
88// ) -> AirInstance<F>
89// where
90// T: proofman_common::trace::Trace<F> + std::ops::Index<usize> + std::ops::IndexMut<usize>,
91// T::Output: Clone,
92// V: proofman_common::trace::Values<F>,
93// {
94// let num_rows = trace.num_rows();
95// let padding = trace[padding_row].clone();
96// let padding_size = num_rows - padding_row;
97// set_padding_size(&mut air_values, F::from_usize(padding_size));
98// trace.row_slice_mut()[padding_row + 1..num_rows]
99// .par_iter_mut()
100// .for_each(|slot| *slot = padding);
101// AirInstance::new_from_trace(FromTrace::new(trace).with_air_values(&mut air_values))
102// }
103
104// Implement the common trait for all trace types
105impl<F: PrimeField64> MemAlignByteRow<F, MemAlignByteTraceType<F>> for MemAlignByteTraceRowType<F> {
106 #[inline(always)]
107 fn set_common_fields(
108 &mut self,
109 sel_high_4b: bool,
110 sel_high_2b: bool,
111 sel_high_b: bool,
112 direct_value: u32,
113 composed_value: u32,
114 value_16b: u16,
115 value_8b: u8,
116 byte_value: u8,
117 addr_w: u32,
118 step: u64,
119 ) {
120 self.set_sel_high_4b(sel_high_4b);
121 self.set_sel_high_2b(sel_high_2b);
122 self.set_sel_high_b(sel_high_b);
123 self.set_direct_value(direct_value);
124 self.set_composed_value(composed_value);
125 self.set_value_16b(value_16b);
126 self.set_value_8b(value_8b);
127 self.set_byte_value(byte_value);
128 self.set_addr_w(addr_w);
129 self.set_step(step);
130 }
131 #[inline(always)]
132 fn set_write_fields(
133 &mut self,
134 is_write: bool,
135 written_composed_value: u32,
136 written_byte_value: u8,
137 mem_write_values: [u32; 2],
138 ) {
139 self.set_is_write(is_write);
140 self.set_written_composed_value(written_composed_value);
141 self.set_written_byte_value(written_byte_value);
142 self.set_bus_byte(if is_write {
143 self.get_written_byte_value()
144 } else {
145 self.get_byte_value()
146 });
147 self.set_mem_write_values(0, mem_write_values[0]);
148 self.set_mem_write_values(1, mem_write_values[1]);
149 }
150 #[inline(always)]
151 fn valid_for_read() -> bool {
152 true
153 }
154 #[inline(always)]
155 fn valid_for_write() -> bool {
156 true
157 }
158 fn create_trace(trace_buffer: Vec<F>) -> ProofmanResult<MemAlignByteTraceType<F>> {
159 MemAlignByteTraceType::new_from_vec(trace_buffer)
160 }
161 fn get_num_rows(trace: &MemAlignByteTraceType<F>) -> usize {
162 trace.num_rows()
163 }
164 fn name() -> &'static str {
165 "MemAlignByteTraceType"
166 }
167 fn get_row_mut(trace: &mut MemAlignByteTraceType<F>, index: usize) -> &mut Self {
168 &mut trace[index]
169 }
170 fn create_instance_from_trace(
171 trace: &mut MemAlignByteTraceType<F>,
172 padding_row: usize,
173 ) -> AirInstance<F> {
174 let num_rows = trace.num_rows();
175 let padding_size = num_rows - padding_row;
176 if padding_size > 0 {
177 let padding = trace[padding_row];
178 trace.buffer[padding_row + 1..num_rows].par_iter_mut().for_each(|slot| *slot = padding);
179 }
180 let mut air_values = MemAlignByteAirValues::<F>::new();
181 air_values.padding_size = F::from_usize(padding_size);
182 AirInstance::new_from_trace(FromTrace::new(trace).with_air_values(&mut air_values))
183 }
184 // fn create_instance_from_trace(
185 // trace: &mut MemAlignByteTrace<F>,
186 // padding_row: usize,
187 // ) -> AirInstance<F> {
188 // create_instance_from_trace_helper(
189 // trace,
190 // padding_row,
191 // MemAlignByteAirValues::<F>::new(),
192 // |air_values, size| air_values.padding_size = size,
193 // )
194 // }
195}
196
197impl<F: PrimeField64> MemAlignByteRow<F, MemAlignReadByteTraceType<F>>
198 for MemAlignReadByteTraceRowType<F>
199{
200 fn set_common_fields(
201 &mut self,
202 sel_high_4b: bool,
203 sel_high_2b: bool,
204 sel_high_b: bool,
205 direct_value: u32,
206 composed_value: u32,
207 value_16b: u16,
208 value_8b: u8,
209 byte_value: u8,
210 addr_w: u32,
211 step: u64,
212 ) {
213 self.set_sel_high_4b(sel_high_4b);
214 self.set_sel_high_2b(sel_high_2b);
215 self.set_sel_high_b(sel_high_b);
216 self.set_direct_value(direct_value);
217 self.set_composed_value(composed_value);
218 self.set_value_16b(value_16b);
219 self.set_value_8b(value_8b);
220 self.set_byte_value(byte_value);
221 self.set_addr_w(addr_w);
222 self.set_step(step);
223 }
224 #[inline(always)]
225 fn set_write_fields(
226 &mut self,
227 _is_write: bool,
228 _written_composed_value: u32,
229 _written_byte_value: u8,
230 _mem_write_values: [u32; 2],
231 ) {
232 }
233 #[inline(always)]
234 fn valid_for_read() -> bool {
235 true
236 }
237 #[inline(always)]
238 fn valid_for_write() -> bool {
239 false
240 }
241 fn create_trace(trace_buffer: Vec<F>) -> ProofmanResult<MemAlignReadByteTraceType<F>> {
242 MemAlignReadByteTraceType::new_from_vec(trace_buffer)
243 }
244 fn get_num_rows(trace: &MemAlignReadByteTraceType<F>) -> usize {
245 trace.num_rows()
246 }
247 fn name() -> &'static str {
248 "MemAlignReadByteTraceType"
249 }
250 fn get_row_mut(trace: &mut MemAlignReadByteTraceType<F>, index: usize) -> &mut Self {
251 &mut trace[index]
252 }
253 fn create_instance_from_trace(
254 trace: &mut MemAlignReadByteTraceType<F>,
255 padding_row: usize,
256 ) -> AirInstance<F> {
257 let num_rows = trace.num_rows();
258 let padding_size = num_rows - padding_row;
259 if padding_size > 0 {
260 let padding = trace[padding_row];
261 trace.buffer[padding_row + 1..num_rows].par_iter_mut().for_each(|slot| *slot = padding);
262 }
263 let mut air_values = MemAlignReadByteAirValues::<F>::new();
264 air_values.padding_size = F::from_usize(padding_size);
265 AirInstance::new_from_trace(FromTrace::new(trace).with_air_values(&mut air_values))
266 }
267 // fn create_instance_from_trace(
268 // trace: &mut MemAlignReadByteTrace<F>,
269 // padding_row: usize,
270 // ) -> AirInstance<F> {
271 // create_instance_from_trace_helper(
272 // trace,
273 // padding_row,
274 // MemAlignReadByteAirValues::<F>::new(),
275 // |air_values, size| air_values.padding_size = size,
276 // )
277 // }
278}
279
280impl<F: PrimeField64> MemAlignByteRow<F, MemAlignWriteByteTraceType<F>>
281 for MemAlignWriteByteTraceRowType<F>
282{
283 fn set_common_fields(
284 &mut self,
285 sel_high_4b: bool,
286 sel_high_2b: bool,
287 sel_high_b: bool,
288 direct_value: u32,
289 composed_value: u32,
290 value_16b: u16,
291 value_8b: u8,
292 byte_value: u8,
293 addr_w: u32,
294 step: u64,
295 ) {
296 self.set_sel_high_4b(sel_high_4b);
297 self.set_sel_high_2b(sel_high_2b);
298 self.set_sel_high_b(sel_high_b);
299 self.set_direct_value(direct_value);
300 self.set_composed_value(composed_value);
301 self.set_value_16b(value_16b);
302 self.set_value_8b(value_8b);
303 self.set_byte_value(byte_value);
304 self.set_addr_w(addr_w);
305 self.set_step(step);
306 }
307 #[inline(always)]
308 fn set_write_fields(
309 &mut self,
310 _is_write: bool,
311 written_composed_value: u32,
312 written_byte_value: u8,
313 mem_write_values: [u32; 2],
314 ) {
315 self.set_written_composed_value(written_composed_value);
316 self.set_written_byte_value(written_byte_value);
317 self.set_mem_write_values(0, mem_write_values[0]);
318 self.set_mem_write_values(1, mem_write_values[1]);
319 }
320 #[inline(always)]
321 fn valid_for_read() -> bool {
322 false
323 }
324 #[inline(always)]
325 fn valid_for_write() -> bool {
326 true
327 }
328 fn create_trace(trace_buffer: Vec<F>) -> ProofmanResult<MemAlignWriteByteTraceType<F>> {
329 MemAlignWriteByteTraceType::new_from_vec(trace_buffer)
330 }
331 fn get_num_rows(trace: &MemAlignWriteByteTraceType<F>) -> usize {
332 trace.num_rows()
333 }
334
335 fn name() -> &'static str {
336 "MemAlignWriteByteTrace"
337 }
338 fn get_row_mut(trace: &mut MemAlignWriteByteTraceType<F>, index: usize) -> &mut Self {
339 &mut trace[index]
340 }
341 fn create_instance_from_trace(
342 trace: &mut MemAlignWriteByteTraceType<F>,
343 padding_row: usize,
344 ) -> AirInstance<F> {
345 let num_rows = trace.num_rows();
346 let padding_size = num_rows - padding_row;
347 if padding_size > 0 {
348 let padding = trace[padding_row];
349 trace.buffer[padding_row + 1..num_rows].par_iter_mut().for_each(|slot| *slot = padding);
350 }
351 let mut air_values = MemAlignWriteByteAirValues::<F>::new();
352 air_values.padding_size = F::from_usize(padding_size);
353 AirInstance::new_from_trace(FromTrace::new(trace).with_air_values(&mut air_values))
354 }
355}
356
357const OFFSET_MASK: u32 = 0x07;
358const OFFSET_BITS: u32 = 3;
359const DUAL_BYTE_TABLE_ID: usize = 88;
360
361pub struct MemAlignByteSM<F: PrimeField64> {
362 /// PIL2 standard library
363 std: Arc<Std<F>>,
364
365 #[cfg(feature = "debug_mem_align")]
366 num_computed_rows: Mutex<usize>,
367
368 /// The table ID for the Mem Align ROM State Machine
369 table_dual_byte_id: usize,
370
371 table_16b_id: usize,
372 table_8b_id: usize,
373}
374
375impl<F: PrimeField64> MemAlignByteSM<F> {
376 pub fn new(std: Arc<Std<F>>) -> Arc<Self> {
377 // Get the table ID
378 Arc::new(Self {
379 std: std.clone(),
380 table_dual_byte_id: std
381 .get_virtual_table_id(DUAL_BYTE_TABLE_ID)
382 .expect("Failed to get dual byte table ID"),
383 table_16b_id: std.get_range_id(0, 0xFFFF, None).expect("Failed to get 16b table ID"),
384 table_8b_id: std.get_range_id(0, 0xFF, None).expect("Failed to get 8b table ID"),
385 #[cfg(feature = "debug_mem_align")]
386 num_computed_rows: Mutex::new(0),
387 })
388 }
389
390 pub fn compute_witness<T, R: MemAlignByteRow<F, T>>(
391 &self,
392 mem_ops: &[Vec<MemAlignInput>],
393 used_rows: usize,
394 trace_buffer: Vec<F>,
395 ) -> ProofmanResult<AirInstance<F>> {
396 let mut trace = R::create_trace(trace_buffer)?;
397 let num_rows = R::get_num_rows(&trace);
398
399 tracing::debug!(
400 "··· Creating {} instance [{} / {} rows filled {:.2}%]",
401 R::name(),
402 used_rows,
403 num_rows,
404 used_rows as f64 / num_rows as f64 * 100.0
405 );
406
407 let mut irow = 0;
408 for inner_memp_ops in mem_ops.iter() {
409 for input in inner_memp_ops.iter() {
410 assert!(irow < num_rows);
411 self.compute_row_witness(input, irow, R::get_row_mut(&mut trace, irow));
412 irow += 1;
413 }
414 }
415 let padding_size = (num_rows - irow) as u64;
416 if padding_size > 0 {
417 let padding_row = R::get_row_mut(&mut trace, irow);
418 self.compute_row_witness(
419 &MemAlignInput {
420 addr: 0,
421 width: 1,
422 is_write: !R::valid_for_read(),
423 value: 0,
424 mem_values: [0, 0],
425 step: 0,
426 },
427 irow,
428 padding_row,
429 );
430 // padding_size - 1, because compute_row_witness call range_check
431 self.std.inc_virtual_row(self.table_dual_byte_id, 0, padding_size - 1);
432 self.std.range_check(self.table_16b_id, 0, padding_size - 1);
433 if R::valid_for_write() {
434 self.std.range_check(self.table_8b_id, 0, padding_size - 1);
435 }
436 }
437 Ok(R::create_instance_from_trace(&mut trace, irow))
438 }
439
440 /// Common logic for computing witness that can be shared across different trace types
441 fn compute_row_witness<T, R: MemAlignByteRow<F, T>>(
442 &self,
443 input: &MemAlignInput,
444 irow: usize,
445 row: &mut R,
446 ) {
447 let addr = input.addr;
448
449 let high_value = (input.mem_values[0] >> 32) as u32;
450 let low_value = (input.mem_values[0] & 0xFFFF_FFFF) as u32;
451 let offset = (addr & OFFSET_MASK) as u8;
452 let addr_w = addr >> OFFSET_BITS;
453 let step = input.step;
454
455 let (
456 sel_high_4b,
457 sel_high_2b,
458 sel_high_b,
459 direct_value,
460 composed_value,
461 byte_value,
462 value_16b,
463 value_8b,
464 ) = match offset {
465 0 => (
466 false,
467 false,
468 false,
469 high_value,
470 low_value,
471 low_value as u8,
472 (low_value >> 16) as u16,
473 (low_value >> 8) as u8,
474 ),
475 1 => (
476 false,
477 false,
478 true,
479 high_value,
480 low_value,
481 (low_value >> 8) as u8,
482 (low_value >> 16) as u16,
483 low_value as u8,
484 ),
485 2 => (
486 false,
487 true,
488 false,
489 high_value,
490 low_value,
491 (low_value >> 16) as u8,
492 low_value as u16,
493 (low_value >> 24) as u8,
494 ),
495 3 => (
496 false,
497 true,
498 true,
499 high_value,
500 low_value,
501 (low_value >> 24) as u8,
502 low_value as u16,
503 (low_value >> 16) as u8,
504 ),
505 4 => (
506 true,
507 false,
508 false,
509 low_value,
510 high_value,
511 high_value as u8,
512 (high_value >> 16) as u16,
513 (high_value >> 8) as u8,
514 ),
515 5 => (
516 true,
517 false,
518 true,
519 low_value,
520 high_value,
521 (high_value >> 8) as u8,
522 (high_value >> 16) as u16,
523 high_value as u8,
524 ),
525 6 => (
526 true,
527 true,
528 false,
529 low_value,
530 high_value,
531 (high_value >> 16) as u8,
532 high_value as u16,
533 (high_value >> 24) as u8,
534 ),
535 7 => (
536 true,
537 true,
538 true,
539 low_value,
540 high_value,
541 (high_value >> 24) as u8,
542 high_value as u16,
543 (high_value >> 16) as u8,
544 ),
545 _ => unreachable!("Invalid offset"),
546 };
547
548 row.set_common_fields(
549 sel_high_4b,
550 sel_high_2b,
551 sel_high_b,
552 direct_value,
553 composed_value,
554 value_16b,
555 value_8b,
556 byte_value,
557 addr_w,
558 step,
559 );
560 self.std.inc_virtual_row(
561 self.table_dual_byte_id,
562 (value_8b as u16 + ((byte_value as u16) << 8)) as u64,
563 1,
564 );
565 self.std.range_check(self.table_16b_id, value_16b as i64, 1);
566
567 let written_byte_value = input.value as u8;
568 let written_composed_value = match offset {
569 0 => (low_value & 0xFFFF_FF00) | (written_byte_value as u32),
570 1 => (low_value & 0xFFFF_00FF) | ((written_byte_value as u32) << 8),
571 2 => (low_value & 0xFF00_FFFF) | ((written_byte_value as u32) << 16),
572 3 => (low_value & 0x00FF_FFFF) | ((written_byte_value as u32) << 24),
573 4 => (high_value & 0xFFFF_FF00) | (written_byte_value as u32),
574 5 => (high_value & 0xFFFF_00FF) | ((written_byte_value as u32) << 8),
575 6 => (high_value & 0xFF00_FFFF) | ((written_byte_value as u32) << 16),
576 7 => (high_value & 0x00FF_FFFF) | ((written_byte_value as u32) << 24),
577 _ => unreachable!("Invalid offset"),
578 };
579 let write_values = if offset < 4 {
580 [written_composed_value, high_value]
581 } else {
582 [low_value, written_composed_value]
583 };
584
585 if R::valid_for_write() {
586 self.std.range_check(self.table_8b_id, written_byte_value as i64, 1);
587 }
588 row.set_write_fields(
589 input.is_write,
590 written_composed_value,
591 written_byte_value,
592 write_values,
593 );
594
595 if input.is_write {
596 assert!(
597 R::valid_for_write(),
598 "Row type does not support write operations ({}) row:{irow} step:{step}",
599 std::any::type_name::<R>()
600 );
601 } else {
602 assert!(
603 R::valid_for_read(),
604 "Row type does not support read operations ({}) row:{irow} step:{step}",
605 std::any::type_name::<R>()
606 );
607 }
608 }
609}