openvm_circuit/arch/
extensions.rs

1//! Traits and builders to compose collections of chips into a virtual machine.
2//!
3//! A full VM extension consists of three components, represented by sub-traits:
4//! - [VmExecutionExtension]
5//! - [VmCircuitExtension]
6//! - [VmProverExtension]: there may be multiple implementations of `VmProverExtension` for the same
7//!   `VmCircuitExtension` for different prover backends.
8//!
9//! It is intended that `VmExecutionExtension` and `VmCircuitExtension` are implemented on the
10//! same struct and `VmProverExtension` is implemented on a separate struct (usually a ZST) to
11//! get around Rust orphan rules.
12use std::{
13    any::{type_name, Any},
14    iter::{self, zip},
15    sync::Arc,
16};
17
18use getset::{CopyGetters, Getters};
19use openvm_circuit_primitives::{
20    var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerAir},
21    AnyChip, Chip, ColumnsAir,
22};
23use openvm_cpu_backend::CpuBackend;
24use openvm_instructions::{PhantomDiscriminant, VmOpcode};
25use openvm_stark_backend::{
26    interaction::BusIndex,
27    keygen::{types::MultiStarkProvingKey, MultiStarkKeygenBuilder},
28    prover::{AirProvingContext, MatrixDimensions, ProverBackend, ProvingContext},
29    AirRef, AnyAir, StarkEngine, StarkProtocolConfig, Val,
30};
31use rustc_hash::FxHashMap;
32use tracing::info_span;
33
34use super::{GenerationError, PhantomSubExecutor, SystemConfig};
35use crate::{
36    arch::Arena,
37    system::{
38        memory::{BOUNDARY_AIR_OFFSET, MERKLE_AIR_OFFSET},
39        phantom::PhantomExecutor,
40        SystemAirInventory, SystemChipComplex, SystemRecords,
41    },
42};
43
44/// Global AIR ID in the VM circuit verifying key.
45pub const PROGRAM_AIR_ID: usize = 0;
46/// ProgramAir is the first AIR so its cached trace should be the first main trace.
47pub const PROGRAM_CACHED_TRACE_INDEX: usize = 0;
48pub const CONNECTOR_AIR_ID: usize = 1;
49/// Starting AIR index of memory AIRs in the VM circuit.
50pub const MEMORY_AIRS_START_IDX: usize = 2;
51/// AIR index of the boundary AIR in the VM circuit.
52pub const BOUNDARY_AIR_ID: usize = MEMORY_AIRS_START_IDX + BOUNDARY_AIR_OFFSET;
53/// If VM has continuations enabled, all AIRs of MemoryController are added after ConnectorChip.
54/// Merkle AIR commits start/final memory states.
55pub const MERKLE_AIR_ID: usize = MEMORY_AIRS_START_IDX + MERKLE_AIR_OFFSET;
56
57pub type ExecutorId = u32;
58
59/// AIR trait object combining [`AnyAir`] (used by stark-backend for proving) with
60/// [`ColumnsAir`] (OpenVM-internal column-name introspection used by external tooling). The
61/// blanket impl below makes every type satisfying both traits also satisfy this one, so existing
62/// concrete AIRs need no changes.
63///
64/// Trait upcasting (stable since Rust 1.86) coerces `Arc<dyn AnyAirWithColumns<SC>>` to
65/// `Arc<dyn AnyAir<SC>>` in argument position, so [`AirRefWithColumns`] passes transparently to
66/// stark-backend APIs that expect [`AirRef`](openvm_stark_backend::AirRef).
67pub trait AnyAirWithColumns<SC: StarkProtocolConfig>: AnyAir<SC> + ColumnsAir {}
68
69impl<SC, T> AnyAirWithColumns<SC> for T
70where
71    SC: StarkProtocolConfig,
72    T: AnyAir<SC> + ColumnsAir,
73{
74}
75
76/// Reference-counted dyn pointer to an AIR with column-name introspection.
77pub type AirRefWithColumns<SC> = Arc<dyn AnyAirWithColumns<SC>>;
78
79// ======================= VM Extension Traits =============================
80
81/// Extension of VM execution. Allows registration of custom execution of new instructions by
82/// opcode.
83pub trait VmExecutionExtension<F> {
84    /// Enum of executor variants
85    type Executor: AnyEnum;
86
87    fn extend_execution(
88        &self,
89        inventory: &mut ExecutorInventoryBuilder<F, Self::Executor>,
90    ) -> Result<(), ExecutorInventoryError>;
91}
92
93/// Extension of the VM circuit. Allows _in-order_ addition of new AIRs with interactions.
94pub trait VmCircuitExtension<SC: StarkProtocolConfig> {
95    fn extend_circuit(&self, inventory: &mut AirInventory<SC>) -> Result<(), AirInventoryError>;
96}
97
98/// Extension of VM trace generation. The generics are `E` for [StarkEngine], `RA` for record arena,
99/// and `EXT` for execution and circuit extension.
100///
101/// Note that this trait differs from [VmExecutionExtension] and [VmCircuitExtension]. This trait is
102/// meant to be implemented on a separate ZST which may be different for different [ProverBackend]s.
103/// This is done to get around Rust orphan rules.
104pub trait VmProverExtension<E, RA, EXT>
105where
106    E: StarkEngine,
107    EXT: VmExecutionExtension<Val<E::SC>> + VmCircuitExtension<E::SC>,
108{
109    /// The chips added to `inventory` should exactly match the order of AIRs in the
110    /// [VmCircuitExtension] implementation of `EXT`.
111    ///
112    /// We do not provide access to the [ExecutorInventory] because the process to find an executor
113    /// from the inventory seems more cumbersome than to simply re-construct any necessary executors
114    /// directly within this function implementation.
115    fn extend_prover(
116        &self,
117        extension: &EXT,
118        inventory: &mut ChipInventory<E::SC, RA, E::PB>,
119    ) -> Result<(), ChipInventoryError>;
120}
121
122// ======================= Different Inventory Struct Definitions =============================
123
124pub struct ExecutorInventory<E> {
125    config: SystemConfig,
126    /// Lookup table to executor ID.
127    /// This is stored in a hashmap because it is _not_ expected to be used in the hot path.
128    /// A direct opcode -> executor mapping should be generated before runtime execution.
129    pub instruction_lookup: FxHashMap<VmOpcode, ExecutorId>,
130    pub executors: Vec<E>,
131    /// `ext_start[i]` will have the starting index in `executors` for extension `i`
132    ext_start: Vec<usize>,
133}
134
135// @dev: We need ExecutorInventoryBuilder separate from ExecutorInventory because of how
136// ExecutorInventory::extend works: we want to build an inventory with some big E3 enum that
137// includes both enum types E1, E2. However the interface for an ExecutionExtension will only know
138// about the enum E2. In order to be able to allow access to the old executors with type E1 without
139// referring to the type E1, we need to create this separate builder struct.
140pub struct ExecutorInventoryBuilder<'a, F, E> {
141    /// Chips that are already included in the chipset and may be used
142    /// as dependencies. The order should be that depended-on chips are ordered
143    /// **before** their dependents.
144    old_executors: Vec<&'a dyn AnyEnum>,
145    new_inventory: ExecutorInventory<E>,
146    phantom_executors: FxHashMap<PhantomDiscriminant, Arc<dyn PhantomSubExecutor<F>>>,
147}
148
149#[derive(Clone, Getters, CopyGetters)]
150pub struct AirInventory<SC: StarkProtocolConfig> {
151    #[get = "pub"]
152    config: SystemConfig,
153    /// The system AIRs required by the circuit architecture.
154    #[get = "pub"]
155    system: SystemAirInventory,
156    /// List of all non-system AIRs in the circuit, in insertion order, which is the **reverse** of
157    /// the order they appear in the verifying key.
158    ///
159    /// Note that the system will ensure that the first AIR in the list is always the
160    /// [VariableRangeCheckerAir].
161    #[get = "pub"]
162    ext_airs: Vec<AirRefWithColumns<SC>>,
163    /// `ext_start[i]` will have the starting index in `ext_airs` for extension `i`
164    ext_start: Vec<usize>,
165
166    bus_idx_mgr: BusIndexManager,
167}
168
169#[derive(Clone, Copy, Debug, Default)]
170pub struct BusIndexManager {
171    /// All existing buses use indices in [0, bus_idx_max)
172    bus_idx_max: BusIndex,
173}
174
175// @dev: ChipInventory does not have the SystemChipComplex because that is custom depending on `PB`.
176// The full struct with SystemChipComplex is VmChipComplex
177#[derive(Getters)]
178pub struct ChipInventory<SC, RA, PB>
179where
180    SC: StarkProtocolConfig,
181    PB: ProverBackend,
182{
183    /// Read-only view of AIRs, as constructed via the [VmCircuitExtension] trait.
184    #[get = "pub"]
185    airs: AirInventory<SC>,
186    /// Chips that are being built.
187    #[get = "pub"]
188    chips: Vec<Box<dyn AnyChip<RA, PB>>>,
189
190    /// Number of extensions that have chips added, including the current one that is still being
191    /// built.
192    cur_num_exts: usize,
193    /// Mapping from executor index to chip insertion index. Chips must be added in order so the
194    /// chip insertion index matches the AIR insertion index. Reminder: this is in **reverse**
195    /// order of the verifying key AIR ordering.
196    ///
197    /// Note: if public values chip exists, then it will be the first entry and point to
198    /// `usize::MAX`. This entry should never be used.
199    pub executor_idx_to_insertion_idx: Vec<usize>,
200}
201
202/// The collection of all chips in the VM. The chips should correspond 1-to-1 with the associated
203/// [AirInventory]. The [VmChipComplex] coordinates the trace generation for all chips in the VM
204/// after construction.
205#[derive(Getters)]
206pub struct VmChipComplex<SC, RA, PB, SCC>
207where
208    SC: StarkProtocolConfig,
209    PB: ProverBackend,
210{
211    /// System chip complex responsible for trace generation of [SystemAirInventory]
212    pub system: SCC,
213    pub inventory: ChipInventory<SC, RA, PB>,
214}
215
216// ======================= Inventory Function Definitions =============================
217
218impl<E> ExecutorInventory<E> {
219    /// Empty inventory should be created at the start of the declaration of a new extension.
220    #[allow(clippy::new_without_default)]
221    pub fn new(config: SystemConfig) -> Self {
222        Self {
223            config,
224            instruction_lookup: Default::default(),
225            executors: Default::default(),
226            ext_start: vec![0],
227        }
228    }
229
230    /// Inserts an executor with the collection of opcodes that it handles.
231    /// If some executor already owns one of the opcodes, an error is returned with the existing
232    /// executor.
233    pub fn add_executor(
234        &mut self,
235        executor: impl Into<E>,
236        opcodes: impl IntoIterator<Item = VmOpcode>,
237    ) -> Result<(), ExecutorInventoryError> {
238        let opcodes: Vec<_> = opcodes.into_iter().collect();
239        for opcode in &opcodes {
240            if let Some(id) = self.instruction_lookup.get(opcode) {
241                return Err(ExecutorInventoryError::ExecutorExists {
242                    opcode: *opcode,
243                    id: *id,
244                });
245            }
246        }
247        let id = self.executors.len();
248        self.executors.push(executor.into());
249        for opcode in opcodes {
250            self.instruction_lookup
251                .insert(opcode, id.try_into().unwrap());
252        }
253        Ok(())
254    }
255
256    /// Extend the inventory with a new extension.
257    /// A new inventory with different type generics is returned with the combined inventory.
258    pub fn extend<F, E3, EXT>(
259        self,
260        other: &EXT,
261    ) -> Result<ExecutorInventory<E3>, ExecutorInventoryError>
262    where
263        F: 'static,
264        E: Into<E3> + AnyEnum,
265        E3: AnyEnum,
266        EXT: VmExecutionExtension<F>,
267        EXT::Executor: Into<E3>,
268    {
269        let mut builder: ExecutorInventoryBuilder<F, EXT::Executor> = self.builder();
270        other.extend_execution(&mut builder)?;
271        let other_inventory = builder.new_inventory;
272        let other_phantom_executors = builder.phantom_executors;
273        let mut inventory_ext = self.transmute();
274        inventory_ext.append(other_inventory.transmute())?;
275        let phantom_chip: &mut PhantomExecutor<F> = inventory_ext
276            .find_executor_mut()
277            .next()
278            .expect("system always has phantom chip");
279        let phantom_executors = &mut phantom_chip.phantom_executors;
280        for (discriminant, sub_executor) in other_phantom_executors {
281            if phantom_executors
282                .insert(discriminant, sub_executor)
283                .is_some()
284            {
285                return Err(ExecutorInventoryError::PhantomSubExecutorExists { discriminant });
286            }
287        }
288
289        Ok(inventory_ext)
290    }
291
292    pub fn builder<F, E2>(&self) -> ExecutorInventoryBuilder<'_, F, E2>
293    where
294        F: 'static,
295        E: AnyEnum,
296    {
297        let old_executors = self.executors.iter().map(|e| e as &dyn AnyEnum).collect();
298        ExecutorInventoryBuilder {
299            old_executors,
300            new_inventory: ExecutorInventory::new(self.config.clone()),
301            phantom_executors: Default::default(),
302        }
303    }
304
305    pub fn transmute<E2>(self) -> ExecutorInventory<E2>
306    where
307        E: Into<E2>,
308    {
309        ExecutorInventory {
310            config: self.config,
311            instruction_lookup: self.instruction_lookup,
312            executors: self.executors.into_iter().map(|e| e.into()).collect(),
313            ext_start: self.ext_start,
314        }
315    }
316
317    /// Append `other` to current inventory. This means `self` comes earlier in the dependency
318    /// chain.
319    fn append(&mut self, mut other: ExecutorInventory<E>) -> Result<(), ExecutorInventoryError> {
320        let num_executors = self.executors.len();
321        for (opcode, mut id) in other.instruction_lookup.into_iter() {
322            id = id.checked_add(num_executors.try_into().unwrap()).unwrap();
323            if let Some(old_id) = self.instruction_lookup.insert(opcode, id) {
324                return Err(ExecutorInventoryError::ExecutorExists { opcode, id: old_id });
325            }
326        }
327        for id in &mut other.ext_start {
328            *id = id.checked_add(num_executors).unwrap();
329        }
330        self.executors.append(&mut other.executors);
331        self.ext_start.append(&mut other.ext_start);
332        Ok(())
333    }
334
335    pub fn get_executor(&self, opcode: VmOpcode) -> Option<&E> {
336        let id = self.instruction_lookup.get(&opcode)?;
337        self.executors.get(*id as usize)
338    }
339
340    pub fn get_mut_executor(&mut self, opcode: &VmOpcode) -> Option<&mut E> {
341        let id = self.instruction_lookup.get(opcode)?;
342        self.executors.get_mut(*id as usize)
343    }
344
345    pub fn executors(&self) -> &[E] {
346        &self.executors
347    }
348
349    pub fn find_executor<EX: 'static>(&self) -> impl Iterator<Item = &'_ EX>
350    where
351        E: AnyEnum,
352    {
353        self.executors
354            .iter()
355            .filter_map(|e| e.as_any_kind().downcast_ref())
356    }
357
358    pub fn find_executor_mut<EX: 'static>(&mut self) -> impl Iterator<Item = &'_ mut EX>
359    where
360        E: AnyEnum,
361    {
362        self.executors
363            .iter_mut()
364            .filter_map(|e| e.as_any_kind_mut().downcast_mut())
365    }
366
367    /// Returns the system config of the inventory.
368    pub fn config(&self) -> &SystemConfig {
369        &self.config
370    }
371}
372
373impl<F, E> ExecutorInventoryBuilder<'_, F, E> {
374    pub fn add_executor(
375        &mut self,
376        executor: impl Into<E>,
377        opcodes: impl IntoIterator<Item = VmOpcode>,
378    ) -> Result<(), ExecutorInventoryError> {
379        self.new_inventory.add_executor(executor, opcodes)
380    }
381
382    pub fn add_phantom_sub_executor<PE>(
383        &mut self,
384        phantom_sub: PE,
385        discriminant: PhantomDiscriminant,
386    ) -> Result<(), ExecutorInventoryError>
387    where
388        E: AnyEnum,
389        F: 'static,
390        PE: PhantomSubExecutor<F> + 'static,
391    {
392        let existing = self
393            .phantom_executors
394            .insert(discriminant, Arc::new(phantom_sub));
395        if existing.is_some() {
396            return Err(ExecutorInventoryError::PhantomSubExecutorExists { discriminant });
397        }
398        Ok(())
399    }
400
401    pub fn find_executor<EX: 'static>(&self) -> impl Iterator<Item = &'_ EX>
402    where
403        E: AnyEnum,
404    {
405        self.old_executors
406            .iter()
407            .filter_map(|e| e.as_any_kind().downcast_ref())
408    }
409
410    /// Returns the maximum number of bits used to represent addresses in memory
411    pub fn pointer_max_bits(&self) -> usize {
412        self.new_inventory.config().memory_config.pointer_max_bits
413    }
414}
415
416impl<SC: StarkProtocolConfig> AirInventory<SC> {
417    /// Outside of this crate, [AirInventory] must be constructed via [SystemConfig].
418    pub(crate) fn new(
419        config: SystemConfig,
420        system: SystemAirInventory,
421        bus_idx_mgr: BusIndexManager,
422    ) -> Self {
423        Self {
424            config,
425            system,
426            ext_start: Vec::new(),
427            ext_airs: Vec::new(),
428            bus_idx_mgr,
429        }
430    }
431
432    /// This should be called **exactly once** at the start of the declaration of a new extension.
433    pub fn start_new_extension(&mut self) {
434        self.ext_start.push(self.ext_airs.len());
435    }
436
437    pub fn new_bus_idx(&mut self) -> BusIndex {
438        self.bus_idx_mgr.new_bus_idx()
439    }
440
441    /// Looks through already-defined AIRs to see if there exists any of type `A` by downcasting.
442    /// Returns all chips of type `A` in the circuit.
443    ///
444    /// This should not be used to look for system AIRs.
445    pub fn find_air<A: 'static>(&self) -> impl Iterator<Item = &'_ A> {
446        self.ext_airs
447            .iter()
448            .filter_map(|air| air.as_any().downcast_ref())
449    }
450
451    pub fn add_air<A: AnyAirWithColumns<SC> + 'static>(&mut self, air: A) {
452        self.add_air_ref(Arc::new(air));
453    }
454
455    pub fn add_air_ref(&mut self, air: AirRefWithColumns<SC>) {
456        self.ext_airs.push(air);
457    }
458
459    pub fn range_checker(&self) -> &VariableRangeCheckerAir {
460        self.find_air()
461            .next()
462            .expect("system always has range checker AIR")
463    }
464
465    /// The AIRs in the order they appear in the verifying key.
466    /// This is the system AIRs, followed by the other AIRs in the **reverse** of the order they
467    /// were added in the VM extension definitions. In particular, the AIRs that have dependencies
468    /// appear later. The system guarantees that the last AIR is the [VariableRangeCheckerAir].
469    pub fn into_airs(self) -> impl Iterator<Item = AirRefWithColumns<SC>> {
470        self.system
471            .into_airs()
472            .into_iter()
473            .chain(self.ext_airs.into_iter().rev())
474    }
475
476    /// Generates the proving key for this circuit, marking the system AIRs that must be present
477    /// in any valid proof (see [`SystemConfig::is_required_air_id`]) as required.
478    pub fn keygen(self, config: &SC) -> MultiStarkProvingKey<SC> {
479        let system_config = self.config.clone();
480        let mut keygen_builder = MultiStarkKeygenBuilder::new(config.clone());
481        for (air_id, air) in self.into_airs().enumerate() {
482            if system_config.is_required_air_id(air_id) {
483                keygen_builder.add_required_air(air as AirRef<_>);
484            } else {
485                keygen_builder.add_air(air as AirRef<_>);
486            }
487        }
488        keygen_builder.generate_pk().unwrap()
489    }
490
491    /// This is O(1). Returns the total number of AIRs and equals the length of [`Self::into_airs`].
492    pub fn num_airs(&self) -> usize {
493        self.config.num_airs() + self.ext_airs.len()
494    }
495
496    /// Returns the maximum number of bits used to represent addresses in memory
497    pub fn pointer_max_bits(&self) -> usize {
498        self.config.memory_config.pointer_max_bits
499    }
500}
501
502impl BusIndexManager {
503    pub fn new() -> Self {
504        Self { bus_idx_max: 0 }
505    }
506
507    pub fn new_bus_idx(&mut self) -> BusIndex {
508        let idx = self.bus_idx_max;
509        self.bus_idx_max = self.bus_idx_max.checked_add(1).unwrap();
510        idx
511    }
512}
513
514impl<SC, RA, PB> ChipInventory<SC, RA, PB>
515where
516    SC: StarkProtocolConfig,
517    PB: ProverBackend,
518{
519    pub fn new(airs: AirInventory<SC>) -> Self {
520        Self {
521            airs,
522            chips: Vec::new(),
523            cur_num_exts: 0,
524            executor_idx_to_insertion_idx: Vec::new(),
525        }
526    }
527
528    pub fn config(&self) -> &SystemConfig {
529        &self.airs.config
530    }
531
532    // NOTE[jpw]: this is currently unused, it is for debugging purposes
533    pub fn start_new_extension(&mut self) -> Result<(), ChipInventoryError> {
534        if self.cur_num_exts >= self.airs.ext_start.len() {
535            return Err(ChipInventoryError::MissingCircuitExtension(
536                self.airs.ext_start.len(),
537            ));
538        }
539        if self.chips.len() != self.airs.ext_start[self.cur_num_exts] {
540            return Err(ChipInventoryError::MissingChip {
541                actual: self.chips.len(),
542                expected: self.airs.ext_start[self.cur_num_exts],
543            });
544        }
545
546        self.cur_num_exts += 1;
547        Ok(())
548    }
549
550    /// Gets the next AIR from the pre-existing AIR inventory according to the index of the next
551    /// chip to be built.
552    pub fn next_air<A: 'static>(&self) -> Result<&A, ChipInventoryError> {
553        let cur_idx = self.chips.len();
554        self.airs
555            .ext_airs
556            .get(cur_idx)
557            .and_then(|air| air.as_any().downcast_ref())
558            .ok_or_else(|| ChipInventoryError::AirNotFound {
559                name: type_name::<A>().to_string(),
560            })
561    }
562
563    /// Looks through built chips to see if there exists any of type `C` by downcasting.
564    /// Returns all chips of type `C` in the chipset.
565    ///
566    /// Note: the type `C` will usually be a smart pointer to a chip.
567    pub fn find_chip<C: 'static>(&self) -> impl Iterator<Item = &'_ C> {
568        self.chips.iter().filter_map(|c| c.as_any().downcast_ref())
569    }
570
571    /// Adds a chip that is not associated with any executor, as defined by the
572    /// [VmExecutionExtension] trait.
573    pub fn add_periphery_chip<C: Chip<RA, PB> + 'static>(&mut self, chip: C) {
574        self.chips.push(Box::new(chip));
575    }
576
577    /// Adds a chip and associates it to the next executor.
578    /// **Caution:** you must add chips in the order matching the order that executors were added in
579    /// the [VmExecutionExtension] implementation.
580    pub fn add_executor_chip<C: Chip<RA, PB> + 'static>(&mut self, chip: C) {
581        tracing::debug!("add_executor_chip: {}", type_name::<C>());
582        self.executor_idx_to_insertion_idx.push(self.chips.len());
583        self.chips.push(Box::new(chip));
584    }
585
586    /// Returns the mapping from executor index to the AIR index, where AIR index is the index of
587    /// the AIR within the verifying key.
588    ///
589    /// This should only be called after the `ChipInventory` is fully built.
590    pub fn executor_idx_to_air_idx(&self) -> Vec<usize> {
591        let num_airs = self.airs.num_airs();
592        assert_eq!(
593            num_airs,
594            self.config().num_airs() + self.chips.len(),
595            "Number of chips does not match number of AIRs"
596        );
597        // system AIRs are at the front of vkey, and then insertion index is the reverse ordering of
598        // AIR index
599        self.executor_idx_to_insertion_idx
600            .iter()
601            .map(|insertion_idx| {
602                num_airs
603                    .checked_sub(insertion_idx.checked_add(1).unwrap())
604                    .unwrap_or_else(|| {
605                        panic!(
606                            "Attempt to subtract num_airs={num_airs} by {}",
607                            insertion_idx + 1
608                        )
609                    })
610            })
611            .collect()
612    }
613
614    pub fn timestamp_max_bits(&self) -> usize {
615        self.airs.config().memory_config.timestamp_max_bits
616    }
617
618    /// Returns constant trace heights for all AIRs in verifying key order.
619    /// System AIRs get `None` (their constant heights are handled separately).
620    /// Extension chips follow in the same order as AIRs in the verifying key
621    /// (reversed insertion order).
622    pub fn constant_trace_heights(&self) -> Vec<Option<usize>> {
623        let num_system = self.airs.config().num_airs();
624        let mut heights = vec![None; num_system];
625        heights.extend(
626            self.chips
627                .iter()
628                .rev()
629                .map(|chip| chip.constant_trace_height()),
630        );
631        heights
632    }
633}
634
635// SharedVariableRangeCheckerChip is only used by the CPU backend.
636impl<SC, RA> ChipInventory<SC, RA, CpuBackend<SC>>
637where
638    SC: StarkProtocolConfig,
639{
640    pub fn range_checker(&self) -> Result<&SharedVariableRangeCheckerChip, ChipInventoryError> {
641        self.find_chip::<SharedVariableRangeCheckerChip>()
642            .next()
643            .ok_or_else(|| ChipInventoryError::ChipNotFound {
644                name: "VariableRangeCheckerChip".to_string(),
645            })
646    }
647}
648
649// ================================== Error Types =====================================
650
651#[derive(thiserror::Error, Debug)]
652pub enum ExecutorInventoryError {
653    #[error("Opcode {opcode} already owned by executor id {id}")]
654    ExecutorExists { opcode: VmOpcode, id: ExecutorId },
655    #[error("Phantom discriminant {} already has sub-executor", .discriminant.0)]
656    PhantomSubExecutorExists { discriminant: PhantomDiscriminant },
657}
658
659#[derive(thiserror::Error, Debug)]
660pub enum AirInventoryError {
661    #[error("AIR {name} not found")]
662    AirNotFound { name: String },
663}
664
665#[derive(thiserror::Error, Debug)]
666pub enum ChipInventoryError {
667    #[error("Air {name} not found")]
668    AirNotFound { name: String },
669    #[error("Chip {name} not found")]
670    ChipNotFound { name: String },
671    #[error("Adding prover extension without execution extension. Number of execution extensions is {0}")]
672    MissingExecutionExtension(usize),
673    #[error(
674        "Adding prover extension without circuit extension. Number of circuit extensions is {0}"
675    )]
676    MissingCircuitExtension(usize),
677    #[error("Missing chip. Number of chips is {actual}, expected number is {expected}")]
678    MissingChip { actual: usize, expected: usize },
679    #[error("Missing executor chip. Number of executors with associated chips is {actual}, expected number is {expected}")]
680    MissingExecutor { actual: usize, expected: usize },
681}
682
683// ======================= VM Chip Complex Implementation =============================
684
685impl<SC, RA, PB, SCC> VmChipComplex<SC, RA, PB, SCC>
686where
687    SC: StarkProtocolConfig,
688    RA: Arena,
689    PB: ProverBackend,
690    SCC: SystemChipComplex<RA, PB>,
691{
692    pub fn system_config(&self) -> &SystemConfig {
693        self.inventory.config()
694    }
695
696    /// `record_arenas` is expected to have length equal to the number of AIRs in the verifying key
697    /// and in the same order as the AIRs appearing in the verifying key, even though some chips may
698    /// not require a record arena.
699    pub(crate) fn generate_proving_ctx(
700        &mut self,
701        system_records: SystemRecords<PB::Val>,
702        record_arenas: Vec<RA>,
703        // trace_height_constraints: &[LinearConstraint],
704    ) -> Result<ProvingContext<PB>, GenerationError> {
705        // ATTENTION: The order of AIR proving context generation MUST be consistent with
706        // `AirInventory::into_airs`.
707
708        // Execution has finished at this point.
709        // ASSUMPTION WHICH MUST HOLD: non-system chips do not have a dependency on the system chips
710        // during trace generation. Given this assumption, we can generate trace on the system chips
711        // first.
712        let num_sys_airs = self.system_config().num_airs();
713        let num_airs = num_sys_airs + self.inventory.chips.len();
714        if num_airs != record_arenas.len() {
715            return Err(GenerationError::UnexpectedNumArenas {
716                actual: record_arenas.len(),
717                expected: num_airs,
718            });
719        }
720        let mut _record_arenas = record_arenas;
721        let record_arenas = _record_arenas.split_off(num_sys_airs);
722        let sys_record_arenas = _record_arenas;
723
724        // First go through all system chips
725        // Then go through all other chips in inventory in **reverse** order they were added (to
726        // resolve dependencies)
727        //
728        // Perf[jpw]: currently we call tracegen on each chip **serially** (although tracegen per
729        // chip is parallelized). We could introduce more parallelism, while potentially increasing
730        // the peak memory usage, by keeping a dependency tree and generating traces at the same
731        // layer of the tree in parallel.
732        let ctx_without_empties: Vec<(usize, AirProvingContext<_>)> = iter::empty()
733            .chain(info_span!("system_trace_gen").in_scope(|| {
734                self.system
735                    .generate_proving_ctx(system_records, sys_record_arenas)
736            }))
737            .chain(
738                zip(self.inventory.chips.iter().enumerate().rev(), record_arenas).map(
739                    |((insertion_idx, chip), records)| {
740                        // Only create a span if record is not empty:
741                        let _span = (!records.is_empty()).then(|| {
742                            let air_name = self.inventory.airs.ext_airs[insertion_idx].name();
743                            info_span!("single_trace_gen", air = air_name).entered()
744                        });
745                        #[cfg(feature = "metrics")]
746                        if let Some(allocated_bytes) = (!records.is_empty())
747                            .then(|| records.allocated_bytes())
748                            .flatten()
749                        {
750                            let air_name = self.inventory.airs.ext_airs[insertion_idx].name();
751                            let labels = [
752                                ("air_name", air_name.to_string()),
753                                ("air_id", (num_sys_airs + insertion_idx).to_string()),
754                            ];
755                            metrics::counter!("trace_gen.record_arena_bytes", &labels)
756                                .absolute(allocated_bytes as u64);
757                        }
758                        chip.generate_proving_ctx(records)
759                    },
760                ),
761            )
762            .enumerate()
763            .filter(|(_air_id, ctx)| ctx.common_main.height() > 0)
764            .collect();
765
766        Ok(ProvingContext::new(ctx_without_empties))
767    }
768}
769
770// ============ Blanket implementation of VM extension traits for Option<E> ===========
771
772impl<F, EXT: VmExecutionExtension<F>> VmExecutionExtension<F> for Option<EXT> {
773    type Executor = EXT::Executor;
774
775    fn extend_execution(
776        &self,
777        inventory: &mut ExecutorInventoryBuilder<F, Self::Executor>,
778    ) -> Result<(), ExecutorInventoryError> {
779        if let Some(extension) = self {
780            extension.extend_execution(inventory)
781        } else {
782            Ok(())
783        }
784    }
785}
786
787impl<SC: StarkProtocolConfig, EXT: VmCircuitExtension<SC>> VmCircuitExtension<SC> for Option<EXT> {
788    fn extend_circuit(&self, inventory: &mut AirInventory<SC>) -> Result<(), AirInventoryError> {
789        if let Some(extension) = self {
790            extension.extend_circuit(inventory)
791        } else {
792            Ok(())
793        }
794    }
795}
796
797/// A helper trait for downcasting types that may be enums.
798pub trait AnyEnum {
799    /// Recursively "unwraps" enum and casts to `Any` for downcasting.
800    fn as_any_kind(&self) -> &dyn Any;
801
802    /// Recursively "unwraps" enum and casts to `Any` for downcasting.
803    fn as_any_kind_mut(&mut self) -> &mut dyn Any;
804}
805
806impl AnyEnum for () {
807    fn as_any_kind(&self) -> &dyn Any {
808        self
809    }
810    fn as_any_kind_mut(&mut self) -> &mut dyn Any {
811        self
812    }
813}
814
815#[cfg(test)]
816mod tests {
817    use openvm_circuit_derive::AnyEnum;
818    use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config;
819
820    use super::*;
821    use crate::arch::VmCircuitConfig;
822
823    #[allow(dead_code)]
824    #[derive(Copy, Clone)]
825    enum EnumA {
826        A(u8),
827        B(u32),
828    }
829
830    enum EnumB {
831        C(u64),
832        D(EnumA),
833    }
834
835    #[derive(AnyEnum)]
836    enum EnumC {
837        C(u64),
838        #[any_enum]
839        D(EnumA),
840    }
841
842    impl AnyEnum for EnumA {
843        fn as_any_kind(&self) -> &dyn Any {
844            match self {
845                EnumA::A(a) => a,
846                EnumA::B(b) => b,
847            }
848        }
849
850        fn as_any_kind_mut(&mut self) -> &mut dyn Any {
851            match self {
852                EnumA::A(a) => a,
853                EnumA::B(b) => b,
854            }
855        }
856    }
857
858    impl AnyEnum for EnumB {
859        fn as_any_kind(&self) -> &dyn Any {
860            match self {
861                EnumB::C(c) => c,
862                EnumB::D(d) => d.as_any_kind(),
863            }
864        }
865
866        fn as_any_kind_mut(&mut self) -> &mut dyn Any {
867            match self {
868                EnumB::C(c) => c,
869                EnumB::D(d) => d.as_any_kind_mut(),
870            }
871        }
872    }
873
874    #[test]
875    fn test_any_enum_downcast() {
876        let a = EnumA::A(1);
877        assert_eq!(a.as_any_kind().downcast_ref::<u8>(), Some(&1));
878        let b = EnumB::D(a);
879        assert!(b.as_any_kind().downcast_ref::<u64>().is_none());
880        assert!(b.as_any_kind().downcast_ref::<EnumA>().is_none());
881        assert_eq!(b.as_any_kind().downcast_ref::<u8>(), Some(&1));
882        let c = EnumB::C(3);
883        assert_eq!(c.as_any_kind().downcast_ref::<u64>(), Some(&3));
884        let d = EnumC::D(a);
885        assert!(d.as_any_kind().downcast_ref::<u64>().is_none());
886        assert!(d.as_any_kind().downcast_ref::<EnumA>().is_none());
887        assert_eq!(d.as_any_kind().downcast_ref::<u8>(), Some(&1));
888        let e = EnumC::C(3);
889        assert_eq!(e.as_any_kind().downcast_ref::<u64>(), Some(&3));
890    }
891
892    #[test]
893    fn test_system_bus_indices() {
894        let config = SystemConfig::default();
895        let inventory: AirInventory<BabyBearPoseidon2Config> = config.create_airs().unwrap();
896        let system = inventory.system();
897        let port = system.port();
898        assert_eq!(port.execution_bus.index(), 0);
899        assert_eq!(port.memory_bridge.memory_bus().index(), 1);
900        assert_eq!(port.program_bus.index(), 2);
901        assert_eq!(port.memory_bridge.range_bus().index(), 3);
902        assert_eq!(system.memory.interface.boundary.merkle_bus.index, 4);
903        assert_eq!(system.memory.interface.boundary.compression_bus.index, 5);
904    }
905}
openvm_circuit/arch/extensions.rs

openvm_circuit/arch/
extensions.rs
