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::PrimeCharacteristicRing,
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: PrimeCharacteristicRing, V: Copy + Into<F>, const N: usize>
129    MemoryReadOperation<'_, F, V, N>
130{
131    /// Evaluate constraints and send/receive interactions.
132    pub fn eval<AB>(self, builder: &mut AB, enabled: impl Into<AB::Expr>)
133    where
134        AB: InteractionBuilder<Var = V, Expr = F>,
135    {
136        let enabled = enabled.into();
137
138        // NOTE: We do not need to constrain `address_space != 0` since this is done implicitly by
139        // the memory interactions argument together with initial/final memory chips.
140
141        self.offline_checker.eval_timestamps(
142            builder,
143            self.timestamp.clone(),
144            &self.aux.base,
145            enabled.clone(),
146        );
147
148        self.offline_checker.eval_bulk_access(
149            builder,
150            self.address,
151            &self.data,
152            &self.data,
153            self.timestamp.clone(),
154            self.aux.base.prev_timestamp,
155            enabled,
156        );
157    }
158}
159
160/// Constraints and interactions for a logical memory read of `(address, data)` at time `timestamp`,
161/// supporting `address.address_space = 0` for immediates.
162///
163/// If `address.address_space` is non-zero, it behaves like `MemoryReadOperation`. Otherwise,
164/// it constrains the immediate value appropriately.
165///
166/// The generic `T` type is intended to be `AB::Expr` where `AB` is the [AirBuilder].
167/// The auxiliary columns are not expected to be expressions, so the generic `V` type is intended
168/// to be `AB::Var`.
169pub struct MemoryReadOrImmediateOperation<'a, T, V> {
170    offline_checker: MemoryOfflineChecker,
171    address: MemoryAddress<T, T>,
172    data: T,
173    timestamp: T,
174    aux: &'a MemoryReadOrImmediateAuxCols<V>,
175}
176
177/// The max degree of constraints is:
178/// IsZeroSubAir.subair_eval:
179///         deg(enabled) + max(deg(address.address_space) + deg(aux.is_immediate),
180///                           deg(address.address_space) + deg(aux.is_zero_aux))
181/// is_immediate check: deg(aux.is_immediate) + max(deg(data), deg(address.pointer))
182/// eval_timestamps: deg(enabled) + max(1, deg(self.timestamp))
183/// eval_bulk_access: refer to private function MemoryOfflineChecker::eval_bulk_access
184impl<F: PrimeCharacteristicRing, V: Copy + Into<F>> MemoryReadOrImmediateOperation<'_, F, V> {
185    /// Evaluate constraints and send/receive interactions.
186    pub fn eval<AB>(self, builder: &mut AB, enabled: impl Into<AB::Expr>)
187    where
188        AB: InteractionBuilder<Var = V, Expr = F>,
189    {
190        let enabled = enabled.into();
191
192        // `is_immediate` should be an indicator for `address_space == 0` (when `enabled`).
193        {
194            let is_zero_io = IsZeroIo::new(
195                self.address.address_space.clone(),
196                self.aux.is_immediate.into(),
197                enabled.clone(),
198            );
199            IsZeroSubAir.eval(builder, (is_zero_io, self.aux.is_zero_aux));
200        }
201        // When `is_immediate`, the data should be the pointer value.
202        builder
203            .when(self.aux.is_immediate)
204            .assert_eq(self.data.clone(), self.address.pointer.clone());
205
206        // Timestamps should be increasing (when enabled).
207        self.offline_checker.eval_timestamps(
208            builder,
209            self.timestamp.clone(),
210            &self.aux.base,
211            enabled.clone(),
212        );
213
214        #[allow(clippy::cloned_ref_to_slice_refs)]
215        self.offline_checker.eval_bulk_access(
216            builder,
217            self.address,
218            #[allow(clippy::cloned_ref_to_slice_refs)]
219            &[self.data.clone()],
220            &[self.data],
221            self.timestamp,
222            self.aux.base.prev_timestamp,
223            enabled * not(self.aux.is_immediate),
224        );
225    }
226}
227
228/// Constraints and interactions for a logical memory write of `(address, data)` at time
229/// `timestamp`. This reads `(address, data_prev, timestamp_prev)` from the memory bus and writes
230/// `(address, data, timestamp)` to the memory bus.
231/// Includes constraints for `timestamp_prev < timestamp`.
232///
233/// **Note:** This can be used as a logical read operation by setting `data_prev = data`.
234pub struct MemoryWriteOperation<'a, T, V, const N: usize> {
235    offline_checker: MemoryOfflineChecker,
236    address: MemoryAddress<T, T>,
237    data: [T; N],
238    /// The timestamp of the current read
239    timestamp: T,
240    aux: &'a MemoryWriteAuxCols<V, N>,
241}
242
243/// The max degree of constraints is:
244/// eval_timestamps: deg(enabled) + max(1, deg(self.timestamp))
245/// eval_bulk_access: refer to private function MemoryOfflineChecker::eval_bulk_access
246impl<T: PrimeCharacteristicRing, V: Copy + Into<T>, const N: usize>
247    MemoryWriteOperation<'_, T, V, N>
248{
249    /// Evaluate constraints and send/receive interactions. `enabled` must be boolean.
250    pub fn eval<AB>(self, builder: &mut AB, enabled: impl Into<AB::Expr>)
251    where
252        AB: InteractionBuilder<Var = V, Expr = T>,
253    {
254        let enabled = enabled.into();
255        self.offline_checker.eval_timestamps(
256            builder,
257            self.timestamp.clone(),
258            &self.aux.base,
259            enabled.clone(),
260        );
261
262        self.offline_checker.eval_bulk_access(
263            builder,
264            self.address,
265            &self.data,
266            &self.aux.prev_data.map(Into::into),
267            self.timestamp,
268            self.aux.base.prev_timestamp,
269            enabled,
270        );
271    }
272}
273
274#[derive(Clone, Copy, Debug, CopyGetters)]
275struct MemoryOfflineChecker {
276    #[get_copy = "pub"]
277    memory_bus: MemoryBus,
278    #[get_copy = "pub"]
279    timestamp_lt_air: AssertLtSubAir,
280}
281
282impl MemoryOfflineChecker {
283    fn new(
284        memory_bus: MemoryBus,
285        timestamp_max_bits: usize,
286        range_bus: VariableRangeCheckerBus,
287    ) -> Self {
288        Self {
289            memory_bus,
290            timestamp_lt_air: AssertLtSubAir::new(range_bus, timestamp_max_bits),
291        }
292    }
293
294    /// The max degree of constraints is:
295    /// deg(enabled) + max(1, deg(timestamp))
296    /// Note: deg(prev_timestamp) = 1 since prev_timestamp is Var
297    fn eval_timestamps<AB: InteractionBuilder>(
298        &self,
299        builder: &mut AB,
300        timestamp: AB::Expr,
301        base: &MemoryBaseAuxCols<AB::Var>,
302        enabled: AB::Expr,
303    ) {
304        let lt_io = AssertLessThanIo::new(base.prev_timestamp, timestamp.clone(), enabled);
305        self.timestamp_lt_air
306            .eval(builder, (lt_io, &base.timestamp_lt_aux.lower_decomp));
307    }
308
309    /// At the core, eval_bulk_access is a bunch of push_sends and push_receives.
310    /// The max constraint degree of expressions in sends/receives is:
311    /// max(max_deg(data), max_deg(prev_data), max_deg(timestamp), max_deg(prev_timestamps))
312    /// Also, each one of them has count with degree: deg(enabled)
313    #[allow(clippy::too_many_arguments)]
314    fn eval_bulk_access<AB, const N: usize>(
315        &self,
316        builder: &mut AB,
317        address: MemoryAddress<AB::Expr, AB::Expr>,
318        data: &[AB::Expr; N],
319        prev_data: &[AB::Expr; N],
320        timestamp: AB::Expr,
321        prev_timestamp: AB::Var,
322        enabled: AB::Expr,
323    ) where
324        AB: InteractionBuilder,
325    {
326        self.memory_bus
327            .receive(address.clone(), prev_data.to_vec(), prev_timestamp)
328            .eval(builder, enabled.clone());
329
330        self.memory_bus
331            .send(address, data.to_vec(), timestamp)
332            .eval(builder, enabled);
333    }
334}