openvm_stark_backend/
rap.rs

1//! # RAP (Randomized Air with Preprocessing)
2//! See <https://hackmd.io/@aztec-network/plonk-arithmetiization-air> for formal definition.
3
4use std::{
5    any::{type_name, Any},
6    sync::Arc,
7};
8
9use p3_air::{BaseAir, PermutationAirBuilder};
10
11use crate::{
12    air_builders::{debug::DebugConstraintBuilder, symbolic::SymbolicRapBuilder},
13    config::{StarkGenericConfig, Val},
14};
15
16/// An AIR with 0 or more public values.
17/// This trait will be merged into Plonky3 in PR: <https://github.com/Plonky3/Plonky3/pull/470>
18pub trait BaseAirWithPublicValues<F>: BaseAir<F> {
19    fn num_public_values(&self) -> usize {
20        0
21    }
22}
23
24/// An AIR with 1 or more main trace partitions.
25pub trait PartitionedBaseAir<F>: BaseAir<F> {
26    /// By default, an AIR has no cached main trace.
27    fn cached_main_widths(&self) -> Vec<usize> {
28        vec![]
29    }
30    /// By default, an AIR has only one private main trace.
31    fn common_main_width(&self) -> usize {
32        self.width()
33    }
34}
35
36/// An AIR that works with a particular `AirBuilder` which allows preprocessing
37/// and injected randomness.
38///
39/// Currently this is not a fully general RAP. Only the following phases are allowed:
40/// - Preprocessing
41/// - Main trace generation and commitment
42/// - Permutation trace generation and commitment
43///
44/// Randomness is drawn after the main trace commitment phase, and used in the permutation trace.
45///
46/// Does not inherit [Air](p3_air::Air) trait to allow overrides for technical reasons
47/// around dynamic dispatch.
48pub trait Rap<AB>: Sync
49where
50    AB: PermutationAirBuilder,
51{
52    fn eval(&self, builder: &mut AB);
53}
54
55/// Permutation AIR builder that exposes certain values to both prover and verifier
56/// _after_ the permutation challenges are drawn. These can be thought of as
57/// "public values" known after the challenges are drawn.
58///
59/// Exposed values are used internally by the prover and verifier
60/// in cross-table permutation arguments.
61pub trait PermutationAirBuilderWithExposedValues: PermutationAirBuilder {
62    fn permutation_exposed_values(&self) -> &[Self::VarEF];
63}
64
65/// Shared reference to any Interactive Air.
66/// This type is the main interface for keygen.
67pub type AirRef<SC> = Arc<dyn AnyRap<SC>>;
68
69/// RAP trait for all-purpose dynamic dispatch use.
70/// This trait is auto-implemented if you implement `Air` and `BaseAirWithPublicValues` and
71/// `PartitionedBaseAir` traits.
72pub trait AnyRap<SC: StarkGenericConfig>:
73Rap<SymbolicRapBuilder<Val<SC>>> // for keygen to extract fixed data about the RAP
74    + for<'a> Rap<DebugConstraintBuilder<'a, SC>> // for debugging
75    + BaseAirWithPublicValues<Val<SC>>
76    + PartitionedBaseAir<Val<SC>>
77    + Send + Sync
78{
79    fn as_any(&self) -> &dyn Any;
80    /// Name for display purposes
81    fn name(&self) -> String;
82}
83
84impl<SC, T> AnyRap<SC> for T
85where
86    SC: StarkGenericConfig,
87    T: Rap<SymbolicRapBuilder<Val<SC>>>
88        + for<'a> Rap<DebugConstraintBuilder<'a, SC>>
89        + BaseAirWithPublicValues<Val<SC>>
90        + PartitionedBaseAir<Val<SC>>
91        + Send
92        + Sync
93        + 'static,
94{
95    fn as_any(&self) -> &dyn Any {
96        self
97    }
98
99    fn name(&self) -> String {
100        get_air_name(self)
101    }
102}
103
104/// Automatically derives the AIR name from the type name for pretty display purposes.
105pub fn get_air_name<T>(_rap: &T) -> String {
106    let full_name = type_name::<T>().to_string();
107    // Split the input by the first '<' to separate the main type from its generics
108    if let Some((main_part, generics_part)) = full_name.split_once('<') {
109        // Extract the last segment of the main type
110        let main_type = main_part.split("::").last().unwrap_or("");
111
112        // Remove the trailing '>' from the generics part and split by ", " to handle multiple
113        // generics
114        let generics: Vec<String> = generics_part
115            .trim_end_matches('>')
116            .split(", ")
117            .map(|generic| {
118                // For each generic type, extract the last segment after "::"
119                generic.split("::").last().unwrap_or("").to_string()
120            })
121            .collect();
122
123        // Join the simplified generics back together with ", " and format the result
124        format!("{}<{}>", main_type, generics.join(", "))
125    } else {
126        // If there's no generic part, just return the last segment after "::"
127        full_name.split("::").last().unwrap_or("").to_string()
128    }
129}