openvm_circuit/system/
mod.rs

1use std::sync::Arc;
2
3use derive_more::derive::From;
4use openvm_circuit_derive::{AnyEnum, Executor, MeteredExecutor, PreflightExecutor};
5use openvm_circuit_primitives::var_range::{
6    SharedVariableRangeCheckerChip, VariableRangeCheckerAir, VariableRangeCheckerBus,
7    VariableRangeCheckerChip,
8};
9use openvm_instructions::{
10    LocalOpcode, PhantomDiscriminant, PublishOpcode, SysPhantom, SystemOpcode,
11};
12use openvm_stark_backend::{
13    config::{StarkGenericConfig, Val},
14    engine::StarkEngine,
15    interaction::{LookupBus, PermutationCheckBus},
16    p3_field::{Field, PrimeField32},
17    prover::{
18        cpu::{CpuBackend, CpuDevice},
19        hal::{MatrixDimensions, ProverBackend},
20        types::{AirProvingContext, CommittedTraceData},
21    },
22    AirRef, Chip,
23};
24use rustc_hash::FxHashMap;
25
26use self::{connector::VmConnectorAir, program::ProgramAir, public_values::PublicValuesAir};
27use crate::{
28    arch::{
29        vm_poseidon2_config, AirInventory, AirInventoryError, BusIndexManager, ChipInventory,
30        ChipInventoryError, DenseRecordArena, ExecutionBridge, ExecutionBus, ExecutionState,
31        ExecutorInventory, ExecutorInventoryError, MatrixRecordArena, PhantomSubExecutor,
32        RowMajorMatrixArena, SystemConfig, VmAirWrapper, VmBuilder, VmChipComplex, VmChipWrapper,
33        VmCircuitConfig, VmExecutionConfig, CONNECTOR_AIR_ID, PROGRAM_AIR_ID, PUBLIC_VALUES_AIR_ID,
34    },
35    system::{
36        connector::VmConnectorChip,
37        memory::{
38            interface::MemoryInterfaceAirs,
39            offline_checker::{MemoryBridge, MemoryBus},
40            online::GuestMemory,
41            MemoryAirInventory, MemoryController, TimestampedEquipartition, CHUNK,
42        },
43        native_adapter::{NativeAdapterAir, NativeAdapterExecutor},
44        phantom::{
45            CycleEndPhantomExecutor, CycleStartPhantomExecutor, NopPhantomExecutor, PhantomAir,
46            PhantomChip, PhantomExecutor, PhantomFiller,
47        },
48        poseidon2::{
49            air::Poseidon2PeripheryAir, new_poseidon2_periphery_air, Poseidon2PeripheryChip,
50        },
51        program::{ProgramBus, ProgramChip},
52        public_values::{
53            PublicValuesChip, PublicValuesCoreAir, PublicValuesExecutor, PublicValuesFiller,
54        },
55    },
56};
57
58pub mod connector;
59#[cfg(feature = "cuda")]
60pub mod cuda;
61pub mod memory;
62// Necessary for the PublicValuesChip
63pub mod native_adapter;
64pub mod phantom;
65pub mod poseidon2;
66pub mod program;
67pub mod public_values;
68
69/// **If** internal poseidon2 chip exists, then its insertion index is 1.
70const POSEIDON2_INSERTION_IDX: usize = 1;
71/// **If** public values chip exists, then its executor index is 0.
72pub(crate) const PV_EXECUTOR_IDX: usize = 0;
73
74/// Trait for trace generation of all system AIRs. The system chip complex is special because we may
75/// not exactly following the exact matching between `Air` and `Chip`. Moreover we may require more
76/// flexibility than what is provided through the trait object [`AnyChip`].
77///
78/// The [SystemChipComplex] is meant to be constructible once the VM configuration is known, and it
79/// can be loaded with arbitrary programs supported by the instruction set available to its
80/// configuration. The [SystemChipComplex] is meant to persistent between instances of proof
81/// generation.
82pub trait SystemChipComplex<RA, PB: ProverBackend> {
83    /// Loads the program in the form of a cached trace with prover data.
84    fn load_program(&mut self, cached_program_trace: CommittedTraceData<PB>);
85
86    /// Transport the initial memory state to device. This may be called before preflight execution
87    /// begins and start async device processes in parallel to execution.
88    fn transport_init_memory_to_device(&mut self, memory: &GuestMemory);
89
90    /// The caller must guarantee that `record_arenas` has length equal to the number of system
91    /// AIRs, although some arenas may be empty if they are unused.
92    fn generate_proving_ctx(
93        &mut self,
94        system_records: SystemRecords<PB::Val>,
95        record_arenas: Vec<RA>,
96    ) -> Vec<AirProvingContext<PB>>;
97
98    /// This function is only used for metric collection purposes and custom implementations are
99    /// free to ignore it.
100    ///
101    /// Since system chips (primarily memory) will only have all information needed to compute the
102    /// true used trace heights after `generate_proving_ctx` is called, this method will be called
103    /// after `generate_proving_ctx` on the trace `heights` of all AIRs (including non-system AIRs)
104    /// in the AIR ID order.
105    ///
106    /// The default implementation does nothing.
107    #[cfg(feature = "metrics")]
108    fn finalize_trace_heights(&self, _heights: &mut [usize]) {}
109}
110
111/// Trait meant to be implemented on a SystemChipComplex.
112pub trait SystemWithFixedTraceHeights {
113    /// `heights` will have length equal to number of system AIRs, in AIR ID order. This function
114    /// must guarantee that the system trace matrices generated have the required heights.
115    fn override_trace_heights(&mut self, heights: &[u32]);
116}
117
118pub struct SystemRecords<F> {
119    pub from_state: ExecutionState<u32>,
120    pub to_state: ExecutionState<u32>,
121    pub exit_code: Option<u32>,
122    /// `i` -> frequency of instruction in `i`th row of trace matrix. This requires filtering
123    /// `program.instructions_and_debug_infos` to remove gaps.
124    pub filtered_exec_frequencies: Vec<u32>,
125    // We always use a [DenseRecordArena] here, regardless of the generic `RA` used for other
126    // execution records.
127    pub access_adapter_records: DenseRecordArena,
128    // Perf[jpw]: this should be computed on-device and changed to just touched blocks
129    pub touched_memory: TouchedMemory<F>,
130    /// The public values of the [PublicValuesChip]. These should only be non-empty if
131    /// continuations are disabled.
132    pub public_values: Vec<F>,
133}
134
135pub enum TouchedMemory<F> {
136    Persistent(TimestampedEquipartition<F, CHUNK>),
137    Volatile(TimestampedEquipartition<F, 1>),
138}
139
140#[derive(Clone, AnyEnum, Executor, MeteredExecutor, PreflightExecutor, From)]
141pub enum SystemExecutor<F: Field> {
142    PublicValues(PublicValuesExecutor<F>),
143    Phantom(PhantomExecutor<F>),
144}
145
146/// SystemPort combines system resources needed by most extensions
147#[derive(Clone, Copy)]
148pub struct SystemPort {
149    pub execution_bus: ExecutionBus,
150    pub program_bus: ProgramBus,
151    pub memory_bridge: MemoryBridge,
152}
153
154#[derive(Clone)]
155pub struct SystemAirInventory<SC: StarkGenericConfig> {
156    pub program: ProgramAir,
157    pub connector: VmConnectorAir,
158    pub memory: MemoryAirInventory<SC>,
159    /// Public values AIR exists if and only if continuations is disabled and `num_public_values`
160    /// is greater than 0.
161    pub public_values: Option<PublicValuesAir>,
162}
163
164impl<SC: StarkGenericConfig> SystemAirInventory<SC> {
165    pub fn new(
166        config: &SystemConfig,
167        port: SystemPort,
168        merkle_compression_buses: Option<(PermutationCheckBus, PermutationCheckBus)>,
169    ) -> Self {
170        let SystemPort {
171            execution_bus,
172            program_bus,
173            memory_bridge,
174        } = port;
175        let range_bus = memory_bridge.range_bus();
176        let program = ProgramAir::new(program_bus);
177        let connector = VmConnectorAir::new(
178            execution_bus,
179            program_bus,
180            range_bus,
181            config.memory_config.timestamp_max_bits,
182        );
183        assert_eq!(
184            config.continuation_enabled,
185            merkle_compression_buses.is_some()
186        );
187
188        let memory = MemoryAirInventory::new(
189            memory_bridge,
190            &config.memory_config,
191            range_bus,
192            merkle_compression_buses,
193        );
194
195        let public_values = if config.has_public_values_chip() {
196            let air = VmAirWrapper::new(
197                NativeAdapterAir::new(
198                    ExecutionBridge::new(execution_bus, program_bus),
199                    memory_bridge,
200                ),
201                PublicValuesCoreAir::new(
202                    config.num_public_values,
203                    config.max_constraint_degree as u32 - 1,
204                ),
205            );
206            Some(air)
207        } else {
208            None
209        };
210
211        Self {
212            program,
213            connector,
214            memory,
215            public_values,
216        }
217    }
218
219    pub fn port(&self) -> SystemPort {
220        SystemPort {
221            memory_bridge: self.memory.bridge,
222            program_bus: self.program.bus,
223            execution_bus: self.connector.execution_bus,
224        }
225    }
226
227    pub fn into_airs(self) -> Vec<AirRef<SC>> {
228        let mut airs: Vec<AirRef<SC>> = Vec::new();
229        airs.push(Arc::new(self.program));
230        airs.push(Arc::new(self.connector));
231        if let Some(public_values) = self.public_values {
232            airs.push(Arc::new(public_values));
233        }
234        airs.extend(self.memory.into_airs());
235        airs
236    }
237}
238
239impl<F: PrimeField32> VmExecutionConfig<F> for SystemConfig {
240    type Executor = SystemExecutor<F>;
241
242    /// The only way to create an [ExecutorInventory] is from a [SystemConfig]. This will add an
243    /// executor for [PublicValuesExecutor] if continuations is disabled. It will always add an
244    /// executor for [PhantomChip], which handles all phantom sub-executors.
245    fn create_executors(
246        &self,
247    ) -> Result<ExecutorInventory<Self::Executor>, ExecutorInventoryError> {
248        let mut inventory = ExecutorInventory::new(self.clone());
249        // PublicValuesChip is required when num_public_values > 0 in single segment mode.
250        if self.has_public_values_chip() {
251            assert_eq!(inventory.executors().len(), PV_EXECUTOR_IDX);
252
253            let public_values = PublicValuesExecutor::new(NativeAdapterExecutor::default());
254            inventory.add_executor(public_values, [PublishOpcode::PUBLISH.global_opcode()])?;
255        }
256        let phantom_opcode = SystemOpcode::PHANTOM.global_opcode();
257        let mut phantom_executors: FxHashMap<PhantomDiscriminant, Arc<dyn PhantomSubExecutor<F>>> =
258            FxHashMap::default();
259        // Use NopPhantomExecutor so the discriminant is set but `DebugPanic` is handled specially.
260        phantom_executors.insert(
261            PhantomDiscriminant(SysPhantom::DebugPanic as u16),
262            Arc::new(NopPhantomExecutor),
263        );
264        phantom_executors.insert(
265            PhantomDiscriminant(SysPhantom::Nop as u16),
266            Arc::new(NopPhantomExecutor),
267        );
268        phantom_executors.insert(
269            PhantomDiscriminant(SysPhantom::CtStart as u16),
270            Arc::new(CycleStartPhantomExecutor),
271        );
272        phantom_executors.insert(
273            PhantomDiscriminant(SysPhantom::CtEnd as u16),
274            Arc::new(CycleEndPhantomExecutor),
275        );
276        let phantom = PhantomExecutor::new(phantom_executors, phantom_opcode);
277        inventory.add_executor(phantom, [phantom_opcode])?;
278
279        Ok(inventory)
280    }
281}
282
283impl<SC: StarkGenericConfig> VmCircuitConfig<SC> for SystemConfig {
284    /// Every VM circuit within the OpenVM circuit architecture **must** be initialized from the
285    /// [SystemConfig].
286    fn create_airs(&self) -> Result<AirInventory<SC>, AirInventoryError> {
287        let mut bus_idx_mgr = BusIndexManager::new();
288        let execution_bus = ExecutionBus::new(bus_idx_mgr.new_bus_idx());
289        let memory_bus = MemoryBus::new(bus_idx_mgr.new_bus_idx());
290        let program_bus = ProgramBus::new(bus_idx_mgr.new_bus_idx());
291        let range_bus =
292            VariableRangeCheckerBus::new(bus_idx_mgr.new_bus_idx(), self.memory_config.decomp);
293
294        let merkle_compression_buses = if self.continuation_enabled {
295            let merkle_bus = PermutationCheckBus::new(bus_idx_mgr.new_bus_idx());
296            let compression_bus = PermutationCheckBus::new(bus_idx_mgr.new_bus_idx());
297            Some((merkle_bus, compression_bus))
298        } else {
299            None
300        };
301        let memory_bridge =
302            MemoryBridge::new(memory_bus, self.memory_config.timestamp_max_bits, range_bus);
303        let system_port = SystemPort {
304            execution_bus,
305            program_bus,
306            memory_bridge,
307        };
308        let system = SystemAirInventory::new(self, system_port, merkle_compression_buses);
309
310        let mut inventory = AirInventory::new(self.clone(), system, bus_idx_mgr);
311
312        let range_checker = VariableRangeCheckerAir::new(range_bus);
313        // Range checker is always the first AIR in the inventory
314        inventory.add_air(range_checker);
315
316        if self.continuation_enabled {
317            assert_eq!(inventory.ext_airs().len(), POSEIDON2_INSERTION_IDX);
318            // Add direct poseidon2 AIR for persistent memory.
319            // Currently we never use poseidon2 opcodes when continuations is enabled: we will need
320            // special handling when that happens
321            let (_, compression_bus) = merkle_compression_buses.unwrap();
322            let direct_bus_idx = compression_bus.index;
323            let air = new_poseidon2_periphery_air(
324                vm_poseidon2_config(),
325                LookupBus::new(direct_bus_idx),
326                self.max_constraint_degree,
327            );
328            inventory.add_air_ref(air);
329        }
330        let execution_bridge = ExecutionBridge::new(execution_bus, program_bus);
331        let phantom = PhantomAir {
332            execution_bridge,
333            phantom_opcode: SystemOpcode::PHANTOM.global_opcode(),
334        };
335        inventory.add_air(phantom);
336
337        Ok(inventory)
338    }
339}
340
341// =================== CPU Backend Specific System Chip Complex Constructor ==================
342
343/// Base system chips for CPU backend. These chips must exactly correspond to the AIRs in
344/// [SystemAirInventory].
345pub struct SystemChipInventory<SC: StarkGenericConfig> {
346    pub program_chip: ProgramChip<SC>,
347    pub connector_chip: VmConnectorChip<Val<SC>>,
348    /// Contains all memory chips
349    pub memory_controller: MemoryController<Val<SC>>,
350    pub public_values_chip: Option<PublicValuesChip<Val<SC>>>,
351}
352
353// Note[jpw]: We could get rid of the `mem_inventory` input because `MemoryController` doesn't need
354// the buses for tracegen. We leave it to use old interfaces.
355impl<SC: StarkGenericConfig> SystemChipInventory<SC>
356where
357    Val<SC>: PrimeField32,
358{
359    pub fn new(
360        config: &SystemConfig,
361        mem_inventory: &MemoryAirInventory<SC>,
362        range_checker: SharedVariableRangeCheckerChip,
363        hasher_chip: Option<Arc<Poseidon2PeripheryChip<Val<SC>>>>,
364    ) -> Self {
365        // We create an empty program chip: the program should be loaded later (and can be swapped
366        // out). The execution frequencies are supplied only after execution.
367        let program_chip = ProgramChip::unloaded();
368        let connector_chip = VmConnectorChip::<Val<SC>>::new(
369            range_checker.clone(),
370            config.memory_config.timestamp_max_bits,
371        );
372        let memory_bus = mem_inventory.bridge.memory_bus();
373        let memory_controller = match &mem_inventory.interface {
374            MemoryInterfaceAirs::Persistent {
375                boundary: _,
376                merkle,
377            } => {
378                assert!(config.continuation_enabled);
379                MemoryController::<Val<SC>>::with_persistent_memory(
380                    memory_bus,
381                    config.memory_config.clone(),
382                    range_checker.clone(),
383                    merkle.merkle_bus,
384                    merkle.compression_bus,
385                    hasher_chip.unwrap(),
386                )
387            }
388            MemoryInterfaceAirs::Volatile { boundary: _ } => {
389                assert!(!config.continuation_enabled);
390                MemoryController::with_volatile_memory(
391                    memory_bus,
392                    config.memory_config.clone(),
393                    range_checker.clone(),
394                )
395            }
396        };
397
398        let public_values_chip = config.has_public_values_chip().then(|| {
399            VmChipWrapper::new(
400                PublicValuesFiller::new(
401                    NativeAdapterExecutor::default(),
402                    config.num_public_values,
403                    (config.max_constraint_degree as u32)
404                        .checked_sub(1)
405                        .unwrap(),
406                ),
407                memory_controller.helper(),
408            )
409        });
410
411        Self {
412            program_chip,
413            connector_chip,
414            memory_controller,
415            public_values_chip,
416        }
417    }
418}
419
420impl<RA, SC> SystemChipComplex<RA, CpuBackend<SC>> for SystemChipInventory<SC>
421where
422    RA: RowMajorMatrixArena<Val<SC>>,
423    SC: StarkGenericConfig,
424    Val<SC>: PrimeField32,
425{
426    fn load_program(&mut self, cached_program_trace: CommittedTraceData<CpuBackend<SC>>) {
427        let _ = self.program_chip.cached.replace(cached_program_trace);
428    }
429
430    fn transport_init_memory_to_device(&mut self, memory: &GuestMemory) {
431        self.memory_controller
432            .set_initial_memory(memory.memory.clone());
433    }
434
435    fn generate_proving_ctx(
436        &mut self,
437        system_records: SystemRecords<Val<SC>>,
438        mut record_arenas: Vec<RA>,
439    ) -> Vec<AirProvingContext<CpuBackend<SC>>> {
440        let SystemRecords {
441            from_state,
442            to_state,
443            exit_code,
444            filtered_exec_frequencies,
445            access_adapter_records,
446            touched_memory,
447            public_values,
448        } = system_records;
449
450        if let Some(chip) = &mut self.public_values_chip {
451            chip.inner.set_public_values(public_values);
452        }
453        self.program_chip.filtered_exec_frequencies = filtered_exec_frequencies;
454        let program_ctx = self.program_chip.generate_proving_ctx(());
455        self.connector_chip.begin(from_state);
456        self.connector_chip.end(to_state, exit_code);
457        let connector_ctx = self.connector_chip.generate_proving_ctx(());
458
459        let pv_ctx = self.public_values_chip.as_ref().map(|chip| {
460            let arena = record_arenas.remove(PUBLIC_VALUES_AIR_ID);
461            chip.generate_proving_ctx(arena)
462        });
463
464        let memory_ctxs = self
465            .memory_controller
466            .generate_proving_ctx(access_adapter_records, touched_memory);
467
468        [program_ctx, connector_ctx]
469            .into_iter()
470            .chain(pv_ctx)
471            .chain(memory_ctxs)
472            .collect()
473    }
474
475    #[cfg(feature = "metrics")]
476    fn finalize_trace_heights(&self, heights: &mut [usize]) {
477        use openvm_stark_backend::ChipUsageGetter;
478
479        use crate::system::memory::interface::MemoryInterface;
480
481        let boundary_idx = PUBLIC_VALUES_AIR_ID + usize::from(self.public_values_chip.is_some());
482        let mut access_adapter_offset = boundary_idx + 1;
483        match &self.memory_controller.interface_chip {
484            MemoryInterface::Volatile { boundary_chip } => {
485                let boundary_height = boundary_chip
486                    .final_memory
487                    .as_ref()
488                    .map(|m| m.len())
489                    .unwrap_or(0);
490                heights[boundary_idx] = boundary_height;
491            }
492            MemoryInterface::Persistent {
493                boundary_chip,
494                merkle_chip,
495                ..
496            } => {
497                let boundary_height = 2 * boundary_chip.touched_labels.len();
498                heights[boundary_idx] = boundary_height;
499                heights[boundary_idx + 1] = merkle_chip.current_height;
500                access_adapter_offset += 1;
501
502                // Poseidon2Periphery height also varies based on memory, so set it now even though
503                // it's not a system chip:
504                let poseidon_chip = self.memory_controller.hasher_chip.as_ref().unwrap();
505                let poseidon_height = poseidon_chip.current_trace_height();
506                // We know the chip insertion index, which starts from *the end* of the the AIR
507                // ordering
508                let poseidon_idx = heights.len() - 1 - POSEIDON2_INSERTION_IDX;
509                heights[poseidon_idx] = poseidon_height;
510            }
511        }
512        let access_heights = &self
513            .memory_controller
514            .access_adapter_inventory
515            .trace_heights;
516        heights[access_adapter_offset..access_adapter_offset + access_heights.len()]
517            .copy_from_slice(access_heights);
518    }
519}
520
521#[derive(Clone)]
522pub struct SystemCpuBuilder;
523
524impl<SC, E> VmBuilder<E> for SystemCpuBuilder
525where
526    SC: StarkGenericConfig,
527    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
528    Val<SC>: PrimeField32,
529{
530    type VmConfig = SystemConfig;
531    type RecordArena = MatrixRecordArena<Val<SC>>;
532    type SystemChipInventory = SystemChipInventory<SC>;
533
534    fn create_chip_complex(
535        &self,
536        config: &SystemConfig,
537        airs: AirInventory<SC>,
538    ) -> Result<
539        VmChipComplex<SC, MatrixRecordArena<Val<SC>>, CpuBackend<SC>, SystemChipInventory<SC>>,
540        ChipInventoryError,
541    > {
542        let range_bus = airs.range_checker().bus;
543        let range_checker = Arc::new(VariableRangeCheckerChip::new(range_bus));
544
545        let mut inventory = ChipInventory::new(airs);
546        // PublicValuesChip is required when num_public_values > 0 in single segment mode.
547        if config.has_public_values_chip() {
548            assert_eq!(
549                inventory.executor_idx_to_insertion_idx.len(),
550                PV_EXECUTOR_IDX
551            );
552            // We set insertion_idx so that air_idx = num_airs - (insertion_idx + 1) =
553            // PUBLIC_VALUES_AIR_ID in `VmChipComplex::executor_idx_to_air_idx`. We need to do this
554            // because this chip is special and not part of the normal inventory.
555            let insertion_idx = inventory
556                .airs()
557                .num_airs()
558                .checked_sub(1 + PUBLIC_VALUES_AIR_ID)
559                .unwrap();
560            inventory.executor_idx_to_insertion_idx.push(insertion_idx);
561        }
562        inventory.next_air::<VariableRangeCheckerAir>()?;
563        inventory.add_periphery_chip(range_checker.clone());
564
565        let hasher_chip = if config.continuation_enabled {
566            assert_eq!(inventory.chips().len(), POSEIDON2_INSERTION_IDX);
567            // ATTENTION: The threshold 7 here must match the one in `new_poseidon2_periphery_air`
568            let direct_bus = if config.max_constraint_degree >= 7 {
569                inventory
570                    .next_air::<Poseidon2PeripheryAir<Val<SC>, 0>>()?
571                    .bus
572            } else {
573                inventory
574                    .next_air::<Poseidon2PeripheryAir<Val<SC>, 1>>()?
575                    .bus
576            };
577            let chip = Arc::new(Poseidon2PeripheryChip::new(
578                vm_poseidon2_config(),
579                direct_bus.index,
580                config.max_constraint_degree,
581            ));
582            inventory.add_periphery_chip(chip.clone());
583            Some(chip)
584        } else {
585            None
586        };
587        let system = SystemChipInventory::new(
588            config,
589            &inventory.airs().system().memory,
590            range_checker,
591            hasher_chip,
592        );
593
594        let phantom_chip = PhantomChip::new(PhantomFiller, system.memory_controller.helper());
595        inventory.add_executor_chip(phantom_chip);
596
597        Ok(VmChipComplex { system, inventory })
598    }
599}
600
601impl<SC: StarkGenericConfig> SystemWithFixedTraceHeights for SystemChipInventory<SC>
602where
603    Val<SC>: PrimeField32,
604{
605    /// Warning: this does not set the override for the PublicValuesChip. The PublicValuesChip
606    /// override must be set via the RecordArena.
607    fn override_trace_heights(&mut self, heights: &[u32]) {
608        assert_eq!(
609            heights[PROGRAM_AIR_ID] as usize,
610            self.program_chip
611                .cached
612                .as_ref()
613                .expect("program not loaded")
614                .trace
615                .height()
616        );
617        assert_eq!(heights[CONNECTOR_AIR_ID], 2);
618        let mut memory_start_idx = PUBLIC_VALUES_AIR_ID;
619        if self.public_values_chip.is_some() {
620            memory_start_idx += 1;
621        }
622        self.memory_controller
623            .set_override_trace_heights(&heights[memory_start_idx..]);
624    }
625}