openvm_circuit/system/memory/controller/
mod.rs

Help
1use std::{
2    array,
3    collections::BTreeMap,
4    iter,
5    marker::PhantomData,
6    mem,
7    sync::{Arc, Mutex},
8};
9
10use getset::{Getters, MutGetters};
11use openvm_circuit_primitives::{
12    assert_less_than::{AssertLtSubAir, LessThanAuxCols},
13    is_zero::IsZeroSubAir,
14    utils::next_power_of_two_or_zero,
15    var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus},
16    TraceSubRowGenerator,
17};
18use openvm_stark_backend::{
19    config::{Domain, StarkGenericConfig},
20    interaction::PermutationCheckBus,
21    p3_commit::PolynomialSpace,
22    p3_field::PrimeField32,
23    p3_maybe_rayon::prelude::{IntoParallelIterator, ParallelIterator},
24    p3_util::{log2_ceil_usize, log2_strict_usize},
25    prover::types::AirProofInput,
26    AirRef, Chip, ChipUsageGetter,
27};
28use serde::{Deserialize, Serialize};
29
30use self::interface::MemoryInterface;
31use super::{
32    paged_vec::{AddressMap, PAGE_SIZE},
33    volatile::VolatileBoundaryChip,
34};
35use crate::{
36    arch::{hasher::HasherChip, MemoryConfig},
37    system::memory::{
38        adapter::AccessAdapterInventory,
39        dimensions::MemoryDimensions,
40        merkle::{MemoryMerkleChip, SerialReceiver},
41        offline::{MemoryRecord, OfflineMemory, INITIAL_TIMESTAMP},
42        offline_checker::{
43            MemoryBaseAuxCols, MemoryBridge, MemoryBus, MemoryReadAuxCols,
44            MemoryReadOrImmediateAuxCols, MemoryWriteAuxCols, AUX_LEN,
45        },
46        online::{Memory, MemoryLogEntry},
47        persistent::PersistentBoundaryChip,
48        tree::MemoryNode,
49    },
50};
51
52pub mod dimensions;
53pub mod interface;
54
55pub const CHUNK: usize = 8;
56/// The offset of the Merkle AIR in AIRs of MemoryController.
57pub const MERKLE_AIR_OFFSET: usize = 1;
58/// The offset of the boundary AIR in AIRs of MemoryController.
59pub const BOUNDARY_AIR_OFFSET: usize = 0;
60
61#[repr(C)]
62#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
63pub struct RecordId(pub usize);
64
65pub type MemoryImage<F> = AddressMap<F, PAGE_SIZE>;
66
67#[repr(C)]
68#[derive(Clone, Copy, Debug, PartialEq, Eq)]
69pub struct TimestampedValues<T, const N: usize> {
70    pub timestamp: u32,
71    pub values: [T; N],
72}
73
74/// An equipartition of memory, with timestamps and values.
75///
76/// The key is a pair `(address_space, label)`, where `label` is the index of the block in the
77/// partition. I.e., the starting address of the block is `(address_space, label * N)`.
78///
79/// If a key is not present in the map, then the block is uninitialized (and therefore zero).
80pub type TimestampedEquipartition<F, const N: usize> =
81    BTreeMap<(u32, u32), TimestampedValues<F, N>>;
82
83/// An equipartition of memory values.
84///
85/// The key is a pair `(address_space, label)`, where `label` is the index of the block in the
86/// partition. I.e., the starting address of the block is `(address_space, label * N)`.
87///
88/// If a key is not present in the map, then the block is uninitialized (and therefore zero).
89pub type Equipartition<F, const N: usize> = BTreeMap<(u32, u32), [F; N]>;
90
91#[derive(Getters, MutGetters)]
92pub struct MemoryController<F> {
93    pub memory_bus: MemoryBus,
94    pub interface_chip: MemoryInterface<F>,
95    #[getset(get = "pub")]
96    pub(crate) mem_config: MemoryConfig,
97    pub range_checker: SharedVariableRangeCheckerChip,
98    // Store separately to avoid smart pointer reference each time
99    range_checker_bus: VariableRangeCheckerBus,
100    // addr_space -> Memory data structure
101    memory: Memory<F>,
102    /// A reference to the `OfflineMemory`. Will be populated after `finalize()`.
103    offline_memory: Arc<Mutex<OfflineMemory<F>>>,
104    pub access_adapters: AccessAdapterInventory<F>,
105    // Filled during finalization.
106    final_state: Option<FinalState<F>>,
107}
108
109#[allow(clippy::large_enum_variant)]
110#[derive(Debug)]
111enum FinalState<F> {
112    Volatile(VolatileFinalState<F>),
113    #[allow(dead_code)]
114    Persistent(PersistentFinalState<F>),
115}
116#[derive(Debug, Default)]
117struct VolatileFinalState<F> {
118    _marker: PhantomData<F>,
119}
120#[allow(dead_code)]
121#[derive(Debug)]
122struct PersistentFinalState<F> {
123    final_memory: Equipartition<F, CHUNK>,
124}
125
126#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
127pub enum MemoryTraceHeights {
128    Volatile(VolatileMemoryTraceHeights),
129    Persistent(PersistentMemoryTraceHeights),
130}
131
132impl MemoryTraceHeights {
133    fn flatten(&self) -> Vec<usize> {
134        match self {
135            MemoryTraceHeights::Volatile(oh) => oh.flatten(),
136            MemoryTraceHeights::Persistent(oh) => oh.flatten(),
137        }
138    }
139
140    /// Round all trace heights to the next power of two. This will round trace heights of 0 to 1.
141    pub fn round_to_next_power_of_two(&mut self) {
142        match self {
143            MemoryTraceHeights::Volatile(oh) => oh.round_to_next_power_of_two(),
144            MemoryTraceHeights::Persistent(oh) => oh.round_to_next_power_of_two(),
145        }
146    }
147
148    /// Round all trace heights to the next power of two, except 0 stays 0.
149    pub fn round_to_next_power_of_two_or_zero(&mut self) {
150        match self {
151            MemoryTraceHeights::Volatile(oh) => oh.round_to_next_power_of_two_or_zero(),
152            MemoryTraceHeights::Persistent(oh) => oh.round_to_next_power_of_two_or_zero(),
153        }
154    }
155}
156
157#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
158pub struct VolatileMemoryTraceHeights {
159    pub boundary: usize,
160    pub access_adapters: Vec<usize>,
161}
162
163impl VolatileMemoryTraceHeights {
164    pub fn flatten(&self) -> Vec<usize> {
165        iter::once(self.boundary)
166            .chain(self.access_adapters.iter().copied())
167            .collect()
168    }
169
170    fn round_to_next_power_of_two(&mut self) {
171        self.boundary = self.boundary.next_power_of_two();
172        self.access_adapters
173            .iter_mut()
174            .for_each(|v| *v = v.next_power_of_two());
175    }
176
177    fn round_to_next_power_of_two_or_zero(&mut self) {
178        self.boundary = next_power_of_two_or_zero(self.boundary);
179        self.access_adapters
180            .iter_mut()
181            .for_each(|v| *v = next_power_of_two_or_zero(*v));
182    }
183}
184
185#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
186pub struct PersistentMemoryTraceHeights {
187    boundary: usize,
188    merkle: usize,
189    access_adapters: Vec<usize>,
190}
191impl PersistentMemoryTraceHeights {
192    pub fn flatten(&self) -> Vec<usize> {
193        vec![self.boundary, self.merkle]
194            .into_iter()
195            .chain(self.access_adapters.iter().copied())
196            .collect()
197    }
198
199    fn round_to_next_power_of_two(&mut self) {
200        self.boundary = self.boundary.next_power_of_two();
201        self.merkle = self.merkle.next_power_of_two();
202        self.access_adapters
203            .iter_mut()
204            .for_each(|v| *v = v.next_power_of_two());
205    }
206
207    fn round_to_next_power_of_two_or_zero(&mut self) {
208        self.boundary = next_power_of_two_or_zero(self.boundary);
209        self.merkle = next_power_of_two_or_zero(self.merkle);
210        self.access_adapters
211            .iter_mut()
212            .for_each(|v| *v = next_power_of_two_or_zero(*v));
213    }
214}
215
216impl<F: PrimeField32> MemoryController<F> {
217    pub fn continuation_enabled(&self) -> bool {
218        match &self.interface_chip {
219            MemoryInterface::Volatile { .. } => false,
220            MemoryInterface::Persistent { .. } => true,
221        }
222    }
223    pub fn with_volatile_memory(
224        memory_bus: MemoryBus,
225        mem_config: MemoryConfig,
226        range_checker: SharedVariableRangeCheckerChip,
227    ) -> Self {
228        let range_checker_bus = range_checker.bus();
229        let initial_memory = AddressMap::from_mem_config(&mem_config);
230        assert!(mem_config.pointer_max_bits <= F::bits() - 2);
231        assert!(mem_config.as_height < F::bits() - 2);
232        let addr_space_max_bits = log2_ceil_usize(
233            (mem_config.as_offset + 2u32.pow(mem_config.as_height as u32)) as usize,
234        );
235        Self {
236            memory_bus,
237            mem_config,
238            interface_chip: MemoryInterface::Volatile {
239                boundary_chip: VolatileBoundaryChip::new(
240                    memory_bus,
241                    addr_space_max_bits,
242                    mem_config.pointer_max_bits,
243                    range_checker.clone(),
244                ),
245            },
246            memory: Memory::new(&mem_config),
247            offline_memory: Arc::new(Mutex::new(OfflineMemory::new(
248                initial_memory,
249                1,
250                memory_bus,
251                range_checker.clone(),
252                mem_config,
253            ))),
254            access_adapters: AccessAdapterInventory::new(
255                range_checker.clone(),
256                memory_bus,
257                mem_config.clk_max_bits,
258                mem_config.max_access_adapter_n,
259            ),
260            range_checker,
261            range_checker_bus,
262            final_state: None,
263        }
264    }
265
266    /// Creates a new memory controller for persistent memory.
267    ///
268    /// Call `set_initial_memory` to set the initial memory state after construction.
269    pub fn with_persistent_memory(
270        memory_bus: MemoryBus,
271        mem_config: MemoryConfig,
272        range_checker: SharedVariableRangeCheckerChip,
273        merkle_bus: PermutationCheckBus,
274        compression_bus: PermutationCheckBus,
275    ) -> Self {
276        assert_eq!(mem_config.as_offset, 1);
277        let memory_dims = MemoryDimensions {
278            as_height: mem_config.as_height,
279            address_height: mem_config.pointer_max_bits - log2_strict_usize(CHUNK),
280            as_offset: 1,
281        };
282        let range_checker_bus = range_checker.bus();
283        let interface_chip = MemoryInterface::Persistent {
284            boundary_chip: PersistentBoundaryChip::new(
285                memory_dims,
286                memory_bus,
287                merkle_bus,
288                compression_bus,
289            ),
290            merkle_chip: MemoryMerkleChip::new(memory_dims, merkle_bus, compression_bus),
291            initial_memory: AddressMap::from_mem_config(&mem_config),
292        };
293        Self {
294            memory_bus,
295            mem_config,
296            interface_chip,
297            memory: Memory::new(&mem_config), // it is expected that the memory will be set later
298            offline_memory: Arc::new(Mutex::new(OfflineMemory::new(
299                AddressMap::from_mem_config(&mem_config),
300                CHUNK,
301                memory_bus,
302                range_checker.clone(),
303                mem_config,
304            ))),
305            access_adapters: AccessAdapterInventory::new(
306                range_checker.clone(),
307                memory_bus,
308                mem_config.clk_max_bits,
309                mem_config.max_access_adapter_n,
310            ),
311            range_checker,
312            range_checker_bus,
313            final_state: None,
314        }
315    }
316
317    pub fn memory_image(&self) -> &MemoryImage<F> {
318        &self.memory.data
319    }
320
321    pub fn set_override_trace_heights(&mut self, overridden_heights: MemoryTraceHeights) {
322        match &mut self.interface_chip {
323            MemoryInterface::Volatile { boundary_chip } => match overridden_heights {
324                MemoryTraceHeights::Volatile(oh) => {
325                    boundary_chip.set_overridden_height(oh.boundary);
326                    self.access_adapters
327                        .set_override_trace_heights(oh.access_adapters);
328                }
329                _ => panic!("Expect overridden_heights to be MemoryTraceHeights::Volatile"),
330            },
331            MemoryInterface::Persistent {
332                boundary_chip,
333                merkle_chip,
334                ..
335            } => match overridden_heights {
336                MemoryTraceHeights::Persistent(oh) => {
337                    boundary_chip.set_overridden_height(oh.boundary);
338                    merkle_chip.set_overridden_height(oh.merkle);
339                    self.access_adapters
340                        .set_override_trace_heights(oh.access_adapters);
341                }
342                _ => panic!("Expect overridden_heights to be MemoryTraceHeights::Persistent"),
343            },
344        }
345    }
346
347    pub fn set_initial_memory(&mut self, memory: MemoryImage<F>) {
348        if self.timestamp() > INITIAL_TIMESTAMP + 1 {
349            panic!("Cannot set initial memory after first timestamp");
350        }
351        let mut offline_memory = self.offline_memory.lock().unwrap();
352        offline_memory.set_initial_memory(memory.clone(), self.mem_config);
353
354        self.memory = Memory::from_image(memory.clone(), self.mem_config.access_capacity);
355
356        match &mut self.interface_chip {
357            MemoryInterface::Volatile { .. } => {
358                assert!(
359                    memory.is_empty(),
360                    "Cannot set initial memory for volatile memory"
361                );
362            }
363            MemoryInterface::Persistent { initial_memory, .. } => {
364                *initial_memory = memory;
365            }
366        }
367    }
368
369    pub fn memory_bridge(&self) -> MemoryBridge {
370        MemoryBridge::new(
371            self.memory_bus,
372            self.mem_config.clk_max_bits,
373            self.range_checker_bus,
374        )
375    }
376
377    pub fn read_cell(&mut self, address_space: F, pointer: F) -> (RecordId, F) {
378        let (record_id, [data]) = self.read(address_space, pointer);
379        (record_id, data)
380    }
381
382    pub fn read<const N: usize>(&mut self, address_space: F, pointer: F) -> (RecordId, [F; N]) {
383        let address_space_u32 = address_space.as_canonical_u32();
384        let ptr_u32 = pointer.as_canonical_u32();
385        assert!(
386            address_space == F::ZERO || ptr_u32 < (1 << self.mem_config.pointer_max_bits),
387            "memory out of bounds: {ptr_u32:?}",
388        );
389
390        let (record_id, values) = self.memory.read::<N>(address_space_u32, ptr_u32);
391
392        (record_id, values)
393    }
394
395    /// Reads a word directly from memory without updating internal state.
396    ///
397    /// Any value returned is unconstrained.
398    pub fn unsafe_read_cell(&self, addr_space: F, ptr: F) -> F {
399        self.unsafe_read::<1>(addr_space, ptr)[0]
400    }
401
402    /// Reads a word directly from memory without updating internal state.
403    ///
404    /// Any value returned is unconstrained.
405    pub fn unsafe_read<const N: usize>(&self, addr_space: F, ptr: F) -> [F; N] {
406        let addr_space = addr_space.as_canonical_u32();
407        let ptr = ptr.as_canonical_u32();
408        array::from_fn(|i| self.memory.get(addr_space, ptr + i as u32))
409    }
410
411    /// Writes `data` to the given cell.
412    ///
413    /// Returns the `RecordId` and previous data.
414    pub fn write_cell(&mut self, address_space: F, pointer: F, data: F) -> (RecordId, F) {
415        let (record_id, [data]) = self.write(address_space, pointer, [data]);
416        (record_id, data)
417    }
418
419    pub fn write<const N: usize>(
420        &mut self,
421        address_space: F,
422        pointer: F,
423        data: [F; N],
424    ) -> (RecordId, [F; N]) {
425        assert_ne!(address_space, F::ZERO);
426        let address_space_u32 = address_space.as_canonical_u32();
427        let ptr_u32 = pointer.as_canonical_u32();
428        assert!(
429            ptr_u32 < (1 << self.mem_config.pointer_max_bits),
430            "memory out of bounds: {ptr_u32:?}",
431        );
432
433        self.memory.write(address_space_u32, ptr_u32, data)
434    }
435
436    pub fn aux_cols_factory(&self) -> MemoryAuxColsFactory<F> {
437        let range_bus = self.range_checker.bus();
438        MemoryAuxColsFactory {
439            range_checker: self.range_checker.clone(),
440            timestamp_lt_air: AssertLtSubAir::new(range_bus, self.mem_config.clk_max_bits),
441            _marker: Default::default(),
442        }
443    }
444
445    pub fn increment_timestamp(&mut self) {
446        self.memory.increment_timestamp_by(1);
447    }
448
449    pub fn increment_timestamp_by(&mut self, change: u32) {
450        self.memory.increment_timestamp_by(change);
451    }
452
453    pub fn timestamp(&self) -> u32 {
454        self.memory.timestamp()
455    }
456
457    fn replay_access_log(&mut self) {
458        let log = mem::take(&mut self.memory.log);
459        if log.is_empty() {
460            // Online memory logs may be empty, but offline memory may be replayed from external
461            // sources. In these cases, we skip the calls to replay access logs because
462            // `set_log_capacity` would panic.
463            tracing::debug!("skipping replay_access_log");
464            return;
465        }
466
467        let mut offline_memory = self.offline_memory.lock().unwrap();
468        offline_memory.set_log_capacity(log.len());
469
470        for entry in log {
471            Self::replay_access(
472                entry,
473                &mut offline_memory,
474                &mut self.interface_chip,
475                &mut self.access_adapters,
476            );
477        }
478    }
479
480    /// Low-level API to replay a single memory access log entry and populate the [OfflineMemory],
481    /// [MemoryInterface], and `AccessAdapterInventory`.
482    pub fn replay_access(
483        entry: MemoryLogEntry<F>,
484        offline_memory: &mut OfflineMemory<F>,
485        interface_chip: &mut MemoryInterface<F>,
486        adapter_records: &mut AccessAdapterInventory<F>,
487    ) {
488        match entry {
489            MemoryLogEntry::Read {
490                address_space,
491                pointer,
492                len,
493            } => {
494                if address_space != 0 {
495                    interface_chip.touch_range(address_space, pointer, len as u32);
496                }
497                offline_memory.read(address_space, pointer, len, adapter_records);
498            }
499            MemoryLogEntry::Write {
500                address_space,
501                pointer,
502                data,
503            } => {
504                if address_space != 0 {
505                    interface_chip.touch_range(address_space, pointer, data.len() as u32);
506                }
507                offline_memory.write(address_space, pointer, data, adapter_records);
508            }
509            MemoryLogEntry::IncrementTimestampBy(amount) => {
510                offline_memory.increment_timestamp_by(amount);
511            }
512        };
513    }
514
515    /// Returns the final memory state if persistent.
516    pub fn finalize<H>(&mut self, hasher: Option<&mut H>)
517    where
518        H: HasherChip<CHUNK, F> + Sync + for<'a> SerialReceiver<&'a [F]>,
519    {
520        if self.final_state.is_some() {
521            return;
522        }
523
524        self.replay_access_log();
525        let mut offline_memory = self.offline_memory.lock().unwrap();
526
527        match &mut self.interface_chip {
528            MemoryInterface::Volatile { boundary_chip } => {
529                let final_memory = offline_memory.finalize::<1>(&mut self.access_adapters);
530                boundary_chip.finalize(final_memory);
531                self.final_state = Some(FinalState::Volatile(VolatileFinalState::default()));
532            }
533            MemoryInterface::Persistent {
534                merkle_chip,
535                boundary_chip,
536                initial_memory,
537            } => {
538                let hasher = hasher.unwrap();
539                let final_partition = offline_memory.finalize::<CHUNK>(&mut self.access_adapters);
540
541                boundary_chip.finalize(initial_memory, &final_partition, hasher);
542                let final_memory_values = final_partition
543                    .into_par_iter()
544                    .map(|(key, value)| (key, value.values))
545                    .collect();
546                let initial_node = MemoryNode::tree_from_memory(
547                    merkle_chip.air.memory_dimensions,
548                    initial_memory,
549                    hasher,
550                );
551                merkle_chip.finalize(&initial_node, &final_memory_values, hasher);
552                self.final_state = Some(FinalState::Persistent(PersistentFinalState {
553                    final_memory: final_memory_values.clone(),
554                }));
555            }
556        };
557    }
558
559    pub fn generate_air_proof_inputs<SC: StarkGenericConfig>(self) -> Vec<AirProofInput<SC>>
560    where
561        Domain<SC>: PolynomialSpace<Val = F>,
562    {
563        let mut ret = Vec::new();
564
565        let Self {
566            interface_chip,
567            access_adapters,
568            ..
569        } = self;
570        match interface_chip {
571            MemoryInterface::Volatile { boundary_chip } => {
572                ret.push(boundary_chip.generate_air_proof_input());
573            }
574            MemoryInterface::Persistent {
575                merkle_chip,
576                boundary_chip,
577                ..
578            } => {
579                debug_assert_eq!(ret.len(), BOUNDARY_AIR_OFFSET);
580                ret.push(boundary_chip.generate_air_proof_input());
581                debug_assert_eq!(ret.len(), MERKLE_AIR_OFFSET);
582                ret.push(merkle_chip.generate_air_proof_input());
583            }
584        }
585        ret.extend(access_adapters.generate_air_proof_inputs());
586        ret
587    }
588
589    pub fn airs<SC: StarkGenericConfig>(&self) -> Vec<AirRef<SC>>
590    where
591        Domain<SC>: PolynomialSpace<Val = F>,
592    {
593        let mut airs = Vec::<AirRef<SC>>::new();
594
595        match &self.interface_chip {
596            MemoryInterface::Volatile { boundary_chip } => {
597                debug_assert_eq!(airs.len(), BOUNDARY_AIR_OFFSET);
598                airs.push(boundary_chip.air())
599            }
600            MemoryInterface::Persistent {
601                boundary_chip,
602                merkle_chip,
603                ..
604            } => {
605                debug_assert_eq!(airs.len(), BOUNDARY_AIR_OFFSET);
606                airs.push(boundary_chip.air());
607                debug_assert_eq!(airs.len(), MERKLE_AIR_OFFSET);
608                airs.push(merkle_chip.air());
609            }
610        }
611        airs.extend(self.access_adapters.airs());
612
613        airs
614    }
615
616    /// Return the number of AIRs in the memory controller.
617    pub fn num_airs(&self) -> usize {
618        let mut num_airs = 1;
619        if self.continuation_enabled() {
620            num_airs += 1;
621        }
622        num_airs += self.access_adapters.num_access_adapters();
623        num_airs
624    }
625
626    pub fn air_names(&self) -> Vec<String> {
627        let mut air_names = vec!["Boundary".to_string()];
628        if self.continuation_enabled() {
629            air_names.push("Merkle".to_string());
630        }
631        air_names.extend(self.access_adapters.air_names());
632        air_names
633    }
634
635    pub fn current_trace_heights(&self) -> Vec<usize> {
636        self.get_memory_trace_heights().flatten()
637    }
638
639    pub fn get_memory_trace_heights(&self) -> MemoryTraceHeights {
640        let access_adapters = self.access_adapters.get_heights();
641        match &self.interface_chip {
642            MemoryInterface::Volatile { boundary_chip } => {
643                MemoryTraceHeights::Volatile(VolatileMemoryTraceHeights {
644                    boundary: boundary_chip.current_trace_height(),
645                    access_adapters,
646                })
647            }
648            MemoryInterface::Persistent {
649                boundary_chip,
650                merkle_chip,
651                ..
652            } => MemoryTraceHeights::Persistent(PersistentMemoryTraceHeights {
653                boundary: boundary_chip.current_trace_height(),
654                merkle: merkle_chip.current_trace_height(),
655                access_adapters,
656            }),
657        }
658    }
659
660    pub fn get_dummy_memory_trace_heights(&self) -> MemoryTraceHeights {
661        let access_adapters = vec![1; self.access_adapters.num_access_adapters()];
662        match &self.interface_chip {
663            MemoryInterface::Volatile { .. } => {
664                MemoryTraceHeights::Volatile(VolatileMemoryTraceHeights {
665                    boundary: 1,
666                    access_adapters,
667                })
668            }
669            MemoryInterface::Persistent { .. } => {
670                MemoryTraceHeights::Persistent(PersistentMemoryTraceHeights {
671                    boundary: 1,
672                    merkle: 1,
673                    access_adapters,
674                })
675            }
676        }
677    }
678
679    pub fn current_trace_cells(&self) -> Vec<usize> {
680        let mut ret = Vec::new();
681        match &self.interface_chip {
682            MemoryInterface::Volatile { boundary_chip } => {
683                ret.push(boundary_chip.current_trace_cells())
684            }
685            MemoryInterface::Persistent {
686                boundary_chip,
687                merkle_chip,
688                ..
689            } => {
690                ret.push(boundary_chip.current_trace_cells());
691                ret.push(merkle_chip.current_trace_cells());
692            }
693        }
694        ret.extend(self.access_adapters.get_cells());
695        ret
696    }
697
698    /// Returns a reference to the offline memory.
699    ///
700    /// Until `finalize` is called, the `OfflineMemory` does not contain useful state, and should
701    /// therefore not be used by any chip during execution. However, to obtain a reference to the
702    /// offline memory that will be useful in trace generation, a chip can call `offline_memory()`
703    /// and store the returned reference for later use.
704    pub fn offline_memory(&self) -> Arc<Mutex<OfflineMemory<F>>> {
705        self.offline_memory.clone()
706    }
707    pub fn get_memory_logs(&self) -> &Vec<MemoryLogEntry<F>> {
708        &self.memory.log
709    }
710    pub fn set_memory_logs(&mut self, logs: Vec<MemoryLogEntry<F>>) {
711        self.memory.log = logs;
712    }
713    pub fn take_memory_logs(&mut self) -> Vec<MemoryLogEntry<F>> {
714        std::mem::take(&mut self.memory.log)
715    }
716}
717
718pub struct MemoryAuxColsFactory<T> {
719    pub(crate) range_checker: SharedVariableRangeCheckerChip,
720    pub(crate) timestamp_lt_air: AssertLtSubAir,
721    pub(crate) _marker: PhantomData<T>,
722}
723
724// NOTE[jpw]: The `make_*_aux_cols` functions should be thread-safe so they can be used in
725// parallelized trace generation.
726impl<F: PrimeField32> MemoryAuxColsFactory<F> {
727    pub fn generate_read_aux(&self, read: &MemoryRecord<F>, buffer: &mut MemoryReadAuxCols<F>) {
728        assert!(
729            !read.address_space.is_zero(),
730            "cannot make `MemoryReadAuxCols` for address space 0"
731        );
732        self.generate_base_aux(read, &mut buffer.base);
733    }
734
735    pub fn generate_read_or_immediate_aux(
736        &self,
737        read: &MemoryRecord<F>,
738        buffer: &mut MemoryReadOrImmediateAuxCols<F>,
739    ) {
740        IsZeroSubAir.generate_subrow(
741            read.address_space,
742            (&mut buffer.is_zero_aux, &mut buffer.is_immediate),
743        );
744        self.generate_base_aux(read, &mut buffer.base);
745    }
746
747    pub fn generate_write_aux<const N: usize>(
748        &self,
749        write: &MemoryRecord<F>,
750        buffer: &mut MemoryWriteAuxCols<F, N>,
751    ) {
752        buffer
753            .prev_data
754            .copy_from_slice(write.prev_data_slice().unwrap());
755        self.generate_base_aux(write, &mut buffer.base);
756    }
757
758    pub fn generate_base_aux(&self, record: &MemoryRecord<F>, buffer: &mut MemoryBaseAuxCols<F>) {
759        buffer.prev_timestamp = F::from_canonical_u32(record.prev_timestamp);
760        self.generate_timestamp_lt(
761            record.prev_timestamp,
762            record.timestamp,
763            &mut buffer.timestamp_lt_aux,
764        );
765    }
766
767    fn generate_timestamp_lt(
768        &self,
769        prev_timestamp: u32,
770        timestamp: u32,
771        buffer: &mut LessThanAuxCols<F, AUX_LEN>,
772    ) {
773        debug_assert!(prev_timestamp < timestamp);
774        self.timestamp_lt_air.generate_subrow(
775            (self.range_checker.as_ref(), prev_timestamp, timestamp),
776            &mut buffer.lower_decomp,
777        );
778    }
779
780    /// In general, prefer `generate_read_aux` which writes in-place rather than this function.
781    pub fn make_read_aux_cols(&self, read: &MemoryRecord<F>) -> MemoryReadAuxCols<F> {
782        assert!(
783            !read.address_space.is_zero(),
784            "cannot make `MemoryReadAuxCols` for address space 0"
785        );
786        MemoryReadAuxCols::new(
787            read.prev_timestamp,
788            self.generate_timestamp_lt_cols(read.prev_timestamp, read.timestamp),
789        )
790    }
791
792    /// In general, prefer `generate_write_aux` which writes in-place rather than this function.
793    pub fn make_write_aux_cols<const N: usize>(
794        &self,
795        write: &MemoryRecord<F>,
796    ) -> MemoryWriteAuxCols<F, N> {
797        let prev_data = write.prev_data_slice().unwrap();
798        MemoryWriteAuxCols::new(
799            prev_data.try_into().unwrap(),
800            F::from_canonical_u32(write.prev_timestamp),
801            self.generate_timestamp_lt_cols(write.prev_timestamp, write.timestamp),
802        )
803    }
804
805    fn generate_timestamp_lt_cols(
806        &self,
807        prev_timestamp: u32,
808        timestamp: u32,
809    ) -> LessThanAuxCols<F, AUX_LEN> {
810        debug_assert!(prev_timestamp < timestamp);
811        let mut decomp = [F::ZERO; AUX_LEN];
812        self.timestamp_lt_air.generate_subrow(
813            (self.range_checker.as_ref(), prev_timestamp, timestamp),
814            &mut decomp,
815        );
816        LessThanAuxCols::new(decomp)
817    }
818}
819
820#[cfg(test)]
821mod tests {
822    use openvm_circuit_primitives::var_range::{
823        SharedVariableRangeCheckerChip, VariableRangeCheckerBus,
824    };
825    use openvm_stark_backend::{interaction::BusIndex, p3_field::FieldAlgebra};
826    use openvm_stark_sdk::p3_baby_bear::BabyBear;
827    use rand::{prelude::SliceRandom, thread_rng, Rng};
828
829    use super::MemoryController;
830    use crate::{
831        arch::{testing::MEMORY_BUS, MemoryConfig},
832        system::memory::offline_checker::MemoryBus,
833    };
834
835    const RANGE_CHECKER_BUS: BusIndex = 3;
836
837    #[test]
838    fn test_no_adapter_records_for_singleton_accesses() {
839        type F = BabyBear;
840
841        let memory_bus = MemoryBus::new(MEMORY_BUS);
842        let memory_config = MemoryConfig::default();
843        let range_bus = VariableRangeCheckerBus::new(RANGE_CHECKER_BUS, memory_config.decomp);
844        let range_checker = SharedVariableRangeCheckerChip::new(range_bus);
845
846        let mut memory_controller = MemoryController::with_volatile_memory(
847            memory_bus,
848            memory_config,
849            range_checker.clone(),
850        );
851
852        let mut rng = thread_rng();
853        for _ in 0..1000 {
854            let address_space = F::from_canonical_u32(*[1, 2].choose(&mut rng).unwrap());
855            let pointer =
856                F::from_canonical_u32(rng.gen_range(0..1 << memory_config.pointer_max_bits));
857
858            if rng.gen_bool(0.5) {
859                let data = F::from_canonical_u32(rng.gen_range(0..1 << 30));
860                memory_controller.write(address_space, pointer, [data]);
861            } else {
862                memory_controller.read::<1>(address_space, pointer);
863            }
864        }
865        assert!(memory_controller
866            .access_adapters
867            .get_heights()
868            .iter()
869            .all(|&h| h == 0));
870    }
871}
openvm_circuit/system/memory/controller/mod.rs

openvm_circuit/system/memory/controller/
mod.rs
