1use std::sync::Arc;
2
3#[cfg(feature = "debug_mem_align")]
4use std::sync::Mutex;
5
6use fields::PrimeField64;
7use pil_std_lib::Std;
8
9use crate::{MemAlignInput, MemAlignRomSM, MemOp};
10use proofman_common::{AirInstance, FromTrace, ProofmanResult};
11use rayon::prelude::*;
12#[cfg(not(feature = "packed"))]
13use zisk_pil::{MemAlignTrace, MemAlignTraceRow};
14#[cfg(feature = "packed")]
15use zisk_pil::{MemAlignTracePacked, MemAlignTraceRowPacked};
16
17#[cfg(feature = "packed")]
18type MemAlignTraceRowType<F> = MemAlignTraceRowPacked<F>;
19#[cfg(feature = "packed")]
20type MemAlignTraceType<F> = MemAlignTracePacked<F>;
21
22#[cfg(not(feature = "packed"))]
23type MemAlignTraceRowType<F> = MemAlignTraceRow<F>;
24#[cfg(not(feature = "packed"))]
25type MemAlignTraceType<F> = MemAlignTrace<F>;
26
27const RC: usize = 2;
28const CHUNK_NUM: usize = 8;
29const CHUNKS_BY_RC: usize = CHUNK_NUM / RC;
30const CHUNK_BITS: usize = 8;
31const RC_BITS: u64 = (CHUNKS_BY_RC * CHUNK_BITS) as u64;
32const RC_MASK: u64 = (1 << RC_BITS) - 1;
33const OFFSET_MASK: u32 = 0x07;
34const OFFSET_BITS: u32 = 3;
35const CHUNK_BITS_MASK: u64 = (1 << CHUNK_BITS) - 1;
36
37const fn generate_allowed_offsets() -> [u8; CHUNK_NUM] {
38 let mut offsets = [0; CHUNK_NUM];
39 let mut i = 0;
40 while i < CHUNK_NUM {
41 offsets[i] = i as u8;
42 i += 1;
43 }
44 offsets
45}
46
47const ALLOWED_OFFSETS: [u8; CHUNK_NUM] = generate_allowed_offsets();
48const ALLOWED_WIDTHS: [u8; 4] = [1, 2, 4, 8];
49const DEFAULT_OFFSET: u8 = 0;
50const DEFAULT_WIDTH: u8 = 8;
51
52pub struct MemAlignSM<F: PrimeField64> {
53 /// PIL2 standard library
54 std: Arc<Std<F>>,
55
56 #[cfg(feature = "debug_mem_align")]
57 num_computed_rows: Mutex<usize>,
58
59 /// The table ID for the Mem Align ROM State Machine
60 table_id: usize,
61
62 /// The range ID for the byte range check
63 range_id: usize,
64}
65
66macro_rules! debug_info {
67 ($prefix:expr, $($arg:tt)*) => {
68 #[cfg(feature = "debug_mem_align")]
69 {
70 tracing::debug!(concat!("MemAlign: ",$prefix), $($arg)*);
71 }
72 };
73}
74
75impl<F: PrimeField64> MemAlignSM<F> {
76 pub fn new(std: Arc<Std<F>>) -> Arc<Self> {
77 // Get the table ID
78 let table_id =
79 std.get_virtual_table_id(MemAlignRomSM::TABLE_ID).expect("Failed to get table ID");
80 let range_id =
81 std.get_range_id(0, CHUNK_BITS_MASK as i64, None).expect("Failed to get range ID");
82
83 Arc::new(Self {
84 std: std.clone(),
85 #[cfg(feature = "debug_mem_align")]
86 num_computed_rows: Mutex::new(0),
87 table_id,
88 range_id,
89 })
90 }
91
92 pub fn prove_mem_align_op(
93 &self,
94 input: &MemAlignInput,
95 trace: &mut [MemAlignTraceRowType<F>],
96 ) -> usize {
97 let addr = input.addr;
98 let width = input.width;
99
100 // Compute the width
101 debug_assert!(
102 ALLOWED_WIDTHS.contains(&width),
103 "Width={width} is not allowed. Allowed widths are {ALLOWED_WIDTHS:?}"
104 );
105 let width = width as usize;
106
107 // Compute the offset
108 let offset = (addr & OFFSET_MASK) as u8;
109 debug_assert!(
110 ALLOWED_OFFSETS.contains(&offset),
111 "Offset={offset} is not allowed. Allowed offsets are {ALLOWED_OFFSETS:?}"
112 );
113 let offset = offset as usize;
114
115 #[cfg(feature = "debug_mem_align")]
116 let num_rows = self.num_computed_rows.lock().unwrap();
117
118 match (input.is_write, offset + width > CHUNK_NUM) {
119 (false, false) => {
120 /* RV with offset=2, width=4
121 +----+----+====+====+====+====+----+----+
122 | R0 | R1 | R2 | R3 | R4 | R5 | R6 | R7 |
123 +----+----+====+====+====+====+----+----+
124 ⇓
125 +----+----+====+====+====+====+----+----+
126 | V6 | V7 | V0 | V1 | V2 | V3 | V4 | V5 |
127 +----+----+====+====+====+====+----+----+
128 */
129 // Unaligned memory op information thrown into the bus
130 let step = input.step;
131 let value = input.value;
132
133 // Get the aligned address
134 let addr_read = addr >> OFFSET_BITS;
135
136 // Get the aligned value
137 let value_read = input.mem_values[0];
138
139 // Get the next pc and op size
140 let (next_pc, op_size) =
141 MemAlignRomSM::calculate_next_pc_and_op_size(MemOp::OneRead, offset, width);
142
143 // Update the row multiplicity of the operation
144 MemAlignRomSM::get_rows(&self.std, self.table_id, next_pc, op_size);
145
146 let mut read_row = MemAlignTraceRowType::default();
147 read_row.set_step(step);
148 read_row.set_addr(addr_read);
149 read_row.set_offset(DEFAULT_OFFSET);
150 read_row.set_width(DEFAULT_WIDTH);
151 read_row.set_reset(true);
152 read_row.set_sel_up_to_down(true);
153
154 let mut value_row = MemAlignTraceRowType::default();
155 value_row.set_step(step);
156 value_row.set_addr(addr_read);
157 value_row.set_offset(offset as u8);
158 value_row.set_width(width as u8);
159 value_row.set_pc(next_pc as u8);
160 value_row.set_sel_prove(true);
161
162 for i in 0..CHUNK_NUM {
163 read_row.set_reg(i, Self::get_byte(value_read, i, 0));
164 if i >= offset && i < offset + width {
165 read_row.set_sel(i, true);
166 }
167
168 value_row.set_reg(i, Self::get_byte(value, i, CHUNK_NUM - offset));
169 if i == offset {
170 value_row.set_sel(i, true);
171 }
172 }
173
174 let mut _value_read = value_read;
175 let mut _value = value;
176 for i in 0..RC {
177 read_row.set_value(i, (_value_read & RC_MASK) as u32);
178 value_row.set_value(i, (_value & RC_MASK) as u32);
179 _value_read >>= RC_BITS;
180 _value >>= RC_BITS;
181 }
182
183 #[rustfmt::skip]
184 debug_info!(
185 "\nOne Word Read\n\
186 Num Rows: {:?}\n\
187 Input: {:?}\n\
188 Value Read: {:?}\n\
189 Value: {:?}\n\
190 Flags Read: {:?}\n\
191 Flags Value: {:?}",
192 [*num_rows, *num_rows + 1],
193 input,
194 value_read.to_le_bytes(),
195 value.to_le_bytes(),
196 [
197 read_row.get_sel(0), read_row.get_sel(1), read_row.get_sel(2), read_row.get_sel(3),
198 read_row.get_sel(4), read_row.get_sel(5), read_row.get_sel(6), read_row.get_sel(7),
199 read_row.get_wr(), read_row.get_reset(), read_row.get_sel_up_to_down(), read_row.get_sel_down_to_up()
200 ],
201 [
202 value_row.get_sel(0), value_row.get_sel(1), value_row.get_sel(2), value_row.get_sel(3),
203 value_row.get_sel(4), value_row.get_sel(5), value_row.get_sel(6), value_row.get_sel(7),
204 value_row.get_wr(), value_row.get_reset(), value_row.get_sel_up_to_down(), value_row.get_sel_down_to_up()
205 ]
206 );
207
208 #[cfg(feature = "debug_mem_align")]
209 drop(num_rows);
210
211 // Prove the generated rows
212 trace[0] = read_row;
213 trace[1] = value_row;
214 2
215 }
216 (true, false) => {
217 /* RWV with offset=3, width=4
218 +----+----+----+====+====+====+====+----+
219 | R0 | R1 | R2 | R3 | R4 | R5 | R6 | R7 |
220 +----+----+----+====+====+====+====+----+
221 ⇓
222 +----+----+----+====+====+====+====+----+
223 | W0 | W1 | W2 | W3 | W4 | W5 | W6 | W7 |
224 +----+----+----+====+====+====+====+----+
225 ⇓
226 +----+----+----+====+====+====+====+----+
227 | V5 | V6 | V7 | V0 | V1 | V2 | V3 | V4 |
228 +----+----+----+====+====+====+====+----+
229 */
230
231 // Unaligned memory op information thrown into the bus
232 let step = input.step;
233 let value = input.value;
234
235 // Get the aligned address
236 let addr_read = addr >> OFFSET_BITS;
237
238 // Get the aligned value
239 let value_read = input.mem_values[0];
240
241 // Get the next pc
242 let (next_pc, op_size) =
243 MemAlignRomSM::calculate_next_pc_and_op_size(MemOp::OneWrite, offset, width);
244
245 // Update the row multiplicity of the operation
246 MemAlignRomSM::get_rows(&self.std, self.table_id, next_pc, op_size);
247
248 // Compute the write value
249 let value_write = {
250 // with:1 offset:4
251 let width_bytes: u64 = (1 << (width * CHUNK_BITS)) - 1;
252
253 let mask: u64 = width_bytes << (offset * CHUNK_BITS);
254
255 // Get the first width bytes of the unaligned value
256 let value_to_write = (value & width_bytes) << (offset * CHUNK_BITS);
257
258 // Write zeroes to value_read from offset to offset + width
259 // and add the value to write to the value read
260 (value_read & !mask) | value_to_write
261 };
262
263 let mut read_row = MemAlignTraceRowType::default();
264 read_row.set_step(step);
265 read_row.set_addr(addr_read);
266 read_row.set_offset(DEFAULT_OFFSET);
267 read_row.set_width(DEFAULT_WIDTH);
268 read_row.set_reset(true);
269 read_row.set_sel_up_to_down(true);
270
271 let mut write_row = MemAlignTraceRowType::default();
272 write_row.set_step(step + 1);
273 write_row.set_addr(addr_read);
274 write_row.set_offset(DEFAULT_OFFSET);
275 write_row.set_width(DEFAULT_WIDTH);
276 write_row.set_wr(true);
277 write_row.set_pc(next_pc as u8);
278 write_row.set_sel_up_to_down(true);
279
280 let mut value_row = MemAlignTraceRowType::default();
281 value_row.set_step(step);
282 value_row.set_addr(addr_read);
283 value_row.set_offset(offset as u8);
284 value_row.set_width(width as u8);
285 value_row.set_wr(true);
286 value_row.set_pc(next_pc as u8 + 1);
287 value_row.set_sel_prove(true);
288
289 for i in 0..CHUNK_NUM {
290 read_row.set_reg(i, Self::get_byte(value_read, i, 0));
291 if i < offset || i >= offset + width {
292 read_row.set_sel(i, true);
293 }
294
295 let write_reg = Self::get_byte(value_write, i, 0);
296 write_row.set_reg(i, write_reg);
297 if i >= offset && i < offset + width {
298 write_row.set_sel(i, true);
299 }
300
301 value_row.set_reg(
302 i,
303 if i >= offset && i < offset + width {
304 write_reg
305 } else {
306 Self::get_byte(value, i, CHUNK_NUM - offset)
307 },
308 );
309 if i == offset {
310 value_row.set_sel(i, true);
311 }
312 }
313
314 let mut _value_read = value_read;
315 let mut _value_write = value_write;
316 let mut _value = value;
317 for i in 0..RC {
318 read_row.set_value(i, (_value_read & RC_MASK) as u32);
319 write_row.set_value(i, (_value_write & RC_MASK) as u32);
320 value_row.set_value(i, (_value & RC_MASK) as u32);
321 _value_read >>= RC_BITS;
322 _value_write >>= RC_BITS;
323 _value >>= RC_BITS;
324 }
325
326 #[rustfmt::skip]
327 debug_info!(
328 "\nOne Word Write\n\
329 Num Rows: {:?}\n\
330 Input: {:?}\n\
331 Value Read: {:?}\n\
332 Value Write: {:?}\n\
333 Value: {:?}\n\
334 Flags Read: {:?}\n\
335 Flags Write: {:?}\n\
336 Flags Value: {:?}",
337 [*num_rows, *num_rows + 2],
338 input,
339 value_read.to_le_bytes(),
340 value_write.to_le_bytes(),
341 value.to_le_bytes(),
342 [
343 read_row.get_sel(0), read_row.get_sel(1), read_row.get_sel(2), read_row.get_sel(3),
344 read_row.get_sel(4), read_row.get_sel(5), read_row.get_sel(6), read_row.get_sel(7),
345 read_row.get_wr(), read_row.get_reset(), read_row.get_sel_up_to_down(), read_row.get_sel_down_to_up()
346 ],
347 [
348 write_row.get_sel(0), write_row.get_sel(1), write_row.get_sel(2), write_row.get_sel(3),
349 write_row.get_sel(4), write_row.get_sel(5), write_row.get_sel(6), write_row.get_sel(7),
350 write_row.get_wr(), write_row.get_reset(), write_row.get_sel_up_to_down(), write_row.get_sel_down_to_up()
351 ],
352 [
353 value_row.get_sel(0), value_row.get_sel(1), value_row.get_sel(2), value_row.get_sel(3),
354 value_row.get_sel(4), value_row.get_sel(5), value_row.get_sel(6), value_row.get_sel(7),
355 value_row.get_wr(), value_row.get_reset(), value_row.get_sel_up_to_down(), value_row.get_sel_down_to_up()
356 ]
357 );
358
359 #[cfg(feature = "debug_mem_align")]
360 drop(num_rows);
361
362 // Prove the generated rows
363 trace[0] = read_row;
364 trace[1] = write_row;
365 trace[2] = value_row;
366 3
367 }
368 (false, true) => {
369 /* RVR with offset=5, width=8
370 +----+----+----+----+----+====+====+====+
371 | R0 | R1 | R2 | R3 | R4 | R5 | R6 | R7 |
372 +----+----+----+----+----+====+====+====+
373 ⇓
374 +====+====+====+====+====+====+====+====+
375 | V3 | V4 | V5 | V6 | V7 | V0 | V1 | V2 |
376 +====+====+====+====+====+====+====+====+
377 ⇓
378 +====+====+====+====+====+----+----+----+
379 | R0 | R1 | R2 | R3 | R4 | R5 | R6 | R7 |
380 +====+====+====+====+====+----+----+----+
381 */
382
383 // Unaligned memory op information thrown into the bus
384 let step = input.step;
385 let value = input.value;
386
387 // Compute the remaining bytes
388 let rem_bytes = (offset + width) % CHUNK_NUM;
389
390 // Get the aligned address
391 let addr_first_read = addr >> OFFSET_BITS;
392 let addr_second_read = addr_first_read + 1;
393
394 // Get the aligned value
395 let value_first_read = input.mem_values[0];
396 let value_second_read = input.mem_values[1];
397
398 // Get the next pc
399 let (next_pc, op_size) =
400 MemAlignRomSM::calculate_next_pc_and_op_size(MemOp::TwoReads, offset, width);
401
402 // Update the row multiplicity of the operation
403 MemAlignRomSM::get_rows(&self.std, self.table_id, next_pc, op_size);
404
405 let mut first_read_row = MemAlignTraceRowType::default();
406 first_read_row.set_step(step);
407 first_read_row.set_addr(addr_first_read);
408 first_read_row.set_offset(DEFAULT_OFFSET);
409 first_read_row.set_width(DEFAULT_WIDTH);
410 first_read_row.set_reset(true);
411 first_read_row.set_sel_up_to_down(true);
412
413 let mut value_row = MemAlignTraceRowType::default();
414 value_row.set_step(step);
415 value_row.set_addr(addr_first_read);
416 value_row.set_offset(offset as u8);
417 value_row.set_width(width as u8);
418 value_row.set_pc(next_pc as u8);
419 value_row.set_sel_prove(true);
420
421 let mut second_read_row = MemAlignTraceRowType::default();
422 second_read_row.set_step(step);
423 second_read_row.set_addr(addr_second_read);
424 second_read_row.set_delta_addr(1);
425 second_read_row.set_offset(DEFAULT_OFFSET);
426 second_read_row.set_width(DEFAULT_WIDTH);
427 second_read_row.set_pc(next_pc as u8 + 1);
428 second_read_row.set_sel_down_to_up(true);
429
430 for i in 0..CHUNK_NUM {
431 first_read_row.set_reg(i, Self::get_byte(value_first_read, i, 0));
432 if i >= offset {
433 first_read_row.set_sel(i, true);
434 }
435
436 value_row.set_reg(i, Self::get_byte(value, i, CHUNK_NUM - offset));
437
438 if i == offset {
439 value_row.set_sel(i, true);
440 }
441
442 second_read_row.set_reg(i, Self::get_byte(value_second_read, i, 0));
443 if i < rem_bytes {
444 second_read_row.set_sel(i, true);
445 }
446 }
447
448 let mut _value_first_read = value_first_read;
449 let mut _value = value;
450 let mut _value_second_read = value_second_read;
451 for i in 0..RC {
452 first_read_row.set_value(i, (_value_first_read & RC_MASK) as u32);
453 value_row.set_value(i, (_value & RC_MASK) as u32);
454 second_read_row.set_value(i, (_value_second_read & RC_MASK) as u32);
455 _value_first_read >>= RC_BITS;
456 _value >>= RC_BITS;
457 _value_second_read >>= RC_BITS;
458 }
459
460 #[rustfmt::skip]
461 debug_info!(
462 "\nTwo Words Read\n\
463 Num Rows: {:?}\n\
464 Input: {:?}\n\
465 Value First Read: {:?}\n\
466 Value: {:?}\n\
467 Value Second Read: {:?}\n\
468 Flags First Read: {:?}\n\
469 Flags Value: {:?}\n\
470 Flags Second Read: {:?}",
471 [*num_rows, *num_rows + 2],
472 input,
473 value_first_read.to_le_bytes(),
474 value.to_le_bytes(),
475 value_second_read.to_le_bytes(),
476 [
477 first_read_row.get_sel(0), first_read_row.get_sel(1), first_read_row.get_sel(2), first_read_row.get_sel(3),
478 first_read_row.get_sel(4), first_read_row.get_sel(5), first_read_row.get_sel(6), first_read_row.get_sel(7),
479 first_read_row.get_wr(), first_read_row.get_reset(), first_read_row.get_sel_up_to_down(), first_read_row.get_sel_down_to_up()
480 ],
481 [
482 value_row.get_sel(0), value_row.get_sel(1), value_row.get_sel(2), value_row.get_sel(3),
483 value_row.get_sel(4), value_row.get_sel(5), value_row.get_sel(6), value_row.get_sel(7),
484 value_row.get_wr(), value_row.get_reset(), value_row.get_sel_up_to_down(), value_row.get_sel_down_to_up()
485 ],
486 [
487 second_read_row.get_sel(0), second_read_row.get_sel(1), second_read_row.get_sel(2), second_read_row.get_sel(3),
488 second_read_row.get_sel(4), second_read_row.get_sel(5), second_read_row.get_sel(6), second_read_row.get_sel(7),
489 second_read_row.get_wr(), second_read_row.get_reset(), second_read_row.get_sel_up_to_down(), second_read_row.get_sel_down_to_up()
490 ]
491 );
492
493 #[cfg(feature = "debug_mem_align")]
494 drop(num_rows);
495
496 // Prove the generated rows
497 trace[0] = first_read_row;
498 trace[1] = value_row;
499 trace[2] = second_read_row;
500 3
501 }
502 (true, true) => {
503 /* RWVWR with offset=6, width=4
504 +----+----+----+----+----+----+====+====+
505 | R0 | R1 | R2 | R3 | R4 | R5 | R6 | R7 |
506 +----+----+----+----+----+----+====+====+
507 ⇓
508 +----+----+----+----+----+----+====+====+
509 | W0 | W1 | W2 | W3 | W4 | W5 | W6 | W7 |
510 +----+----+----+----+----+----+====+====+
511 ⇓
512 +====+====+----+----+----+----+====+====+
513 | V2 | V3 | V4 | V5 | V6 | V7 | V0 | V1 |
514 +====+====+----+----+----+----+====+====+
515 ⇓
516 +====+====+----+----+----+----+----+----+
517 | W0 | W1 | W2 | W3 | W4 | W5 | W6 | W7 |
518 +====+====+----+----+----+----+----+----+
519 ⇓
520 +====+====+----+----+----+----+----+----+
521 | R0 | R1 | R2 | R3 | R4 | R5 | R6 | R7 |
522 +====+====+----+----+----+----+----+----+
523 */
524 // Unaligned memory op information thrown into the bus
525 let step = input.step;
526 let value = input.value;
527
528 // Compute the shift
529 let rem_bytes = (offset + width) % CHUNK_NUM;
530
531 // Get the aligned address
532 let addr_first_read_write = addr >> OFFSET_BITS;
533 let addr_second_read_write = addr_first_read_write + 1;
534
535 // Get the first aligned value
536 let value_first_read = input.mem_values[0];
537
538 // Recompute the first write value
539 let value_first_write = {
540 // Normalize the width
541 let width_norm = CHUNK_NUM - offset;
542
543 let width_bytes: u64 = (1 << (width_norm * CHUNK_BITS)) - 1;
544
545 let mask: u64 = width_bytes << (offset * CHUNK_BITS);
546
547 // Get the first width bytes of the unaligned value
548 let value_to_write = (value & width_bytes) << (offset * CHUNK_BITS);
549
550 // Write zeroes to value_read from offset to offset + width
551 // and add the value to write to the value read
552 (value_first_read & !mask) | value_to_write
553 };
554
555 // Get the second aligned value
556 let value_second_read = input.mem_values[1];
557
558 // Compute the second write value
559 let value_second_write = {
560 // Normalize the width
561 let width_norm = CHUNK_NUM - offset;
562
563 let mask: u64 = (1 << (rem_bytes * CHUNK_BITS)) - 1;
564
565 // Get the first width bytes of the unaligned value
566 let value_to_write = (value >> (width_norm * CHUNK_BITS)) & mask;
567
568 // Write zeroes to value_read from 0 to offset + width
569 // and add the value to write to the value read
570 (value_second_read & !mask) | value_to_write
571 };
572
573 // Get the next pc
574 let (next_pc, op_size) =
575 MemAlignRomSM::calculate_next_pc_and_op_size(MemOp::TwoWrites, offset, width);
576
577 // Update the row multiplicity of the operation
578 MemAlignRomSM::get_rows(&self.std, self.table_id, next_pc, op_size);
579
580 // RWVWR
581 let mut first_read_row = MemAlignTraceRowType::default();
582 first_read_row.set_step(step);
583 first_read_row.set_addr(addr_first_read_write);
584 first_read_row.set_offset(DEFAULT_OFFSET);
585 first_read_row.set_width(DEFAULT_WIDTH);
586 first_read_row.set_reset(true);
587 first_read_row.set_sel_up_to_down(true);
588
589 let mut first_write_row = MemAlignTraceRowType::<F>::default();
590 first_write_row.set_step(step + 1);
591 first_write_row.set_addr(addr_first_read_write);
592 first_write_row.set_offset(DEFAULT_OFFSET);
593 first_write_row.set_width(DEFAULT_WIDTH);
594 first_write_row.set_wr(true);
595 first_write_row.set_pc(next_pc as u8);
596 first_write_row.set_sel_up_to_down(true);
597
598 let mut value_row = MemAlignTraceRowType::default();
599 value_row.set_step(step);
600 value_row.set_addr(addr_first_read_write);
601 value_row.set_offset(offset as u8);
602 value_row.set_width(width as u8);
603 value_row.set_wr(true);
604 value_row.set_pc(next_pc as u8 + 1);
605 value_row.set_sel_prove(true);
606
607 let mut second_write_row = MemAlignTraceRowType::default();
608 second_write_row.set_step(step + 1);
609 second_write_row.set_addr(addr_second_read_write);
610 second_write_row.set_delta_addr(1);
611 second_write_row.set_offset(DEFAULT_OFFSET);
612 second_write_row.set_width(DEFAULT_WIDTH);
613 second_write_row.set_wr(true);
614 second_write_row.set_pc(next_pc as u8 + 2);
615 second_write_row.set_sel_down_to_up(true);
616
617 let mut second_read_row = MemAlignTraceRowType::default();
618 second_read_row.set_step(step);
619 second_read_row.set_addr(addr_second_read_write);
620 second_read_row.set_offset(DEFAULT_OFFSET);
621 second_read_row.set_width(DEFAULT_WIDTH);
622 second_read_row.set_pc(next_pc as u8 + 3);
623 second_read_row.set_reset(false);
624 second_read_row.set_sel_down_to_up(true);
625
626 for i in 0..CHUNK_NUM {
627 first_read_row.set_reg(i, Self::get_byte(value_first_read, i, 0));
628 if i < offset {
629 first_read_row.set_sel(i, true);
630 }
631
632 first_write_row.set_reg(i, Self::get_byte(value_first_write, i, 0));
633 if i >= offset {
634 first_write_row.set_sel(i, true);
635 }
636
637 value_row.set_reg(i, {
638 if i < rem_bytes {
639 second_write_row.get_reg(i)
640 } else if i >= offset {
641 first_write_row.get_reg(i)
642 } else {
643 Self::get_byte(value, i, CHUNK_NUM - offset)
644 }
645 });
646 if i == offset {
647 value_row.set_sel(i, true);
648 }
649
650 second_write_row.set_reg(i, Self::get_byte(value_second_write, i, 0));
651 if i < rem_bytes {
652 second_write_row.set_sel(i, true);
653 }
654
655 second_read_row.set_reg(i, Self::get_byte(value_second_read, i, 0));
656 if i >= rem_bytes {
657 second_read_row.set_sel(i, true);
658 }
659 }
660
661 let mut _value_first_read = value_first_read;
662 let mut _value_first_write = value_first_write;
663 let mut _value = value;
664 let mut _value_second_write = value_second_write;
665 let mut _value_second_read = value_second_read;
666 for i in 0..RC {
667 first_read_row.set_value(i, (_value_first_read & RC_MASK) as u32);
668 first_write_row.set_value(i, (_value_first_write & RC_MASK) as u32);
669 value_row.set_value(i, (_value & RC_MASK) as u32);
670 second_write_row.set_value(i, (_value_second_write & RC_MASK) as u32);
671 second_read_row.set_value(i, (_value_second_read & RC_MASK) as u32);
672 _value_first_read >>= RC_BITS;
673 _value_first_write >>= RC_BITS;
674 _value >>= RC_BITS;
675 _value_second_write >>= RC_BITS;
676 _value_second_read >>= RC_BITS;
677 }
678
679 #[rustfmt::skip]
680 debug_info!(
681 "\nTwo Words Write\n\
682 Num Rows: {:?}\n\
683 Input: {:?}\n\
684 Value First Read: {:?}\n\
685 Value First Write: {:?}\n\
686 Value: {:?}\n\
687 Value Second Read: {:?}\n\
688 Value Second Write: {:?}\n\
689 Flags First Read: {:?}\n\
690 Flags First Write: {:?}\n\
691 Flags Value: {:?}\n\
692 Flags Second Write: {:?}\n\
693 Flags Second Read: {:?}",
694 [*num_rows, *num_rows + 4],
695 input,
696 value_first_read.to_le_bytes(),
697 value_first_write.to_le_bytes(),
698 value.to_le_bytes(),
699 value_second_write.to_le_bytes(),
700 value_second_read.to_le_bytes(),
701 [
702 first_read_row.get_sel(0), first_read_row.get_sel(1), first_read_row.get_sel(2), first_read_row.get_sel(3),
703 first_read_row.get_sel(4), first_read_row.get_sel(5), first_read_row.get_sel(6), first_read_row.get_sel(7),
704 first_read_row.get_wr(), first_read_row.get_reset(), first_read_row.get_sel_up_to_down(), first_read_row.get_sel_down_to_up()
705 ],
706 [
707 first_write_row.get_sel(0), first_write_row.get_sel(1), first_write_row.get_sel(2), first_write_row.get_sel(3),
708 first_write_row.get_sel(4), first_write_row.get_sel(5), first_write_row.get_sel(6), first_write_row.get_sel(7),
709 first_write_row.get_wr(), first_write_row.get_reset(), first_write_row.get_sel_up_to_down(), first_write_row.get_sel_down_to_up()
710 ],
711 [
712 value_row.get_sel(0), value_row.get_sel(1), value_row.get_sel(2), value_row.get_sel(3),
713 value_row.get_sel(4), value_row.get_sel(5), value_row.get_sel(6), value_row.get_sel(7),
714 value_row.get_wr(), value_row.get_reset(), value_row.get_sel_up_to_down(), value_row.get_sel_down_to_up()
715 ],
716 [
717 second_write_row.get_sel(0), second_write_row.get_sel(1), second_write_row.get_sel(2), second_write_row.get_sel(3),
718 second_write_row.get_sel(4), second_write_row.get_sel(5), second_write_row.get_sel(6), second_write_row.get_sel(7),
719 second_write_row.get_wr(), second_write_row.get_reset(), second_write_row.get_sel_up_to_down(), second_write_row.get_sel_down_to_up()
720 ],
721 [
722 second_read_row.get_sel(0), second_read_row.get_sel(1), second_read_row.get_sel(2), second_read_row.get_sel(3),
723 second_read_row.get_sel(4), second_read_row.get_sel(5), second_read_row.get_sel(6), second_read_row.get_sel(7),
724 second_read_row.get_wr(), second_read_row.get_reset(), second_read_row.get_sel_up_to_down(), second_read_row.get_sel_down_to_up()
725 ]
726 );
727
728 #[cfg(feature = "debug_mem_align")]
729 drop(num_rows);
730
731 // Prove the generated rows
732 trace[0] = first_read_row;
733 trace[1] = first_write_row;
734 trace[2] = value_row;
735 trace[3] = second_write_row;
736 trace[4] = second_read_row;
737 5
738 }
739 }
740 }
741
742 fn get_byte(value: u64, index: usize, offset: usize) -> u8 {
743 let chunk = (offset + index) % CHUNK_NUM;
744 ((value >> (chunk * CHUNK_BITS)) & CHUNK_BITS_MASK) as u8
745 }
746
747 pub fn compute_witness(
748 &self,
749 mem_ops: &[Vec<MemAlignInput>],
750 used_rows: usize,
751 trace_buffer: Vec<F>,
752 ) -> ProofmanResult<AirInstance<F>> {
753 let mut trace = MemAlignTraceType::new_from_vec(trace_buffer)?;
754 let mut reg_range_check = vec![0u32; 1 << CHUNK_BITS];
755
756 let num_rows = trace.num_rows();
757
758 tracing::debug!(
759 "··· Creating Mem Align instance [{} / {} rows filled {:.2}%]",
760 used_rows,
761 num_rows,
762 used_rows as f64 / num_rows as f64 * 100.0
763 );
764
765 let mut trace_rows = &mut trace.buffer[..];
766 let mut par_traces = Vec::new();
767 let mut inputs_indexes = Vec::new();
768 let mut total_index = 0;
769 for (i, inner_memp_ops) in mem_ops.iter().enumerate() {
770 for (j, input) in inner_memp_ops.iter().enumerate() {
771 let addr = input.addr;
772 let width = input.width as usize;
773 let offset = (addr & OFFSET_MASK) as usize;
774 let n_rows = match (input.is_write, offset + width > CHUNK_NUM) {
775 (false, false) => 2,
776 (true, false) => 3,
777 (false, true) => 3,
778 (true, true) => 5,
779 };
780 total_index += n_rows;
781 let (head, tail) = trace_rows.split_at_mut(n_rows);
782 par_traces.push(head);
783 inputs_indexes.push((i, j));
784 trace_rows = tail;
785 }
786 }
787
788 // Prove the memory operations in parallel
789 par_traces.into_par_iter().enumerate().for_each(|(index, trace)| {
790 let input_index = inputs_indexes[index];
791 let input = &mem_ops[input_index.0][input_index.1];
792 self.prove_mem_align_op(input, trace);
793 });
794
795 // Iterate over all traces to set range checks
796 trace.buffer[0..total_index].iter_mut().for_each(|row| {
797 for j in 0..CHUNK_NUM {
798 reg_range_check[row.get_reg(j) as usize] += 1;
799 }
800 });
801
802 let padding_size = num_rows - total_index;
803 let mut padding_row = MemAlignTraceRowType::default();
804 padding_row.set_reset(true);
805
806 // Store the padding rows
807 trace.buffer[total_index..num_rows].par_iter_mut().for_each(|slot| *slot = padding_row);
808
809 // Compute the program multiplicity
810 self.std.inc_virtual_row(self.table_id, MemAlignRomSM::PADDING_ROW, padding_size as u64);
811
812 reg_range_check[0] += CHUNK_NUM as u32 * padding_size as u32;
813 self.update_std_range_check(reg_range_check);
814
815 Ok(AirInstance::new_from_trace(FromTrace::new(&mut trace)))
816 }
817
818 fn update_std_range_check(&self, reg_range_check: Vec<u32>) {
819 // Perform the range checks
820 self.std.range_checks(self.range_id, reg_range_check);
821 }
822}