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 Program = 0,
40 ShaExtend = 1,
42 ShaExtendControl = 2,
44 ShaExtendControlUser = 3,
46 ShaCompress = 4,
48 ShaCompressControl = 5,
50 ShaCompressControlUser = 6,
52 EdAddAssign = 7,
54 EdAddAssignUser = 8,
56 EdDecompress = 9,
58 EdDecompressUser = 10,
60 Secp256k1AddAssign = 11,
62 Secp256k1AddAssignUser = 12,
64 Secp256k1DoubleAssign = 13,
66 Secp256k1DoubleAssignUser = 14,
68 Secp256r1AddAssign = 15,
70 Secp256r1AddAssignUser = 16,
72 Secp256r1DoubleAssign = 17,
74 Secp256r1DoubleAssignUser = 18,
76 KeccakPermute = 19,
78 KeccakPermuteControl = 20,
80 KeccakPermuteControlUser = 21,
82 Bn254AddAssign = 22,
84 Bn254AddAssignUser = 23,
86 Bn254DoubleAssign = 24,
88 Bn254DoubleAssignUser = 25,
90 Bls12381AddAssign = 26,
92 Bls12381AddAssignUser = 27,
94 Bls12381DoubleAssign = 28,
96 Bls12381DoubleAssignUser = 29,
98 Uint256MulMod = 30,
100 Uint256MulModUser = 31,
102 Uint256Ops = 32,
104 Uint256OpsUser = 33,
106 Bls12381FpOpAssign = 34,
108 Bls12381FpOpAssignUser = 35,
110 Bls12381Fp2AddSubAssign = 36,
112 Bls12381Fp2AddSubAssignUser = 37,
114 Bls12381Fp2MulAssign = 38,
116 Bls12381Fp2MulAssignUser = 39,
118 Bn254FpOpAssign = 40,
120 Bn254FpOpAssignUser = 41,
122 Bn254Fp2AddSubAssign = 42,
124 Bn254Fp2AddSubAssignUser = 43,
126 Bn254Fp2MulAssign = 44,
128 Bn254Fp2MulAssignUser = 45,
130 Poseidon2 = 46,
132 Poseidon2User = 47,
134 #[subenum(CoreAirId)]
136 SyscallCore = 48,
137 #[subenum(CoreAirId)]
139 SyscallCoreUser = 49,
140 SyscallPrecompile = 50,
142 SyscallPrecompileUser = 51,
144 #[subenum(CoreAirId)]
146 DivRem = 52,
147 #[subenum(CoreAirId)]
149 DivRemUser = 53,
150 #[subenum(CoreAirId)]
152 Add = 54,
153 #[subenum(CoreAirId)]
155 AddUser = 55,
156 #[subenum(CoreAirId)]
158 Addi = 56,
159 #[subenum(CoreAirId)]
161 AddiUser = 57,
162 #[subenum(CoreAirId)]
164 Addw = 58,
165 #[subenum(CoreAirId)]
167 AddwUser = 59,
168 #[subenum(CoreAirId)]
170 Sub = 60,
171 #[subenum(CoreAirId)]
173 SubUser = 61,
174 #[subenum(CoreAirId)]
176 Subw = 62,
177 #[subenum(CoreAirId)]
179 SubwUser = 63,
180 #[subenum(CoreAirId)]
182 Bitwise = 64,
183 #[subenum(CoreAirId)]
185 BitwiseUser = 65,
186 #[subenum(CoreAirId)]
188 Mul = 66,
189 #[subenum(CoreAirId)]
191 MulUser = 67,
192 #[subenum(CoreAirId)]
194 ShiftRight = 68,
195 #[subenum(CoreAirId)]
197 ShiftRightUser = 69,
198 #[subenum(CoreAirId)]
200 ShiftLeft = 70,
201 #[subenum(CoreAirId)]
203 ShiftLeftUser = 71,
204 #[subenum(CoreAirId)]
206 Lt = 72,
207 #[subenum(CoreAirId)]
209 LtUser = 73,
210 #[subenum(CoreAirId)]
212 LoadByte = 74,
213 #[subenum(CoreAirId)]
215 LoadByteUser = 75,
216 #[subenum(CoreAirId)]
218 LoadHalf = 76,
219 #[subenum(CoreAirId)]
221 LoadHalfUser = 77,
222 #[subenum(CoreAirId)]
224 LoadWord = 78,
225 #[subenum(CoreAirId)]
227 LoadWordUser = 79,
228 #[subenum(CoreAirId)]
230 LoadX0 = 80,
231 #[subenum(CoreAirId)]
233 LoadX0User = 81,
234 #[subenum(CoreAirId)]
236 LoadDouble = 82,
237 #[subenum(CoreAirId)]
239 LoadDoubleUser = 83,
240 #[subenum(CoreAirId)]
242 StoreByte = 84,
243 #[subenum(CoreAirId)]
245 StoreByteUser = 85,
246 #[subenum(CoreAirId)]
248 StoreHalf = 86,
249 #[subenum(CoreAirId)]
251 StoreHalfUser = 87,
252 #[subenum(CoreAirId)]
254 StoreWord = 88,
255 #[subenum(CoreAirId)]
257 StoreWordUser = 89,
258 #[subenum(CoreAirId)]
260 StoreDouble = 90,
261 #[subenum(CoreAirId)]
263 StoreDoubleUser = 91,
264 #[subenum(CoreAirId)]
266 UType = 92,
267 #[subenum(CoreAirId)]
269 UTypeUser = 93,
270 #[subenum(CoreAirId)]
272 Branch = 94,
273 #[subenum(CoreAirId)]
275 BranchUser = 95,
276 #[subenum(CoreAirId)]
278 Jal = 96,
279 #[subenum(CoreAirId)]
281 JalUser = 97,
282 #[subenum(CoreAirId)]
284 Jalr = 98,
285 #[subenum(CoreAirId)]
287 JalrUser = 99,
288 #[subenum(CoreAirId)]
290 SyscallInstrs = 100,
291 #[subenum(CoreAirId)]
293 SyscallInstrsUser = 101,
294 #[subenum(CoreAirId)]
296 MemoryBump = 102,
297 #[subenum(CoreAirId)]
299 StateBump = 103,
300 MemoryGlobalInit = 104,
302 MemoryGlobalFinalize = 105,
304 #[subenum(CoreAirId)]
306 MemoryLocal = 106,
307 #[subenum(CoreAirId)]
309 Global = 107,
310 Byte = 108,
312 Range = 109,
314 #[subenum(CoreAirId)]
316 AluX0 = 110,
317 #[subenum(CoreAirId)]
319 AluX0User = 111,
320 #[subenum(CoreAirId)]
322 Mprotect = 112,
323 SigReturn = 113,
325 #[subenum(CoreAirId)]
327 InstructionDecode = 114,
328 #[subenum(CoreAirId)]
330 InstructionFetch = 115,
331 #[subenum(CoreAirId)]
333 PageProt = 116,
334 #[subenum(CoreAirId)]
336 PageProtLocal = 117,
337 PageProtGlobalInit = 118,
339 PageProtGlobalFinalize = 119,
341 #[subenum(CoreAirId)]
343 TrapExec = 120,
344 #[subenum(CoreAirId)]
346 TrapMem = 121,
347}
348
349impl RiscvAirId {
350 #[must_use]
352 pub fn core() -> Vec<RiscvAirId> {
353 vec![
354 RiscvAirId::Add,
355 RiscvAirId::AddUser,
356 RiscvAirId::Addi,
357 RiscvAirId::AddiUser,
358 RiscvAirId::Addw,
359 RiscvAirId::AddwUser,
360 RiscvAirId::Sub,
361 RiscvAirId::SubUser,
362 RiscvAirId::Subw,
363 RiscvAirId::SubwUser,
364 RiscvAirId::Mul,
365 RiscvAirId::MulUser,
366 RiscvAirId::Bitwise,
367 RiscvAirId::BitwiseUser,
368 RiscvAirId::ShiftLeft,
369 RiscvAirId::ShiftLeftUser,
370 RiscvAirId::ShiftRight,
371 RiscvAirId::ShiftRightUser,
372 RiscvAirId::DivRem,
373 RiscvAirId::DivRemUser,
374 RiscvAirId::Lt,
375 RiscvAirId::LtUser,
376 RiscvAirId::UType,
377 RiscvAirId::UTypeUser,
378 RiscvAirId::MemoryLocal,
379 RiscvAirId::MemoryBump,
380 RiscvAirId::StateBump,
381 RiscvAirId::LoadByte,
382 RiscvAirId::LoadByteUser,
383 RiscvAirId::LoadHalf,
384 RiscvAirId::LoadHalfUser,
385 RiscvAirId::LoadWord,
386 RiscvAirId::LoadWordUser,
387 RiscvAirId::LoadDouble,
388 RiscvAirId::LoadDoubleUser,
389 RiscvAirId::LoadX0,
390 RiscvAirId::LoadX0User,
391 RiscvAirId::StoreByte,
392 RiscvAirId::StoreByteUser,
393 RiscvAirId::StoreHalf,
394 RiscvAirId::StoreHalfUser,
395 RiscvAirId::StoreWord,
396 RiscvAirId::StoreWordUser,
397 RiscvAirId::StoreDouble,
398 RiscvAirId::StoreDoubleUser,
399 RiscvAirId::Branch,
400 RiscvAirId::BranchUser,
401 RiscvAirId::Jal,
402 RiscvAirId::JalUser,
403 RiscvAirId::Jalr,
404 RiscvAirId::JalrUser,
405 RiscvAirId::SyscallCore,
406 RiscvAirId::SyscallInstrs,
407 RiscvAirId::SyscallInstrsUser,
408 RiscvAirId::Global,
409 RiscvAirId::AluX0,
410 RiscvAirId::AluX0User,
411 RiscvAirId::Mprotect,
412 RiscvAirId::InstructionDecode,
413 RiscvAirId::InstructionFetch,
414 RiscvAirId::PageProt,
415 RiscvAirId::PageProtLocal,
416 RiscvAirId::TrapExec,
417 RiscvAirId::TrapMem,
418 ]
419 }
420
421 #[must_use]
423 pub fn is_core(self) -> bool {
424 CoreAirId::try_from(self).is_ok()
425 }
426
427 #[must_use]
429 pub fn is_memory(self) -> bool {
430 matches!(
431 self,
432 RiscvAirId::MemoryGlobalInit
433 | RiscvAirId::MemoryGlobalFinalize
434 | RiscvAirId::PageProtGlobalInit
435 | RiscvAirId::PageProtGlobalFinalize
436 | RiscvAirId::Global
437 )
438 }
439
440 #[must_use]
442 pub fn is_precompile(self) -> bool {
443 matches!(
444 self,
445 RiscvAirId::ShaExtend
446 | RiscvAirId::ShaCompress
447 | RiscvAirId::EdAddAssign
448 | RiscvAirId::EdDecompress
449 | RiscvAirId::Secp256k1AddAssign
450 | RiscvAirId::Secp256k1DoubleAssign
451 | RiscvAirId::Secp256r1AddAssign
452 | RiscvAirId::Secp256r1DoubleAssign
453 | RiscvAirId::KeccakPermute
454 | RiscvAirId::Bn254AddAssign
455 | RiscvAirId::Bn254DoubleAssign
456 | RiscvAirId::Bls12381AddAssign
457 | RiscvAirId::Bls12381DoubleAssign
458 | RiscvAirId::Uint256MulMod
459 | RiscvAirId::Uint256Ops
460 | RiscvAirId::Bls12381FpOpAssign
461 | RiscvAirId::Bls12381Fp2AddSubAssign
462 | RiscvAirId::Bls12381Fp2MulAssign
463 | RiscvAirId::Bn254FpOpAssign
464 | RiscvAirId::Bn254Fp2AddSubAssign
465 | RiscvAirId::Bn254Fp2MulAssign
466 | RiscvAirId::Poseidon2
467 | RiscvAirId::SigReturn
468 )
469 }
470
471 #[must_use]
473 pub fn rows_per_event(&self) -> usize {
474 match self {
475 Self::ShaCompress => 80,
476 Self::ShaExtend => 48,
477 Self::KeccakPermute => 24,
478 _ => 1,
479 }
480 }
481
482 #[must_use]
484 pub fn control_air_id(self, page_protect_enabled: bool) -> Option<RiscvAirId> {
485 if page_protect_enabled {
486 return match self {
487 RiscvAirId::ShaCompress => Some(RiscvAirId::ShaCompressControlUser),
488 RiscvAirId::ShaExtend => Some(RiscvAirId::ShaExtendControlUser),
489 RiscvAirId::KeccakPermute => Some(RiscvAirId::KeccakPermuteControlUser),
490 _ => None,
491 };
492 }
493 match self {
494 RiscvAirId::ShaCompress => Some(RiscvAirId::ShaCompressControl),
495 RiscvAirId::ShaExtend => Some(RiscvAirId::ShaExtendControl),
496 RiscvAirId::KeccakPermute => Some(RiscvAirId::KeccakPermuteControl),
497 _ => None,
498 }
499 }
500
501 #[must_use]
503 pub fn as_str(&self) -> &'static str {
504 self.into()
505 }
506}
507
508impl FromStr for RiscvAirId {
509 type Err = String;
510
511 fn from_str(s: &str) -> Result<Self, Self::Err> {
512 let air = Self::iter().find(|chip| chip.as_str() == s);
513 match air {
514 Some(air) => Ok(air),
515 None => Err(format!("Invalid RV64IMAir: {s}")),
516 }
517 }
518}
519
520impl Display for RiscvAirId {
521 fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
522 write!(f, "{}", self.as_str())
523 }
524}
525
526#[derive(Debug, Clone, Serialize, Deserialize)]
528pub struct MaximalShapes {
529 inner: Vec<EnumMap<CoreAirId, u32>>,
530}
531
532impl FromIterator<Shape<RiscvAirId>> for MaximalShapes {
533 fn from_iter<T: IntoIterator<Item = Shape<RiscvAirId>>>(iter: T) -> Self {
534 let mut maximal_shapes = Vec::new();
535 for shape in iter {
536 let mut maximal_shape = EnumMap::<CoreAirId, u32>::default();
537 for (air, height) in shape {
538 if let Ok(core_air) = CoreAirId::try_from(air) {
539 maximal_shape[core_air] = height as u32;
540 } else if air != RiscvAirId::Program
541 && air != RiscvAirId::Byte
542 && air != RiscvAirId::Range
543 {
544 tracing::warn!("Invalid core air: {air}");
545 }
546 }
547 maximal_shapes.push(maximal_shape);
548 }
549 Self { inner: maximal_shapes }
550 }
551}
552
553impl MaximalShapes {
554 pub fn iter(&self) -> impl Iterator<Item = &EnumMap<CoreAirId, u32>> {
556 self.inner.iter()
557 }
558}