sp1_core_executor/
air.rs

1use std::{
2    fmt::{Display, Formatter, Result as FmtResult},
3    str::FromStr,
4};
5
6use enum_map::{Enum, EnumMap};
7use serde::{Deserialize, Serialize};
8use sp1_stark::shape::Shape;
9use strum::{EnumIter, IntoEnumIterator, IntoStaticStr};
10use subenum::subenum;
11
12/// RV32IM AIR Identifiers.
13///
14/// These identifiers are for the various chips in the rv32im prover. We need them in the
15/// executor to compute the memory cost of the current shard of execution.
16///
17/// The [`CoreAirId`]s are the AIRs that are not part of precompile shards and not the program or
18/// byte AIR.
19#[subenum(CoreAirId)]
20#[derive(
21    Debug,
22    Clone,
23    Copy,
24    PartialEq,
25    Eq,
26    Hash,
27    Serialize,
28    Deserialize,
29    EnumIter,
30    IntoStaticStr,
31    PartialOrd,
32    Ord,
33    Enum,
34)]
35pub enum RiscvAirId {
36    /// The CPU chip.
37    #[subenum(CoreAirId)]
38    Cpu = 0,
39    /// The program chip.
40    Program = 1,
41    /// The SHA-256 extend chip.
42    ShaExtend = 2,
43    /// The SHA-256 compress chip.
44    ShaCompress = 3,
45    /// The Edwards add assign chip.
46    EdAddAssign = 4,
47    /// The Edwards decompress chip.
48    EdDecompress = 5,
49    /// The secp256k1 decompress chip.
50    Secp256k1Decompress = 6,
51    /// The secp256k1 add assign chip.
52    Secp256k1AddAssign = 7,
53    /// The secp256k1 double assign chip.
54    Secp256k1DoubleAssign = 8,
55    /// The secp256r1 decompress chip.
56    Secp256r1Decompress = 9,
57    /// The secp256r1 add assign chip.
58    Secp256r1AddAssign = 10,
59    /// The secp256r1 double assign chip.
60    Secp256r1DoubleAssign = 11,
61    /// The Keccak permute chip.
62    KeccakPermute = 12,
63    /// The bn254 add assign chip.
64    Bn254AddAssign = 13,
65    /// The bn254 double assign chip.
66    Bn254DoubleAssign = 14,
67    /// The bls12-381 add assign chip.
68    Bls12381AddAssign = 15,
69    /// The bls12-381 double assign chip.
70    Bls12381DoubleAssign = 16,
71    /// The uint256 mul mod chip.
72    Uint256MulMod = 17,
73    /// The u256 xu2048 mul chip.
74    U256XU2048Mul = 18,
75    /// The bls12-381 fp op assign chip.
76    Bls12381FpOpAssign = 19,
77    /// The bls12-831 fp2 add sub assign chip.
78    Bls12381Fp2AddSubAssign = 20,
79    /// The bls12-831 fp2 mul assign chip.
80    Bls12381Fp2MulAssign = 21,
81    /// The bn254 fp2 add sub assign chip.
82    Bn254FpOpAssign = 22,
83    /// The bn254 fp op assign chip.
84    Bn254Fp2AddSubAssign = 23,
85    /// The bn254 fp2 mul assign chip.
86    Bn254Fp2MulAssign = 24,
87    /// The bls12-381 decompress chip.
88    Bls12381Decompress = 25,
89    /// The syscall core chip.
90    #[subenum(CoreAirId)]
91    SyscallCore = 26,
92    /// The syscall precompile chip.
93    SyscallPrecompile = 27,
94    /// The div rem chip.
95    #[subenum(CoreAirId)]
96    DivRem = 28,
97    /// The add sub chip.
98    #[subenum(CoreAirId)]
99    AddSub = 29,
100    /// The bitwise chip.
101    #[subenum(CoreAirId)]
102    Bitwise = 30,
103    /// The mul chip.
104    #[subenum(CoreAirId)]
105    Mul = 31,
106    /// The shift right chip.
107    #[subenum(CoreAirId)]
108    ShiftRight = 32,
109    /// The shift left chip.
110    #[subenum(CoreAirId)]
111    ShiftLeft = 33,
112    /// The lt chip.
113    #[subenum(CoreAirId)]
114    Lt = 34,
115    /// The memory instructions chip.
116    #[subenum(CoreAirId)]
117    MemoryInstrs = 35,
118    /// The auipc chip.
119    #[subenum(CoreAirId)]
120    Auipc = 36,
121    /// The branch chip.
122    #[subenum(CoreAirId)]
123    Branch = 37,
124    /// The jump chip.
125    #[subenum(CoreAirId)]
126    Jump = 38,
127    /// The syscall instructions chip.
128    #[subenum(CoreAirId)]
129    SyscallInstrs = 39,
130    /// The memory global init chip.
131    MemoryGlobalInit = 40,
132    /// The memory global finalize chip.
133    MemoryGlobalFinalize = 41,
134    /// The memory local chip.
135    #[subenum(CoreAirId)]
136    MemoryLocal = 42,
137    /// The global chip.
138    #[subenum(CoreAirId)]
139    Global = 43,
140    /// The byte chip.
141    Byte = 44,
142}
143
144impl RiscvAirId {
145    /// Returns the AIRs that are not part of precompile shards and not the program or byte AIR.
146    #[must_use]
147    pub fn core() -> Vec<RiscvAirId> {
148        vec![
149            RiscvAirId::Cpu,
150            RiscvAirId::AddSub,
151            RiscvAirId::Mul,
152            RiscvAirId::Bitwise,
153            RiscvAirId::ShiftLeft,
154            RiscvAirId::ShiftRight,
155            RiscvAirId::DivRem,
156            RiscvAirId::Lt,
157            RiscvAirId::Auipc,
158            RiscvAirId::MemoryLocal,
159            RiscvAirId::MemoryInstrs,
160            RiscvAirId::Branch,
161            RiscvAirId::Jump,
162            RiscvAirId::SyscallCore,
163            RiscvAirId::SyscallInstrs,
164            RiscvAirId::Global,
165        ]
166    }
167
168    /// TODO replace these three with subenums or something
169    /// Whether the ID represents a core AIR.
170    #[must_use]
171    pub fn is_core(self) -> bool {
172        CoreAirId::try_from(self).is_ok()
173    }
174
175    /// Whether the ID represents a memory AIR.
176    #[must_use]
177    pub fn is_memory(self) -> bool {
178        matches!(
179            self,
180            RiscvAirId::MemoryGlobalInit | RiscvAirId::MemoryGlobalFinalize | RiscvAirId::Global
181        )
182    }
183
184    /// Whether the ID represents a precompile AIR.
185    #[must_use]
186    pub fn is_precompile(self) -> bool {
187        matches!(
188            self,
189            RiscvAirId::ShaExtend |
190                RiscvAirId::ShaCompress |
191                RiscvAirId::EdAddAssign |
192                RiscvAirId::EdDecompress |
193                RiscvAirId::Secp256k1Decompress |
194                RiscvAirId::Secp256k1AddAssign |
195                RiscvAirId::Secp256k1DoubleAssign |
196                RiscvAirId::Secp256r1Decompress |
197                RiscvAirId::Secp256r1AddAssign |
198                RiscvAirId::Secp256r1DoubleAssign |
199                RiscvAirId::KeccakPermute |
200                RiscvAirId::Bn254AddAssign |
201                RiscvAirId::Bn254DoubleAssign |
202                RiscvAirId::Bls12381AddAssign |
203                RiscvAirId::Bls12381DoubleAssign |
204                RiscvAirId::Uint256MulMod |
205                RiscvAirId::U256XU2048Mul |
206                RiscvAirId::Bls12381FpOpAssign |
207                RiscvAirId::Bls12381Fp2AddSubAssign |
208                RiscvAirId::Bls12381Fp2MulAssign |
209                RiscvAirId::Bn254FpOpAssign |
210                RiscvAirId::Bn254Fp2AddSubAssign |
211                RiscvAirId::Bn254Fp2MulAssign |
212                RiscvAirId::Bls12381Decompress
213        )
214    }
215
216    /// The number of rows in the AIR produced by each event.
217    #[must_use]
218    pub fn rows_per_event(&self) -> usize {
219        match self {
220            Self::ShaCompress => 80,
221            Self::ShaExtend => 48,
222            Self::KeccakPermute => 24,
223            _ => 1,
224        }
225    }
226
227    /// Returns the string representation of the AIR.
228    #[must_use]
229    pub fn as_str(&self) -> &'static str {
230        self.into()
231    }
232}
233
234impl FromStr for RiscvAirId {
235    type Err = String;
236
237    fn from_str(s: &str) -> Result<Self, Self::Err> {
238        let air = Self::iter().find(|chip| chip.as_str() == s);
239        match air {
240            Some(air) => Ok(air),
241            None => Err(format!("Invalid RV32IMAir: {s}")),
242        }
243    }
244}
245
246impl Display for RiscvAirId {
247    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
248        write!(f, "{}", self.as_str())
249    }
250}
251
252/// Defines a set of maximal shapes for generating core proofs.
253#[derive(Debug, Clone, Serialize, Deserialize)]
254pub struct MaximalShapes {
255    inner: Vec<EnumMap<CoreAirId, u32>>,
256}
257
258impl FromIterator<Shape<RiscvAirId>> for MaximalShapes {
259    fn from_iter<T: IntoIterator<Item = Shape<RiscvAirId>>>(iter: T) -> Self {
260        let mut maximal_shapes = Vec::new();
261        for shape in iter {
262            let mut maximal_shape = EnumMap::<CoreAirId, u32>::default();
263            for (air, height) in shape {
264                if let Ok(core_air) = CoreAirId::try_from(air) {
265                    maximal_shape[core_air] = height as u32;
266                } else if air != RiscvAirId::Program && air != RiscvAirId::Byte {
267                    tracing::warn!("Invalid core air: {air}");
268                }
269            }
270            maximal_shapes.push(maximal_shape);
271        }
272        Self { inner: maximal_shapes }
273    }
274}
275
276impl MaximalShapes {
277    /// Returns an iterator over the maximal shapes.
278    pub fn iter(&self) -> impl Iterator<Item = &EnumMap<CoreAirId, u32>> {
279        self.inner.iter()
280    }
281}