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#[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 #[subenum(CoreAirId)]
40 Cpu = 0,
41 Program = 1,
43 ShaExtend = 2,
45 ShaExtendControl = 3,
47 ShaCompress = 4,
49 ShaCompressControl = 5,
51 EdAddAssign = 6,
53 EdDecompress = 7,
55 Secp256k1Decompress = 8,
57 Secp256k1AddAssign = 9,
59 Secp256k1DoubleAssign = 10,
61 Secp256r1Decompress = 11,
63 Secp256r1AddAssign = 12,
65 Secp256r1DoubleAssign = 13,
67 KeccakPermute = 14,
69 KeccakPermuteControl = 15,
71 Bn254AddAssign = 16,
73 Bn254DoubleAssign = 17,
75 Bls12381AddAssign = 18,
77 Bls12381DoubleAssign = 19,
79 Uint256MulMod = 20,
81 Uint256Ops = 21,
83 U256XU2048Mul = 22,
85 Bls12381FpOpAssign = 23,
87 Bls12381Fp2AddSubAssign = 24,
89 Bls12381Fp2MulAssign = 25,
91 Bn254FpOpAssign = 26,
93 Bn254Fp2AddSubAssign = 27,
95 Bn254Fp2MulAssign = 28,
97 Bls12381Decompress = 29,
99 #[subenum(CoreAirId)]
101 SyscallCore = 30,
102 SyscallPrecompile = 31,
104 #[subenum(CoreAirId)]
106 DivRem = 32,
107 #[subenum(CoreAirId)]
109 Add = 33,
110 #[subenum(CoreAirId)]
112 Addi = 34,
113 #[subenum(CoreAirId)]
115 Addw = 35,
116 #[subenum(CoreAirId)]
118 Sub = 36,
119 #[subenum(CoreAirId)]
121 Subw = 37,
122 #[subenum(CoreAirId)]
124 Bitwise = 38,
125 #[subenum(CoreAirId)]
127 Mul = 39,
128 #[subenum(CoreAirId)]
130 ShiftRight = 40,
131 #[subenum(CoreAirId)]
133 ShiftLeft = 41,
134 #[subenum(CoreAirId)]
136 Lt = 42,
137 #[subenum(CoreAirId)]
139 LoadByte = 43,
140 #[subenum(CoreAirId)]
142 LoadHalf = 44,
143 #[subenum(CoreAirId)]
145 LoadWord = 45,
146 #[subenum(CoreAirId)]
148 LoadX0 = 46,
149 #[subenum(CoreAirId)]
151 LoadDouble = 47,
152 #[subenum(CoreAirId)]
154 StoreByte = 48,
155 #[subenum(CoreAirId)]
157 StoreHalf = 49,
158 #[subenum(CoreAirId)]
160 StoreWord = 50,
161 #[subenum(CoreAirId)]
163 StoreDouble = 51,
164 #[subenum(CoreAirId)]
166 UType = 52,
167 #[subenum(CoreAirId)]
169 Branch = 53,
170 #[subenum(CoreAirId)]
172 Jal = 54,
173 #[subenum(CoreAirId)]
175 Jalr = 55,
176 #[subenum(CoreAirId)]
178 SyscallInstrs = 56,
179 #[subenum(CoreAirId)]
181 MemoryBump = 57,
182 #[subenum(CoreAirId)]
184 StateBump = 58,
185 MemoryGlobalInit = 59,
187 MemoryGlobalFinalize = 60,
189 #[subenum(CoreAirId)]
191 MemoryLocal = 61,
192 #[subenum(CoreAirId)]
194 Global = 62,
195 Byte = 63,
197 Range = 64,
199 #[subenum(CoreAirId)]
201 Mprotect = 65,
202 #[subenum(CoreAirId)]
204 InstructionDecode = 66,
205 #[subenum(CoreAirId)]
207 InstructionFetch = 67,
208 #[subenum(CoreAirId)]
210 PageProt = 68,
211 #[subenum(CoreAirId)]
213 PageProtLocal = 69,
214 PageProtGlobalInit = 70,
216 PageProtGlobalFinalize = 71,
218 Poseidon2 = 72,
220}
221
222impl RiscvAirId {
223 #[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 #[must_use]
268 pub fn is_core(self) -> bool {
269 CoreAirId::try_from(self).is_ok()
270 }
271
272 #[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 #[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 #[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 #[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 #[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#[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 pub fn iter(&self) -> impl Iterator<Item = &EnumMap<CoreAirId, u32>> {
396 self.inner.iter()
397 }
398}