openvm_circuit/system/memory/offline_checker/
bridge.rs

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