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 ShaCompress = 3,
46 ShaCompressControl = 4,
48 EdAddAssign = 5,
50 EdDecompress = 6,
52 Secp256k1AddAssign = 7,
54 Secp256k1DoubleAssign = 8,
56 Secp256r1AddAssign = 9,
58 Secp256r1DoubleAssign = 10,
60 KeccakPermute = 11,
62 KeccakPermuteControl = 12,
64 Bn254AddAssign = 13,
66 Bn254DoubleAssign = 14,
68 Bls12381AddAssign = 15,
70 Bls12381DoubleAssign = 16,
72 Uint256MulMod = 17,
74 Uint256Ops = 18,
76 Bls12381FpOpAssign = 19,
78 Bls12381Fp2AddSubAssign = 20,
80 Bls12381Fp2MulAssign = 21,
82 Bn254FpOpAssign = 22,
84 Bn254Fp2AddSubAssign = 23,
86 Bn254Fp2MulAssign = 24,
88 #[subenum(CoreAirId)]
90 SyscallCore = 25,
91 SyscallPrecompile = 26,
93 #[subenum(CoreAirId)]
95 DivRem = 27,
96 #[subenum(CoreAirId)]
98 Add = 28,
99 #[subenum(CoreAirId)]
101 Addi = 29,
102 #[subenum(CoreAirId)]
104 Addw = 30,
105 #[subenum(CoreAirId)]
107 Sub = 31,
108 #[subenum(CoreAirId)]
110 Subw = 32,
111 #[subenum(CoreAirId)]
113 Bitwise = 33,
114 #[subenum(CoreAirId)]
116 Mul = 34,
117 #[subenum(CoreAirId)]
119 ShiftRight = 35,
120 #[subenum(CoreAirId)]
122 ShiftLeft = 36,
123 #[subenum(CoreAirId)]
125 Lt = 37,
126 #[subenum(CoreAirId)]
128 LoadByte = 38,
129 #[subenum(CoreAirId)]
131 LoadHalf = 39,
132 #[subenum(CoreAirId)]
134 LoadWord = 40,
135 #[subenum(CoreAirId)]
137 LoadX0 = 41,
138 #[subenum(CoreAirId)]
140 LoadDouble = 42,
141 #[subenum(CoreAirId)]
143 StoreByte = 43,
144 #[subenum(CoreAirId)]
146 StoreHalf = 44,
147 #[subenum(CoreAirId)]
149 StoreWord = 45,
150 #[subenum(CoreAirId)]
152 StoreDouble = 46,
153 #[subenum(CoreAirId)]
155 UType = 47,
156 #[subenum(CoreAirId)]
158 Branch = 48,
159 #[subenum(CoreAirId)]
161 Jal = 49,
162 #[subenum(CoreAirId)]
164 Jalr = 50,
165 #[subenum(CoreAirId)]
167 SyscallInstrs = 51,
168 #[subenum(CoreAirId)]
170 MemoryBump = 52,
171 #[subenum(CoreAirId)]
173 StateBump = 53,
174 MemoryGlobalInit = 54,
176 MemoryGlobalFinalize = 55,
178 #[subenum(CoreAirId)]
180 MemoryLocal = 56,
181 #[subenum(CoreAirId)]
183 Global = 57,
184 Byte = 58,
186 Range = 59,
188 Poseidon2 = 60,
190 #[subenum(CoreAirId)]
192 AluX0 = 61,
193}
194
195impl RiscvAirId {
196 #[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 #[must_use]
237 pub fn is_core(self) -> bool {
238 CoreAirId::try_from(self).is_ok()
239 }
240
241 #[must_use]
243 pub fn is_memory(self) -> bool {
244 matches!(
245 self,
246 RiscvAirId::MemoryGlobalInit | RiscvAirId::MemoryGlobalFinalize | RiscvAirId::Global
247 )
248 }
249
250 #[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 #[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 #[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 #[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#[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 pub fn iter(&self) -> impl Iterator<Item = &EnumMap<CoreAirId, u32>> {
357 self.inner.iter()
358 }
359}