Skip to main content

openvm_stark_backend/
any_air.rs

1//! Traits for [Air] trait objects.
2
3use std::{
4    any::{type_name, Any},
5    sync::Arc,
6};
7
8// Re-export for backwards compatibility
9pub use p3_air::BaseAirWithPublicValues;
10use p3_air::{Air, BaseAir};
11
12use crate::{
13    air_builders::{debug::DebugConstraintBuilder, symbolic::SymbolicRapBuilder},
14    config::StarkProtocolConfig,
15};
16
17/// An AIR with 1 or more main trace partitions.
18pub trait PartitionedBaseAir<F>: BaseAir<F> {
19    /// By default, an AIR has no cached main trace.
20    fn cached_main_widths(&self) -> Vec<usize> {
21        vec![]
22    }
23    /// By default, an AIR has only one private main trace.
24    fn common_main_width(&self) -> usize {
25        self.width()
26    }
27}
28
29/// Shared reference to any Interactive Air.
30/// This type is the main interface for keygen.
31pub type AirRef<SC> = Arc<dyn AnyAir<SC>>;
32
33/// RAP trait for all-purpose dynamic dispatch use.
34/// This trait is auto-implemented if you implement `Air` and `BaseAirWithPublicValues` and
35/// `PartitionedBaseAir` traits.
36pub trait AnyAir<SC: StarkProtocolConfig>:
37Air<SymbolicRapBuilder<SC::F>> // for keygen to extract fixed data about the RAP
38    + for<'a> Air<DebugConstraintBuilder<'a, SC>> // for debugging
39    + BaseAirWithPublicValues<SC::F>
40    + PartitionedBaseAir<SC::F>
41    + Send + Sync
42{
43    fn as_any(&self) -> &dyn Any;
44    /// Name for display purposes
45    fn name(&self) -> String;
46}
47
48impl<SC, T> AnyAir<SC> for T
49where
50    SC: StarkProtocolConfig,
51    T: Air<SymbolicRapBuilder<SC::F>>
52        + for<'a> Air<DebugConstraintBuilder<'a, SC>>
53        + BaseAirWithPublicValues<SC::F>
54        + PartitionedBaseAir<SC::F>
55        + Send
56        + Sync
57        + 'static,
58{
59    fn as_any(&self) -> &dyn Any {
60        self
61    }
62
63    fn name(&self) -> String {
64        get_air_name(self)
65    }
66}
67
68/// Automatically derives the AIR name from the type name for pretty display purposes.
69pub fn get_air_name<T>(_rap: &T) -> String {
70    let full_name = type_name::<T>().to_string();
71    // Split the input by the first '<' to separate the main type from its generics
72    if let Some((main_part, generics_part)) = full_name.split_once('<') {
73        // Extract the last segment of the main type
74        let main_type = main_part.split("::").last().unwrap_or("");
75
76        // Remove the trailing '>' from the generics part and split by ", " to handle multiple
77        // generics
78        let generics: Vec<String> = generics_part
79            .trim_end_matches('>')
80            .split(", ")
81            .map(|generic| {
82                // For each generic type, extract the last segment after "::"
83                generic.split("::").last().unwrap_or("").to_string()
84            })
85            .collect();
86
87        // Join the simplified generics back together with ", " and format the result
88        format!("{}<{}>", main_type, generics.join(", "))
89    } else {
90        // If there's no generic part, just return the last segment after "::"
91        full_name.split("::").last().unwrap_or("").to_string()
92    }
93}