openvm_circuit/system/memory/offline_checker/
bridge.rs

1use getset::CopyGetters;
2use openvm_circuit_primitives::{
3    assert_less_than::{AssertLessThanIo, AssertLtSubAir},
4    is_zero::{IsZeroIo, IsZeroSubAir},
5    utils::not,
6    var_range::VariableRangeCheckerBus,
7    SubAir,
8};
9use openvm_stark_backend::{
10    interaction::InteractionBuilder, p3_air::AirBuilder, p3_field::FieldAlgebra,
11};
12
13use super::bus::MemoryBus;
14use crate::system::memory::{
15    offline_checker::columns::{
16        MemoryBaseAuxCols, MemoryReadAuxCols, MemoryReadOrImmediateAuxCols, MemoryWriteAuxCols,
17    },
18    MemoryAddress,
19};
20
21/// AUX_LEN is the number of auxiliary columns (aka the number of limbs that the input numbers will
22/// be decomposed into) for the `AssertLtSubAir` in the `MemoryOfflineChecker`.
23/// Warning: This requires that (timestamp_max_bits + decomp - 1) / decomp = AUX_LEN
24///         in MemoryOfflineChecker (or whenever AssertLtSubAir is used)
25pub const AUX_LEN: usize = 2;
26
27/// The [MemoryBridge] is used within AIR evaluation functions to constrain logical memory
28/// operations (read/write). It adds all necessary constraints and interactions.
29#[derive(Clone, Copy, Debug)]
30pub struct MemoryBridge {
31    offline_checker: MemoryOfflineChecker,
32}
33
34impl MemoryBridge {
35    /// Create a new [MemoryBridge] with the provided offline_checker.
36    pub fn new(
37        memory_bus: MemoryBus,
38        timestamp_max_bits: usize,
39        range_bus: VariableRangeCheckerBus,
40    ) -> Self {
41        Self {
42            offline_checker: MemoryOfflineChecker::new(memory_bus, timestamp_max_bits, range_bus),
43        }
44    }
45
46    pub fn memory_bus(&self) -> MemoryBus {
47        self.offline_checker.memory_bus
48    }
49
50    pub fn range_bus(&self) -> VariableRangeCheckerBus {
51        self.offline_checker.timestamp_lt_air.bus
52    }
53
54    /// Prepare a logical memory read operation.
55    #[must_use]
56    pub fn read<'a, T, V, const N: usize>(
57        &self,
58        address: MemoryAddress<impl Into<T>, impl Into<T>>,
59        data: [impl Into<T>; N],
60        timestamp: impl Into<T>,
61        aux: &'a MemoryReadAuxCols<V>,
62    ) -> MemoryReadOperation<'a, T, V, N> {
63        MemoryReadOperation {
64            offline_checker: self.offline_checker,
65            address: MemoryAddress::from(address),
66            data: data.map(Into::into),
67            timestamp: timestamp.into(),
68            aux,
69        }
70    }
71
72    /// Prepare a logical memory read or immediate operation.
73    #[must_use]
74    pub fn read_or_immediate<'a, T, V>(
75        &self,
76        address: MemoryAddress<impl Into<T>, impl Into<T>>,
77        data: impl Into<T>,
78        timestamp: impl Into<T>,
79        aux: &'a MemoryReadOrImmediateAuxCols<V>,
80    ) -> MemoryReadOrImmediateOperation<'a, T, V> {
81        MemoryReadOrImmediateOperation {
82            offline_checker: self.offline_checker,
83            address: MemoryAddress::from(address),
84            data: data.into(),
85            timestamp: timestamp.into(),
86            aux,
87        }
88    }
89
90    /// Prepare a logical memory write operation.
91    #[must_use]
92    pub fn write<'a, T, V, const N: usize>(
93        &self,
94        address: MemoryAddress<impl Into<T>, impl Into<T>>,
95        data: [impl Into<T>; N],
96        timestamp: impl Into<T>,
97        aux: &'a MemoryWriteAuxCols<V, N>,
98    ) -> MemoryWriteOperation<'a, T, V, N> {
99        MemoryWriteOperation {
100            offline_checker: self.offline_checker,
101            address: MemoryAddress::from(address),
102            data: data.map(Into::into),
103            timestamp: timestamp.into(),
104            aux,
105        }
106    }
107}
108
109/// Constraints and interactions for a logical memory read of `(address, data)` at time `timestamp`.
110/// This reads `(address, data, timestamp_prev)` from the memory bus and writes
111/// `(address, data, timestamp)` to the memory bus.
112/// Includes constraints for `timestamp_prev < timestamp`.
113///
114/// The generic `T` type is intended to be `AB::Expr` where `AB` is the [AirBuilder].
115/// The auxiliary columns are not expected to be expressions, so the generic `V` type is intended
116/// to be `AB::Var`.
117pub struct MemoryReadOperation<'a, T, V, const N: usize> {
118    offline_checker: MemoryOfflineChecker,
119    address: MemoryAddress<T, T>,
120    data: [T; N],
121    timestamp: T,
122    aux: &'a MemoryReadAuxCols<V>,
123}
124
125/// The max degree of constraints is:
126/// eval_timestamps: deg(enabled) + max(1, deg(self.timestamp))
127/// eval_bulk_access: refer to private function MemoryOfflineChecker::eval_bulk_access
128impl<F: FieldAlgebra, V: Copy + Into<F>, const N: usize> MemoryReadOperation<'_, F, V, N> {
129    /// Evaluate constraints and send/receive interactions.
130    pub fn eval<AB>(self, builder: &mut AB, enabled: impl Into<AB::Expr>)
131    where
132        AB: InteractionBuilder<Var = V, Expr = F>,
133    {
134        let enabled = enabled.into();
135
136        // NOTE: We do not need to constrain `address_space != 0` since this is done implicitly by
137        // the memory interactions argument together with initial/final memory chips.
138
139        self.offline_checker.eval_timestamps(
140            builder,
141            self.timestamp.clone(),
142            &self.aux.base,
143            enabled.clone(),
144        );
145
146        self.offline_checker.eval_bulk_access(
147            builder,
148            self.address,
149            &self.data,
150            &self.data,
151            self.timestamp.clone(),
152            self.aux.base.prev_timestamp,
153            enabled,
154        );
155    }
156}
157
158/// Constraints and interactions for a logical memory read of `(address, data)` at time `timestamp`,
159/// supporting `address.address_space = 0` for immediates.
160///
161/// If `address.address_space` is non-zero, it behaves like `MemoryReadOperation`. Otherwise,
162/// it constrains the immediate value appropriately.
163///
164/// The generic `T` type is intended to be `AB::Expr` where `AB` is the [AirBuilder].
165/// The auxiliary columns are not expected to be expressions, so the generic `V` type is intended
166/// to be `AB::Var`.
167pub struct MemoryReadOrImmediateOperation<'a, T, V> {
168    offline_checker: MemoryOfflineChecker,
169    address: MemoryAddress<T, T>,
170    data: T,
171    timestamp: T,
172    aux: &'a MemoryReadOrImmediateAuxCols<V>,
173}
174
175/// The max degree of constraints is:
176/// IsZeroSubAir.subair_eval:
177///         deg(enabled) + max(deg(address.address_space) + deg(aux.is_immediate),
178///                           deg(address.address_space) + deg(aux.is_zero_aux))
179/// is_immediate check: deg(aux.is_immediate) + max(deg(data), deg(address.pointer))
180/// eval_timestamps: deg(enabled) + max(1, deg(self.timestamp))
181/// eval_bulk_access: refer to private function MemoryOfflineChecker::eval_bulk_access
182impl<F: FieldAlgebra, V: Copy + Into<F>> MemoryReadOrImmediateOperation<'_, F, V> {
183    /// Evaluate constraints and send/receive interactions.
184    pub fn eval<AB>(self, builder: &mut AB, enabled: impl Into<AB::Expr>)
185    where
186        AB: InteractionBuilder<Var = V, Expr = F>,
187    {
188        let enabled = enabled.into();
189
190        // `is_immediate` should be an indicator for `address_space == 0` (when `enabled`).
191        {
192            let is_zero_io = IsZeroIo::new(
193                self.address.address_space.clone(),
194                self.aux.is_immediate.into(),
195                enabled.clone(),
196            );
197            IsZeroSubAir.eval(builder, (is_zero_io, self.aux.is_zero_aux));
198        }
199        // When `is_immediate`, the data should be the pointer value.
200        builder
201            .when(self.aux.is_immediate)
202            .assert_eq(self.data.clone(), self.address.pointer.clone());
203
204        // Timestamps should be increasing (when enabled).
205        self.offline_checker.eval_timestamps(
206            builder,
207            self.timestamp.clone(),
208            &self.aux.base,
209            enabled.clone(),
210        );
211
212        self.offline_checker.eval_bulk_access(
213            builder,
214            self.address,
215            &[self.data.clone()],
216            &[self.data],
217            self.timestamp,
218            self.aux.base.prev_timestamp,
219            enabled * not(self.aux.is_immediate),
220        );
221    }
222}
223
224/// Constraints and interactions for a logical memory write of `(address, data)` at time
225/// `timestamp`. This reads `(address, data_prev, timestamp_prev)` from the memory bus and writes
226/// `(address, data, timestamp)` to the memory bus.
227/// Includes constraints for `timestamp_prev < timestamp`.
228///
229/// **Note:** This can be used as a logical read operation by setting `data_prev = data`.
230pub struct MemoryWriteOperation<'a, T, V, const N: usize> {
231    offline_checker: MemoryOfflineChecker,
232    address: MemoryAddress<T, T>,
233    data: [T; N],
234    /// The timestamp of the current read
235    timestamp: T,
236    aux: &'a MemoryWriteAuxCols<V, N>,
237}
238
239/// The max degree of constraints is:
240/// eval_timestamps: deg(enabled) + max(1, deg(self.timestamp))
241/// eval_bulk_access: refer to private function MemoryOfflineChecker::eval_bulk_access
242impl<T: FieldAlgebra, V: Copy + Into<T>, const N: usize> MemoryWriteOperation<'_, T, V, N> {
243    /// Evaluate constraints and send/receive interactions. `enabled` must be boolean.
244    pub fn eval<AB>(self, builder: &mut AB, enabled: impl Into<AB::Expr>)
245    where
246        AB: InteractionBuilder<Var = V, Expr = T>,
247    {
248        let enabled = enabled.into();
249        self.offline_checker.eval_timestamps(
250            builder,
251            self.timestamp.clone(),
252            &self.aux.base,
253            enabled.clone(),
254        );
255
256        self.offline_checker.eval_bulk_access(
257            builder,
258            self.address,
259            &self.data,
260            &self.aux.prev_data.map(Into::into),
261            self.timestamp,
262            self.aux.base.prev_timestamp,
263            enabled,
264        );
265    }
266}
267
268#[derive(Clone, Copy, Debug, CopyGetters)]
269struct MemoryOfflineChecker {
270    #[get_copy = "pub"]
271    memory_bus: MemoryBus,
272    #[get_copy = "pub"]
273    timestamp_lt_air: AssertLtSubAir,
274}
275
276impl MemoryOfflineChecker {
277    fn new(
278        memory_bus: MemoryBus,
279        timestamp_max_bits: usize,
280        range_bus: VariableRangeCheckerBus,
281    ) -> Self {
282        Self {
283            memory_bus,
284            timestamp_lt_air: AssertLtSubAir::new(range_bus, timestamp_max_bits),
285        }
286    }
287
288    /// The max degree of constraints is:
289    /// deg(enabled) + max(1, deg(timestamp))
290    /// Note: deg(prev_timestamp) = 1 since prev_timestamp is Var
291    fn eval_timestamps<AB: InteractionBuilder>(
292        &self,
293        builder: &mut AB,
294        timestamp: AB::Expr,
295        base: &MemoryBaseAuxCols<AB::Var>,
296        enabled: AB::Expr,
297    ) {
298        let lt_io = AssertLessThanIo::new(base.prev_timestamp, timestamp.clone(), enabled);
299        self.timestamp_lt_air
300            .eval(builder, (lt_io, &base.timestamp_lt_aux.lower_decomp));
301    }
302
303    /// At the core, eval_bulk_access is a bunch of push_sends and push_receives.
304    /// The max constraint degree of expressions in sends/receives is:
305    /// max(max_deg(data), max_deg(prev_data), max_deg(timestamp), max_deg(prev_timestamps))
306    /// Also, each one of them has count with degree: deg(enabled)
307    #[allow(clippy::too_many_arguments)]
308    fn eval_bulk_access<AB, const N: usize>(
309        &self,
310        builder: &mut AB,
311        address: MemoryAddress<AB::Expr, AB::Expr>,
312        data: &[AB::Expr; N],
313        prev_data: &[AB::Expr; N],
314        timestamp: AB::Expr,
315        prev_timestamp: AB::Var,
316        enabled: AB::Expr,
317    ) where
318        AB: InteractionBuilder,
319    {
320        self.memory_bus
321            .receive(address.clone(), prev_data.to_vec(), prev_timestamp)
322            .eval(builder, enabled.clone());
323
324        self.memory_bus
325            .send(address, data.to_vec(), timestamp)
326            .eval(builder, enabled);
327    }
328}