use std::{
borrow::{Borrow, BorrowMut},
mem::{size_of, MaybeUninit},
};
use crate::{air::SP1CoreAirBuilder, utils::next_multiple_of_32};
use super::MemoryAccessCols;
use hashbrown::HashMap;
use itertools::Itertools;
use slop_air::{Air, BaseAir};
use slop_algebra::{AbstractField, PrimeField32};
use slop_matrix::Matrix;
use sp1_core_executor::{
events::{ByteLookupEvent, ByteRecord, MemoryReadRecord, MemoryRecordEnum},
ByteOpcode, ExecutionRecord, Program,
};
use sp1_derive::AlignedBorrow;
use sp1_hypercube::air::MachineAir;
pub(crate) const NUM_MEMORY_BUMP_COLS: usize = size_of::<MemoryBumpCols<u8>>();
#[derive(AlignedBorrow, Clone, Copy)]
#[repr(C)]
pub struct MemoryBumpCols<T: Copy> {
pub access: MemoryAccessCols<T>,
pub clk_32_48: T,
pub clk_24_32: T,
pub clk_16_24: T,
pub clk_0_16: T,
pub addr: T,
pub is_real: T,
}
pub struct MemoryBumpChip {}
impl MemoryBumpChip {
pub const fn new() -> Self {
Self {}
}
}
impl<F> BaseAir<F> for MemoryBumpChip {
fn width(&self) -> usize {
NUM_MEMORY_BUMP_COLS
}
}
impl<F: PrimeField32> MachineAir<F> for MemoryBumpChip {
type Record = ExecutionRecord;
type Program = Program;
fn name(&self) -> &'static str {
"MemoryBump"
}
fn generate_dependencies(&self, input: &Self::Record, output: &mut Self::Record) {
let chunk_size = 1;
let event_iter = input.bump_memory_events.chunks(chunk_size);
let blu_batches = event_iter
.map(|events| {
let mut blu: HashMap<ByteLookupEvent, usize> = HashMap::new();
events.iter().for_each(|(event, addr, is_refresh)| {
let mut row = [F::zero(); NUM_MEMORY_BUMP_COLS];
let cols: &mut MemoryBumpCols<F> = row.as_mut_slice().borrow_mut();
let value = event.prev_value();
let prev_timestamp = event.previous_record().timestamp;
let mut timestamp = event.current_record().timestamp;
if !is_refresh {
timestamp = (timestamp >> 24) << 24;
}
let bump_event = MemoryRecordEnum::Read(MemoryReadRecord {
value,
prev_timestamp,
timestamp,
prev_page_prot_record: None,
});
cols.access.populate(bump_event, &mut blu);
blu.add_u16_range_checks(&[
(timestamp & 0xFFFF) as u16,
((timestamp >> 32) & 0xFFFF) as u16,
]);
blu.add_u8_range_checks(&[
((timestamp >> 16) & 0xFF) as u8,
((timestamp >> 24) & 0xFF) as u8,
]);
blu.add_byte_lookup_event(ByteLookupEvent {
opcode: ByteOpcode::LTU,
a: 1,
b: *addr as u8,
c: 32,
});
cols.is_real = F::one();
});
blu
})
.collect::<Vec<_>>();
output.add_byte_lookup_events_from_maps(blu_batches.iter().collect_vec());
}
fn num_rows(&self, input: &Self::Record) -> Option<usize> {
let nb_rows = input.bump_memory_events.len();
let size_log2 = input.fixed_log2_rows::<F, _>(self);
Some(next_multiple_of_32(nb_rows, size_log2))
}
fn generate_trace_into(
&self,
input: &ExecutionRecord,
_output: &mut ExecutionRecord,
buffer: &mut [MaybeUninit<F>],
) {
let chunk_size = 1;
let padded_nb_rows = <MemoryBumpChip as MachineAir<F>>::num_rows(self, input).unwrap();
let num_event_rows = input.bump_memory_events.len();
unsafe {
let padding_start = num_event_rows * NUM_MEMORY_BUMP_COLS;
let padding_size = (padded_nb_rows - num_event_rows) * NUM_MEMORY_BUMP_COLS;
if padding_size > 0 {
core::ptr::write_bytes(buffer[padding_start..].as_mut_ptr(), 0, padding_size);
}
}
let buffer_ptr = buffer.as_mut_ptr() as *mut F;
let values = unsafe {
core::slice::from_raw_parts_mut(buffer_ptr, num_event_rows * NUM_MEMORY_BUMP_COLS)
};
values.chunks_mut(chunk_size * NUM_MEMORY_BUMP_COLS).enumerate().for_each(|(i, rows)| {
rows.chunks_mut(NUM_MEMORY_BUMP_COLS).enumerate().for_each(|(j, row)| {
let idx = i * chunk_size + j;
let cols: &mut MemoryBumpCols<F> = row.borrow_mut();
if idx < input.bump_memory_events.len() {
let mut byte_lookup_events = Vec::new();
let (event, addr, is_refresh) = input.bump_memory_events[idx];
let value = event.prev_value();
let prev_timestamp = event.previous_record().timestamp;
let mut timestamp = event.current_record().timestamp;
if !is_refresh {
timestamp = (timestamp >> 24) << 24;
}
let bump_event = MemoryRecordEnum::Read(MemoryReadRecord {
value,
prev_timestamp,
timestamp,
prev_page_prot_record: None,
});
cols.access.populate(bump_event, &mut byte_lookup_events);
cols.clk_0_16 = F::from_canonical_u16((timestamp & 0xFFFF) as u16);
cols.clk_16_24 = F::from_canonical_u8(((timestamp >> 16) & 0xFF) as u8);
cols.clk_24_32 = F::from_canonical_u8(((timestamp >> 24) & 0xFF) as u8);
cols.clk_32_48 = F::from_canonical_u16(((timestamp >> 32) & 0xFFFF) as u16);
cols.addr = F::from_canonical_u64(addr);
cols.is_real = F::one();
}
})
});
}
fn included(&self, shard: &Self::Record) -> bool {
shard.cpu_event_count != 0
}
}
impl<AB> Air<AB> for MemoryBumpChip
where
AB: SP1CoreAirBuilder,
{
fn eval(&self, builder: &mut AB) {
let main = builder.main();
let local = main.row_slice(0);
let local: &MemoryBumpCols<AB::Var> = (*local).borrow();
builder.assert_bool(local.is_real);
builder.slice_range_check_u16(&[local.clk_0_16, local.clk_32_48], local.is_real);
builder.slice_range_check_u8(&[local.clk_16_24, local.clk_24_32], local.is_real);
builder.send_byte(
AB::Expr::from_canonical_u32(ByteOpcode::LTU as u32),
AB::Expr::one(),
local.addr,
AB::Expr::from_canonical_u8(32),
local.is_real,
);
builder.eval_memory_access_read(
local.clk_24_32 + local.clk_32_48 * AB::Expr::from_canonical_u32(1 << 8),
local.clk_0_16 + local.clk_16_24 * AB::Expr::from_canonical_u32(1 << 16),
&[local.addr.into(), AB::Expr::zero(), AB::Expr::zero()],
local.access,
local.is_real,
);
}
}