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