use std::{
array::from_fn,
borrow::{Borrow, BorrowMut},
cell::RefCell,
iter::{once, zip},
marker::PhantomData,
sync::Arc,
};
use itertools::izip;
use openvm_circuit::{
arch::{
AdapterAirContext, AdapterRuntimeContext, ExecutionBridge, ExecutionBus, ExecutionState,
Result, VecHeapAdapterInterface, VmAdapterAir, VmAdapterChip, VmAdapterInterface,
},
system::{
memory::{
offline_checker::{MemoryBridge, MemoryReadAuxCols, MemoryWriteAuxCols},
MemoryAddress, MemoryAuxColsFactory, MemoryController, MemoryControllerRef,
MemoryReadRecord, MemoryWriteRecord,
},
program::ProgramBus,
},
};
use openvm_circuit_primitives::bitwise_op_lookup::{
BitwiseOperationLookupBus, BitwiseOperationLookupChip,
};
use openvm_circuit_primitives_derive::AlignedBorrow;
use openvm_instructions::{
instruction::Instruction,
riscv::{RV32_MEMORY_AS, RV32_REGISTER_AS},
};
use openvm_rv32im_circuit::adapters::{
abstract_compose, read_rv32_register, RV32_CELL_BITS, RV32_REGISTER_NUM_LIMBS,
};
use openvm_stark_backend::{
interaction::InteractionBuilder,
p3_air::BaseAir,
p3_field::{AbstractField, Field, PrimeField32},
};
#[derive(Debug, Clone)]
pub struct Rv32VecHeapAdapterChip<
F: Field,
const NUM_READS: usize,
const BLOCKS_PER_READ: usize,
const BLOCKS_PER_WRITE: usize,
const READ_SIZE: usize,
const WRITE_SIZE: usize,
> {
pub air:
Rv32VecHeapAdapterAir<NUM_READS, BLOCKS_PER_READ, BLOCKS_PER_WRITE, READ_SIZE, WRITE_SIZE>,
pub bitwise_lookup_chip: Arc<BitwiseOperationLookupChip<RV32_CELL_BITS>>,
_marker: PhantomData<F>,
}
impl<
F: PrimeField32,
const NUM_READS: usize,
const BLOCKS_PER_READ: usize,
const BLOCKS_PER_WRITE: usize,
const READ_SIZE: usize,
const WRITE_SIZE: usize,
>
Rv32VecHeapAdapterChip<F, NUM_READS, BLOCKS_PER_READ, BLOCKS_PER_WRITE, READ_SIZE, WRITE_SIZE>
{
pub fn new(
execution_bus: ExecutionBus,
program_bus: ProgramBus,
memory_controller: MemoryControllerRef<F>,
bitwise_lookup_chip: Arc<BitwiseOperationLookupChip<RV32_CELL_BITS>>,
) -> Self {
assert!(NUM_READS <= 2);
let memory_controller = RefCell::borrow(&memory_controller);
let memory_bridge = memory_controller.memory_bridge();
let address_bits = memory_controller.mem_config().pointer_max_bits;
assert!(
RV32_CELL_BITS * RV32_REGISTER_NUM_LIMBS - address_bits < RV32_CELL_BITS,
"address_bits={address_bits} needs to be large enough for high limb range check"
);
Self {
air: Rv32VecHeapAdapterAir {
execution_bridge: ExecutionBridge::new(execution_bus, program_bus),
memory_bridge,
bus: bitwise_lookup_chip.bus(),
address_bits,
},
bitwise_lookup_chip,
_marker: PhantomData,
}
}
}
#[derive(Clone, Debug)]
pub struct Rv32VecHeapReadRecord<
F: Field,
const NUM_READS: usize,
const BLOCKS_PER_READ: usize,
const READ_SIZE: usize,
> {
pub rs: [MemoryReadRecord<F, RV32_REGISTER_NUM_LIMBS>; NUM_READS],
pub rd: MemoryReadRecord<F, RV32_REGISTER_NUM_LIMBS>,
pub rd_val: F,
pub reads: [[MemoryReadRecord<F, READ_SIZE>; BLOCKS_PER_READ]; NUM_READS],
}
#[derive(Clone, Debug)]
pub struct Rv32VecHeapWriteRecord<F: Field, const BLOCKS_PER_WRITE: usize, const WRITE_SIZE: usize>
{
pub from_state: ExecutionState<u32>,
pub writes: [MemoryWriteRecord<F, WRITE_SIZE>; BLOCKS_PER_WRITE],
}
#[repr(C)]
#[derive(AlignedBorrow)]
pub struct Rv32VecHeapAdapterCols<
T,
const NUM_READS: usize,
const BLOCKS_PER_READ: usize,
const BLOCKS_PER_WRITE: usize,
const READ_SIZE: usize,
const WRITE_SIZE: usize,
> {
pub from_state: ExecutionState<T>,
pub rs_ptr: [T; NUM_READS],
pub rd_ptr: T,
pub rs_val: [[T; RV32_REGISTER_NUM_LIMBS]; NUM_READS],
pub rd_val: [T; RV32_REGISTER_NUM_LIMBS],
pub rs_read_aux: [MemoryReadAuxCols<T, RV32_REGISTER_NUM_LIMBS>; NUM_READS],
pub rd_read_aux: MemoryReadAuxCols<T, RV32_REGISTER_NUM_LIMBS>,
pub reads_aux: [[MemoryReadAuxCols<T, READ_SIZE>; BLOCKS_PER_READ]; NUM_READS],
pub writes_aux: [MemoryWriteAuxCols<T, WRITE_SIZE>; BLOCKS_PER_WRITE],
}
#[allow(dead_code)]
#[derive(Clone, Copy, Debug, derive_new::new)]
pub struct Rv32VecHeapAdapterAir<
const NUM_READS: usize,
const BLOCKS_PER_READ: usize,
const BLOCKS_PER_WRITE: usize,
const READ_SIZE: usize,
const WRITE_SIZE: usize,
> {
pub(super) execution_bridge: ExecutionBridge,
pub(super) memory_bridge: MemoryBridge,
pub bus: BitwiseOperationLookupBus,
address_bits: usize,
}
impl<
F: Field,
const NUM_READS: usize,
const BLOCKS_PER_READ: usize,
const BLOCKS_PER_WRITE: usize,
const READ_SIZE: usize,
const WRITE_SIZE: usize,
> BaseAir<F>
for Rv32VecHeapAdapterAir<NUM_READS, BLOCKS_PER_READ, BLOCKS_PER_WRITE, READ_SIZE, WRITE_SIZE>
{
fn width(&self) -> usize {
Rv32VecHeapAdapterCols::<
F,
NUM_READS,
BLOCKS_PER_READ,
BLOCKS_PER_WRITE,
READ_SIZE,
WRITE_SIZE,
>::width()
}
}
impl<
AB: InteractionBuilder,
const NUM_READS: usize,
const BLOCKS_PER_READ: usize,
const BLOCKS_PER_WRITE: usize,
const READ_SIZE: usize,
const WRITE_SIZE: usize,
> VmAdapterAir<AB>
for Rv32VecHeapAdapterAir<NUM_READS, BLOCKS_PER_READ, BLOCKS_PER_WRITE, READ_SIZE, WRITE_SIZE>
{
type Interface = VecHeapAdapterInterface<
AB::Expr,
NUM_READS,
BLOCKS_PER_READ,
BLOCKS_PER_WRITE,
READ_SIZE,
WRITE_SIZE,
>;
fn eval(
&self,
builder: &mut AB,
local: &[AB::Var],
ctx: AdapterAirContext<AB::Expr, Self::Interface>,
) {
let cols: &Rv32VecHeapAdapterCols<
_,
NUM_READS,
BLOCKS_PER_READ,
BLOCKS_PER_WRITE,
READ_SIZE,
WRITE_SIZE,
> = local.borrow();
let timestamp = cols.from_state.timestamp;
let mut timestamp_delta: usize = 0;
let mut timestamp_pp = || {
timestamp_delta += 1;
timestamp + AB::F::from_canonical_usize(timestamp_delta - 1)
};
for (ptr, val, aux) in izip!(cols.rs_ptr, cols.rs_val, &cols.rs_read_aux).chain(once((
cols.rd_ptr,
cols.rd_val,
&cols.rd_read_aux,
))) {
self.memory_bridge
.read(
MemoryAddress::new(AB::F::from_canonical_u32(RV32_REGISTER_AS), ptr),
val,
timestamp_pp(),
aux,
)
.eval(builder, ctx.instruction.is_valid.clone());
}
let need_range_check: Vec<AB::Var> = cols
.rs_val
.iter()
.chain(std::iter::repeat(&cols.rd_val).take(2))
.map(|val| val[RV32_REGISTER_NUM_LIMBS - 1])
.collect();
let limb_shift = AB::F::from_canonical_usize(
1 << (RV32_CELL_BITS * RV32_REGISTER_NUM_LIMBS - self.address_bits),
);
for pair in need_range_check.chunks_exact(2) {
self.bus
.send_range(pair[0] * limb_shift, pair[1] * limb_shift)
.eval(builder, ctx.instruction.is_valid.clone());
}
let rd_val_f: AB::Expr = abstract_compose(cols.rd_val);
let rs_val_f: [AB::Expr; NUM_READS] = cols.rs_val.map(abstract_compose);
let e = AB::F::from_canonical_u32(RV32_MEMORY_AS);
for (address, reads, reads_aux) in izip!(rs_val_f, ctx.reads, &cols.reads_aux,) {
for (i, (read, aux)) in zip(reads, reads_aux).enumerate() {
self.memory_bridge
.read(
MemoryAddress::new(
e,
address.clone() + AB::Expr::from_canonical_usize(i * READ_SIZE),
),
read,
timestamp_pp(),
aux,
)
.eval(builder, ctx.instruction.is_valid.clone());
}
}
for (i, (write, aux)) in zip(ctx.writes, &cols.writes_aux).enumerate() {
self.memory_bridge
.write(
MemoryAddress::new(
e,
rd_val_f.clone() + AB::Expr::from_canonical_usize(i * WRITE_SIZE),
),
write,
timestamp_pp(),
aux,
)
.eval(builder, ctx.instruction.is_valid.clone());
}
self.execution_bridge
.execute_and_increment_or_set_pc(
ctx.instruction.opcode,
[
cols.rd_ptr.into(),
cols.rs_ptr
.first()
.map(|&x| x.into())
.unwrap_or(AB::Expr::ZERO),
cols.rs_ptr
.get(1)
.map(|&x| x.into())
.unwrap_or(AB::Expr::ZERO),
AB::Expr::from_canonical_u32(RV32_REGISTER_AS),
e.into(),
],
cols.from_state,
AB::F::from_canonical_usize(timestamp_delta),
(4, ctx.to_pc),
)
.eval(builder, ctx.instruction.is_valid.clone());
}
fn get_from_pc(&self, local: &[AB::Var]) -> AB::Var {
let cols: &Rv32VecHeapAdapterCols<
_,
NUM_READS,
BLOCKS_PER_READ,
BLOCKS_PER_WRITE,
READ_SIZE,
WRITE_SIZE,
> = local.borrow();
cols.from_state.pc
}
}
impl<
F: PrimeField32,
const NUM_READS: usize,
const BLOCKS_PER_READ: usize,
const BLOCKS_PER_WRITE: usize,
const READ_SIZE: usize,
const WRITE_SIZE: usize,
> VmAdapterChip<F>
for Rv32VecHeapAdapterChip<
F,
NUM_READS,
BLOCKS_PER_READ,
BLOCKS_PER_WRITE,
READ_SIZE,
WRITE_SIZE,
>
{
type ReadRecord = Rv32VecHeapReadRecord<F, NUM_READS, BLOCKS_PER_READ, READ_SIZE>;
type WriteRecord = Rv32VecHeapWriteRecord<F, BLOCKS_PER_WRITE, WRITE_SIZE>;
type Air =
Rv32VecHeapAdapterAir<NUM_READS, BLOCKS_PER_READ, BLOCKS_PER_WRITE, READ_SIZE, WRITE_SIZE>;
type Interface = VecHeapAdapterInterface<
F,
NUM_READS,
BLOCKS_PER_READ,
BLOCKS_PER_WRITE,
READ_SIZE,
WRITE_SIZE,
>;
fn preprocess(
&mut self,
memory: &mut MemoryController<F>,
instruction: &Instruction<F>,
) -> Result<(
<Self::Interface as VmAdapterInterface<F>>::Reads,
Self::ReadRecord,
)> {
let Instruction { a, b, c, d, e, .. } = *instruction;
debug_assert_eq!(d.as_canonical_u32(), RV32_REGISTER_AS);
debug_assert_eq!(e.as_canonical_u32(), RV32_MEMORY_AS);
let mut rs_vals = [0; NUM_READS];
let rs_records: [_; NUM_READS] = from_fn(|i| {
let addr = if i == 0 { b } else { c };
let (record, val) = read_rv32_register(memory, d, addr);
rs_vals[i] = val;
record
});
let (rd_record, rd_val) = read_rv32_register(memory, d, a);
let read_records = rs_vals.map(|address| {
assert!(
address as usize + READ_SIZE * BLOCKS_PER_READ - 1 < (1 << self.air.address_bits)
);
from_fn(|i| {
memory.read::<READ_SIZE>(e, F::from_canonical_u32(address + (i * READ_SIZE) as u32))
})
});
let read_data = read_records.map(|r| r.map(|x| x.data));
assert!(rd_val as usize + WRITE_SIZE * BLOCKS_PER_WRITE - 1 < (1 << self.air.address_bits));
let record = Rv32VecHeapReadRecord {
rs: rs_records,
rd: rd_record,
rd_val: F::from_canonical_u32(rd_val),
reads: read_records,
};
Ok((read_data, record))
}
fn postprocess(
&mut self,
memory: &mut MemoryController<F>,
instruction: &Instruction<F>,
from_state: ExecutionState<u32>,
output: AdapterRuntimeContext<F, Self::Interface>,
read_record: &Self::ReadRecord,
) -> Result<(ExecutionState<u32>, Self::WriteRecord)> {
let e = instruction.e;
let mut i = 0;
let writes = output.writes.map(|write| {
let record = memory.write(
e,
read_record.rd_val + F::from_canonical_u32((i * WRITE_SIZE) as u32),
write,
);
i += 1;
record
});
Ok((
ExecutionState {
pc: from_state.pc + 4,
timestamp: memory.timestamp(),
},
Self::WriteRecord { from_state, writes },
))
}
fn generate_trace_row(
&self,
row_slice: &mut [F],
read_record: Self::ReadRecord,
write_record: Self::WriteRecord,
aux_cols_factory: &MemoryAuxColsFactory<F>,
) {
vec_heap_generate_trace_row_impl(
row_slice,
&read_record,
&write_record,
aux_cols_factory,
&self.bitwise_lookup_chip,
self.air.address_bits,
)
}
fn air(&self) -> &Self::Air {
&self.air
}
}
pub(super) fn vec_heap_generate_trace_row_impl<
F: PrimeField32,
const NUM_READS: usize,
const BLOCKS_PER_READ: usize,
const BLOCKS_PER_WRITE: usize,
const READ_SIZE: usize,
const WRITE_SIZE: usize,
>(
row_slice: &mut [F],
read_record: &Rv32VecHeapReadRecord<F, NUM_READS, BLOCKS_PER_READ, READ_SIZE>,
write_record: &Rv32VecHeapWriteRecord<F, BLOCKS_PER_WRITE, WRITE_SIZE>,
aux_cols_factory: &MemoryAuxColsFactory<F>,
bitwise_lookup_chip: &BitwiseOperationLookupChip<RV32_CELL_BITS>,
address_bits: usize,
) {
let row_slice: &mut Rv32VecHeapAdapterCols<
F,
NUM_READS,
BLOCKS_PER_READ,
BLOCKS_PER_WRITE,
READ_SIZE,
WRITE_SIZE,
> = row_slice.borrow_mut();
row_slice.from_state = write_record.from_state.map(F::from_canonical_u32);
row_slice.rd_ptr = read_record.rd.pointer;
row_slice.rs_ptr = read_record.rs.map(|r| r.pointer);
row_slice.rd_val = read_record.rd.data;
row_slice.rs_val = read_record.rs.map(|r| r.data);
row_slice.rs_read_aux = read_record
.rs
.map(|r| aux_cols_factory.make_read_aux_cols(r));
row_slice.rd_read_aux = aux_cols_factory.make_read_aux_cols(read_record.rd);
row_slice.reads_aux = read_record
.reads
.map(|r| r.map(|x| aux_cols_factory.make_read_aux_cols(x)));
row_slice.writes_aux = write_record
.writes
.map(|w| aux_cols_factory.make_write_aux_cols(w));
let need_range_check: Vec<u32> = read_record
.rs
.iter()
.chain(std::iter::repeat(&read_record.rd).take(2))
.map(|record| record.data[RV32_REGISTER_NUM_LIMBS - 1].as_canonical_u32())
.collect();
debug_assert!(address_bits <= RV32_CELL_BITS * RV32_REGISTER_NUM_LIMBS);
let limb_shift_bits = RV32_CELL_BITS * RV32_REGISTER_NUM_LIMBS - address_bits;
for pair in need_range_check.chunks_exact(2) {
bitwise_lookup_chip.request_range(pair[0] << limb_shift_bits, pair[1] << limb_shift_bits);
}
}