use std::iter;
use openvm_stark_backend::{
interaction::{BusIndex, InteractionBuilder, PermutationCheckBus},
p3_field::PrimeCharacteristicRing,
};
use crate::system::memory::MemoryAddress;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MemoryBus {
pub inner: PermutationCheckBus,
}
impl MemoryBus {
pub const fn new(index: BusIndex) -> Self {
Self {
inner: PermutationCheckBus::new(index),
}
}
}
impl MemoryBus {
#[inline(always)]
pub fn index(&self) -> BusIndex {
self.inner.index
}
pub fn send<T: Clone>(
&self,
address: MemoryAddress<impl Into<T>, impl Into<T>>,
data: Vec<impl Into<T>>,
timestamp: impl Into<T>,
) -> MemoryBusInteraction<T> {
self.push(true, address, data, timestamp)
}
pub fn receive<T: Clone>(
&self,
address: MemoryAddress<impl Into<T>, impl Into<T>>,
data: Vec<impl Into<T>>,
timestamp: impl Into<T>,
) -> MemoryBusInteraction<T> {
self.push(false, address, data, timestamp)
}
fn push<T: Clone>(
&self,
is_send: bool,
address: MemoryAddress<impl Into<T>, impl Into<T>>,
data: Vec<impl Into<T>>,
timestamp: impl Into<T>,
) -> MemoryBusInteraction<T> {
MemoryBusInteraction {
bus: self.inner,
is_send,
address: MemoryAddress::new(address.address_space.into(), address.pointer.into()),
data: data.into_iter().map(|item| item.into()).collect(),
timestamp: timestamp.into(),
}
}
}
#[derive(Clone, Debug)]
pub struct MemoryBusInteraction<T> {
pub bus: PermutationCheckBus,
pub is_send: bool,
pub address: MemoryAddress<T, T>,
pub data: Vec<T>,
pub timestamp: T,
}
impl<T: PrimeCharacteristicRing> MemoryBusInteraction<T> {
pub fn eval<AB>(self, builder: &mut AB, direction: impl Into<AB::Expr>)
where
AB: InteractionBuilder<Expr = T>,
{
let fields = iter::empty()
.chain(iter::once(self.address.address_space))
.chain(iter::once(self.address.pointer))
.chain(self.data)
.chain(iter::once(self.timestamp));
if self.is_send {
self.bus.interact(builder, fields, direction);
} else {
self.bus
.interact(builder, fields, AB::Expr::NEG_ONE * direction.into());
}
}
}