Skip to main content

sp1_core_executor/
air.rs

1use std::{
2    fmt::{Display, Formatter, Result as FmtResult},
3    str::FromStr,
4};
5
6use deepsize2::DeepSizeOf;
7use enum_map::{Enum, EnumMap};
8use serde::{Deserialize, Serialize};
9use sp1_hypercube::shape::Shape;
10use strum::{EnumIter, IntoEnumIterator, IntoStaticStr};
11use subenum::subenum;
12
13/// RV64IM AIR Identifiers.
14///
15/// These identifiers are for the various chips in the rv64im prover. We need them in the
16/// executor to compute the memory cost of the current shard of execution.
17///
18/// The [`CoreAirId`]s are the AIRs that are not part of precompile shards and not the program or
19/// byte AIR.
20#[subenum(CoreAirId)]
21#[derive(
22    Debug,
23    Clone,
24    Copy,
25    PartialEq,
26    Eq,
27    Hash,
28    Serialize,
29    Deserialize,
30    EnumIter,
31    IntoStaticStr,
32    PartialOrd,
33    Ord,
34    Enum,
35    DeepSizeOf,
36)]
37pub enum RiscvAirId {
38    /// The cpu chip, which is a dummy for now, needed for shape loading.
39    #[subenum(CoreAirId)]
40    Cpu = 0,
41    /// The program chip.
42    Program = 1,
43    /// The SHA-256 extend chip.
44    ShaExtend = 2,
45    /// The sha extend control chip.
46    ShaExtendControl = 3,
47    /// The SHA-256 compress chip.
48    ShaCompress = 4,
49    /// The sha compress control chip.
50    ShaCompressControl = 5,
51    /// The Edwards add assign chip.
52    EdAddAssign = 6,
53    /// The Edwards decompress chip.
54    EdDecompress = 7,
55    /// The secp256k1 decompress chip.
56    Secp256k1Decompress = 8,
57    /// The secp256k1 add assign chip.
58    Secp256k1AddAssign = 9,
59    /// The secp256k1 double assign chip.
60    Secp256k1DoubleAssign = 10,
61    /// The secp256r1 decompress chip.
62    Secp256r1Decompress = 11,
63    /// The secp256r1 add assign chip.
64    Secp256r1AddAssign = 12,
65    /// The secp256r1 double assign chip.
66    Secp256r1DoubleAssign = 13,
67    /// The Keccak permute chip.
68    KeccakPermute = 14,
69    /// The keccak permute control chip.
70    KeccakPermuteControl = 15,
71    /// The bn254 add assign chip.
72    Bn254AddAssign = 16,
73    /// The bn254 double assign chip.
74    Bn254DoubleAssign = 17,
75    /// The bls12-381 add assign chip.
76    Bls12381AddAssign = 18,
77    /// The bls12-381 double assign chip.
78    Bls12381DoubleAssign = 19,
79    /// The uint256 mul mod chip.
80    Uint256MulMod = 20,
81    /// The uint256 ops chip.
82    Uint256Ops = 21,
83    /// The u256 xu2048 mul chip.
84    U256XU2048Mul = 22,
85    /// The bls12-381 fp op assign chip.
86    Bls12381FpOpAssign = 23,
87    /// The bls12-831 fp2 add sub assign chip.
88    Bls12381Fp2AddSubAssign = 24,
89    /// The bls12-831 fp2 mul assign chip.
90    Bls12381Fp2MulAssign = 25,
91    /// The bn254 fp2 add sub assign chip.
92    Bn254FpOpAssign = 26,
93    /// The bn254 fp op assign chip.
94    Bn254Fp2AddSubAssign = 27,
95    /// The bn254 fp2 mul assign chip.
96    Bn254Fp2MulAssign = 28,
97    /// The bls12-381 decompress chip.
98    Bls12381Decompress = 29,
99    /// The syscall core chip.
100    #[subenum(CoreAirId)]
101    SyscallCore = 30,
102    /// The syscall precompile chip.
103    SyscallPrecompile = 31,
104    /// The div rem chip.
105    #[subenum(CoreAirId)]
106    DivRem = 32,
107    /// The add chip.
108    #[subenum(CoreAirId)]
109    Add = 33,
110    /// The addi chip.
111    #[subenum(CoreAirId)]
112    Addi = 34,
113    /// The addw chip.
114    #[subenum(CoreAirId)]
115    Addw = 35,
116    /// The sub chip.
117    #[subenum(CoreAirId)]
118    Sub = 36,
119    /// The subw chip.
120    #[subenum(CoreAirId)]
121    Subw = 37,
122    /// The bitwise chip.
123    #[subenum(CoreAirId)]
124    Bitwise = 38,
125    /// The mul chip.
126    #[subenum(CoreAirId)]
127    Mul = 39,
128    /// The shift right chip.
129    #[subenum(CoreAirId)]
130    ShiftRight = 40,
131    /// The shift left chip.
132    #[subenum(CoreAirId)]
133    ShiftLeft = 41,
134    /// The lt chip.
135    #[subenum(CoreAirId)]
136    Lt = 42,
137    /// The load byte chip.
138    #[subenum(CoreAirId)]
139    LoadByte = 43,
140    /// The load half chip.
141    #[subenum(CoreAirId)]
142    LoadHalf = 44,
143    /// The load word chip.
144    #[subenum(CoreAirId)]
145    LoadWord = 45,
146    /// The load x0 chip.
147    #[subenum(CoreAirId)]
148    LoadX0 = 46,
149    /// The load double chip.
150    #[subenum(CoreAirId)]
151    LoadDouble = 47,
152    /// The store byte chip.
153    #[subenum(CoreAirId)]
154    StoreByte = 48,
155    /// The store half chip.
156    #[subenum(CoreAirId)]
157    StoreHalf = 49,
158    /// The store word chip.
159    #[subenum(CoreAirId)]
160    StoreWord = 50,
161    /// The store double chip.
162    #[subenum(CoreAirId)]
163    StoreDouble = 51,
164    /// The utype chip.
165    #[subenum(CoreAirId)]
166    UType = 52,
167    /// The branch chip.
168    #[subenum(CoreAirId)]
169    Branch = 53,
170    /// The jal chip.
171    #[subenum(CoreAirId)]
172    Jal = 54,
173    /// The jalr chip.
174    #[subenum(CoreAirId)]
175    Jalr = 55,
176    /// The syscall instructions chip.
177    #[subenum(CoreAirId)]
178    SyscallInstrs = 56,
179    /// The memory bump chip.
180    #[subenum(CoreAirId)]
181    MemoryBump = 57,
182    /// The state bump chip.
183    #[subenum(CoreAirId)]
184    StateBump = 58,
185    /// The memory global init chip.
186    MemoryGlobalInit = 59,
187    /// The memory global finalize chip.
188    MemoryGlobalFinalize = 60,
189    /// The memory local chip.
190    #[subenum(CoreAirId)]
191    MemoryLocal = 61,
192    /// The global chip.
193    #[subenum(CoreAirId)]
194    Global = 62,
195    /// The byte chip.
196    Byte = 63,
197    /// The range chip.
198    Range = 64,
199    /// The mprotect chip.
200    #[subenum(CoreAirId)]
201    Mprotect = 65,
202    /// The instruction decode chip.
203    #[subenum(CoreAirId)]
204    InstructionDecode = 66,
205    /// The instruction fetch chip.
206    #[subenum(CoreAirId)]
207    InstructionFetch = 67,
208    /// The page prot chip.
209    #[subenum(CoreAirId)]
210    PageProt = 68,
211    /// The page prot local chip.
212    #[subenum(CoreAirId)]
213    PageProtLocal = 69,
214    /// The page prot global init chip.
215    PageProtGlobalInit = 70,
216    /// The page prot global finalize chip.
217    PageProtGlobalFinalize = 71,
218    /// The poseidon2 chip.
219    Poseidon2 = 72,
220}
221
222impl RiscvAirId {
223    /// Returns the AIRs that are not part of precompile shards and not the program or byte AIR.
224    #[must_use]
225    pub fn core() -> Vec<RiscvAirId> {
226        vec![
227            RiscvAirId::Add,
228            RiscvAirId::Addi,
229            RiscvAirId::Addw,
230            RiscvAirId::Sub,
231            RiscvAirId::Subw,
232            RiscvAirId::Mul,
233            RiscvAirId::Bitwise,
234            RiscvAirId::ShiftLeft,
235            RiscvAirId::ShiftRight,
236            RiscvAirId::DivRem,
237            RiscvAirId::Lt,
238            RiscvAirId::UType,
239            RiscvAirId::MemoryLocal,
240            RiscvAirId::MemoryBump,
241            RiscvAirId::StateBump,
242            RiscvAirId::LoadByte,
243            RiscvAirId::LoadHalf,
244            RiscvAirId::LoadWord,
245            RiscvAirId::LoadDouble,
246            RiscvAirId::LoadX0,
247            RiscvAirId::StoreByte,
248            RiscvAirId::StoreHalf,
249            RiscvAirId::StoreWord,
250            RiscvAirId::StoreDouble,
251            RiscvAirId::Branch,
252            RiscvAirId::Jal,
253            RiscvAirId::Jalr,
254            RiscvAirId::PageProt,
255            RiscvAirId::PageProtLocal,
256            RiscvAirId::SyscallCore,
257            RiscvAirId::SyscallInstrs,
258            RiscvAirId::Global,
259            RiscvAirId::Mprotect,
260            RiscvAirId::InstructionDecode,
261            RiscvAirId::InstructionFetch,
262        ]
263    }
264
265    /// TODO replace these three with subenums or something
266    /// Whether the ID represents a core AIR.
267    #[must_use]
268    pub fn is_core(self) -> bool {
269        CoreAirId::try_from(self).is_ok()
270    }
271
272    /// Whether the ID represents a memory AIR.
273    #[must_use]
274    pub fn is_memory(self) -> bool {
275        matches!(
276            self,
277            RiscvAirId::MemoryGlobalInit
278                | RiscvAirId::MemoryGlobalFinalize
279                | RiscvAirId::Global
280                | RiscvAirId::PageProtGlobalInit
281                | RiscvAirId::PageProtGlobalFinalize
282        )
283    }
284
285    /// Whether the ID represents a precompile AIR.
286    #[must_use]
287    pub fn is_precompile(self) -> bool {
288        matches!(
289            self,
290            RiscvAirId::ShaExtend
291                | RiscvAirId::ShaCompress
292                | RiscvAirId::EdAddAssign
293                | RiscvAirId::EdDecompress
294                | RiscvAirId::Secp256k1Decompress
295                | RiscvAirId::Secp256k1AddAssign
296                | RiscvAirId::Secp256k1DoubleAssign
297                | RiscvAirId::Secp256r1Decompress
298                | RiscvAirId::Secp256r1AddAssign
299                | RiscvAirId::Secp256r1DoubleAssign
300                | RiscvAirId::KeccakPermute
301                | RiscvAirId::Bn254AddAssign
302                | RiscvAirId::Bn254DoubleAssign
303                | RiscvAirId::Bls12381AddAssign
304                | RiscvAirId::Bls12381DoubleAssign
305                | RiscvAirId::Uint256MulMod
306                | RiscvAirId::Uint256Ops
307                | RiscvAirId::U256XU2048Mul
308                | RiscvAirId::Bls12381FpOpAssign
309                | RiscvAirId::Bls12381Fp2AddSubAssign
310                | RiscvAirId::Bls12381Fp2MulAssign
311                | RiscvAirId::Bn254FpOpAssign
312                | RiscvAirId::Bn254Fp2AddSubAssign
313                | RiscvAirId::Bn254Fp2MulAssign
314                | RiscvAirId::Bls12381Decompress
315                | RiscvAirId::Poseidon2
316        )
317    }
318
319    /// The number of rows in the AIR produced by each event.
320    #[must_use]
321    pub fn rows_per_event(&self) -> usize {
322        match self {
323            Self::ShaCompress => 80,
324            Self::ShaExtend => 48,
325            Self::KeccakPermute => 24,
326            _ => 1,
327        }
328    }
329
330    /// Get the ID of the AIR used in the syscall control implementation.
331    #[must_use]
332    pub fn control_air_id(self) -> Option<RiscvAirId> {
333        match self {
334            RiscvAirId::ShaCompress => Some(RiscvAirId::ShaCompressControl),
335            RiscvAirId::ShaExtend => Some(RiscvAirId::ShaExtendControl),
336            RiscvAirId::KeccakPermute => Some(RiscvAirId::KeccakPermuteControl),
337            _ => None,
338        }
339    }
340
341    /// Returns the string representation of the AIR.
342    #[must_use]
343    pub fn as_str(&self) -> &'static str {
344        self.into()
345    }
346}
347
348impl FromStr for RiscvAirId {
349    type Err = String;
350
351    fn from_str(s: &str) -> Result<Self, Self::Err> {
352        let air = Self::iter().find(|chip| chip.as_str() == s);
353        match air {
354            Some(air) => Ok(air),
355            None => Err(format!("Invalid RV64IMAir: {s}")),
356        }
357    }
358}
359
360impl Display for RiscvAirId {
361    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
362        write!(f, "{}", self.as_str())
363    }
364}
365
366/// Defines a set of maximal shapes for generating core proofs.
367#[derive(Debug, Clone, Serialize, Deserialize)]
368pub struct MaximalShapes {
369    inner: Vec<EnumMap<CoreAirId, u32>>,
370}
371
372impl FromIterator<Shape<RiscvAirId>> for MaximalShapes {
373    fn from_iter<T: IntoIterator<Item = Shape<RiscvAirId>>>(iter: T) -> Self {
374        let mut maximal_shapes = Vec::new();
375        for shape in iter {
376            let mut maximal_shape = EnumMap::<CoreAirId, u32>::default();
377            for (air, height) in shape {
378                if let Ok(core_air) = CoreAirId::try_from(air) {
379                    maximal_shape[core_air] = height as u32;
380                } else if air != RiscvAirId::Program
381                    && air != RiscvAirId::Byte
382                    && air != RiscvAirId::Range
383                {
384                    tracing::warn!("Invalid core air: {air}");
385                }
386            }
387            maximal_shapes.push(maximal_shape);
388        }
389        Self { inner: maximal_shapes }
390    }
391}
392
393impl MaximalShapes {
394    /// Returns an iterator over the maximal shapes.
395    pub fn iter(&self) -> impl Iterator<Item = &EnumMap<CoreAirId, u32>> {
396        self.inner.iter()
397    }
398}