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            #[allow(clippy::cloned_ref_to_slice_refs)]
216            &[self.data.clone()],
217            &[self.data],
218            self.timestamp,
219            self.aux.base.prev_timestamp,
220            enabled * not(self.aux.is_immediate),
221        );
222    }
223}
224
225/// Constraints and interactions for a logical memory write of `(address, data)` at time
226/// `timestamp`. This reads `(address, data_prev, timestamp_prev)` from the memory bus and writes
227/// `(address, data, timestamp)` to the memory bus.
228/// Includes constraints for `timestamp_prev < timestamp`.
229///
230/// **Note:** This can be used as a logical read operation by setting `data_prev = data`.
231pub struct MemoryWriteOperation<'a, T, V, const N: usize> {
232    offline_checker: MemoryOfflineChecker,
233    address: MemoryAddress<T, T>,
234    data: [T; N],
235    /// The timestamp of the current read
236    timestamp: T,
237    aux: &'a MemoryWriteAuxCols<V, N>,
238}
239
240/// The max degree of constraints is:
241/// eval_timestamps: deg(enabled) + max(1, deg(self.timestamp))
242/// eval_bulk_access: refer to private function MemoryOfflineChecker::eval_bulk_access
243impl<T: FieldAlgebra, V: Copy + Into<T>, const N: usize> MemoryWriteOperation<'_, T, V, N> {
244    /// Evaluate constraints and send/receive interactions. `enabled` must be boolean.
245    pub fn eval<AB>(self, builder: &mut AB, enabled: impl Into<AB::Expr>)
246    where
247        AB: InteractionBuilder<Var = V, Expr = T>,
248    {
249        let enabled = enabled.into();
250        self.offline_checker.eval_timestamps(
251            builder,
252            self.timestamp.clone(),
253            &self.aux.base,
254            enabled.clone(),
255        );
256
257        self.offline_checker.eval_bulk_access(
258            builder,
259            self.address,
260            &self.data,
261            &self.aux.prev_data.map(Into::into),
262            self.timestamp,
263            self.aux.base.prev_timestamp,
264            enabled,
265        );
266    }
267}
268
269#[derive(Clone, Copy, Debug, CopyGetters)]
270struct MemoryOfflineChecker {
271    #[get_copy = "pub"]
272    memory_bus: MemoryBus,
273    #[get_copy = "pub"]
274    timestamp_lt_air: AssertLtSubAir,
275}
276
277impl MemoryOfflineChecker {
278    fn new(
279        memory_bus: MemoryBus,
280        timestamp_max_bits: usize,
281        range_bus: VariableRangeCheckerBus,
282    ) -> Self {
283        Self {
284            memory_bus,
285            timestamp_lt_air: AssertLtSubAir::new(range_bus, timestamp_max_bits),
286        }
287    }
288
289    /// The max degree of constraints is:
290    /// deg(enabled) + max(1, deg(timestamp))
291    /// Note: deg(prev_timestamp) = 1 since prev_timestamp is Var
292    fn eval_timestamps<AB: InteractionBuilder>(
293        &self,
294        builder: &mut AB,
295        timestamp: AB::Expr,
296        base: &MemoryBaseAuxCols<AB::Var>,
297        enabled: AB::Expr,
298    ) {
299        let lt_io = AssertLessThanIo::new(base.prev_timestamp, timestamp.clone(), enabled);
300        self.timestamp_lt_air
301            .eval(builder, (lt_io, &base.timestamp_lt_aux.lower_decomp));
302    }
303
304    /// At the core, eval_bulk_access is a bunch of push_sends and push_receives.
305    /// The max constraint degree of expressions in sends/receives is:
306    /// max(max_deg(data), max_deg(prev_data), max_deg(timestamp), max_deg(prev_timestamps))
307    /// Also, each one of them has count with degree: deg(enabled)
308    #[allow(clippy::too_many_arguments)]
309    fn eval_bulk_access<AB, const N: usize>(
310        &self,
311        builder: &mut AB,
312        address: MemoryAddress<AB::Expr, AB::Expr>,
313        data: &[AB::Expr; N],
314        prev_data: &[AB::Expr; N],
315        timestamp: AB::Expr,
316        prev_timestamp: AB::Var,
317        enabled: AB::Expr,
318    ) where
319        AB: InteractionBuilder,
320    {
321        self.memory_bus
322            .receive(address.clone(), prev_data.to_vec(), prev_timestamp)
323            .eval(builder, enabled.clone());
324
325        self.memory_bus
326            .send(address, data.to_vec(), timestamp)
327            .eval(builder, enabled);
328    }
329}