1use std::{
2 fmt::{Display, Formatter, Result as FmtResult},
3 str::FromStr,
4};
5
6use enum_map::{Enum, EnumMap};
7use serde::{Deserialize, Serialize};
8use sp1_stark::shape::Shape;
9use strum::{EnumIter, IntoEnumIterator, IntoStaticStr};
10use subenum::subenum;
11
12#[subenum(CoreAirId)]
20#[derive(
21 Debug,
22 Clone,
23 Copy,
24 PartialEq,
25 Eq,
26 Hash,
27 Serialize,
28 Deserialize,
29 EnumIter,
30 IntoStaticStr,
31 PartialOrd,
32 Ord,
33 Enum,
34)]
35pub enum RiscvAirId {
36 #[subenum(CoreAirId)]
38 Cpu = 0,
39 Program = 1,
41 ShaExtend = 2,
43 ShaCompress = 3,
45 EdAddAssign = 4,
47 EdDecompress = 5,
49 Secp256k1Decompress = 6,
51 Secp256k1AddAssign = 7,
53 Secp256k1DoubleAssign = 8,
55 Secp256r1Decompress = 9,
57 Secp256r1AddAssign = 10,
59 Secp256r1DoubleAssign = 11,
61 KeccakPermute = 12,
63 Bn254AddAssign = 13,
65 Bn254DoubleAssign = 14,
67 Bls12381AddAssign = 15,
69 Bls12381DoubleAssign = 16,
71 Uint256MulMod = 17,
73 U256XU2048Mul = 18,
75 Bls12381FpOpAssign = 19,
77 Bls12381Fp2AddSubAssign = 20,
79 Bls12381Fp2MulAssign = 21,
81 Bn254FpOpAssign = 22,
83 Bn254Fp2AddSubAssign = 23,
85 Bn254Fp2MulAssign = 24,
87 Bls12381Decompress = 25,
89 #[subenum(CoreAirId)]
91 SyscallCore = 26,
92 SyscallPrecompile = 27,
94 #[subenum(CoreAirId)]
96 DivRem = 28,
97 #[subenum(CoreAirId)]
99 AddSub = 29,
100 #[subenum(CoreAirId)]
102 Bitwise = 30,
103 #[subenum(CoreAirId)]
105 Mul = 31,
106 #[subenum(CoreAirId)]
108 ShiftRight = 32,
109 #[subenum(CoreAirId)]
111 ShiftLeft = 33,
112 #[subenum(CoreAirId)]
114 Lt = 34,
115 #[subenum(CoreAirId)]
117 MemoryInstrs = 35,
118 #[subenum(CoreAirId)]
120 Auipc = 36,
121 #[subenum(CoreAirId)]
123 Branch = 37,
124 #[subenum(CoreAirId)]
126 Jump = 38,
127 #[subenum(CoreAirId)]
129 SyscallInstrs = 39,
130 MemoryGlobalInit = 40,
132 MemoryGlobalFinalize = 41,
134 #[subenum(CoreAirId)]
136 MemoryLocal = 42,
137 #[subenum(CoreAirId)]
139 Global = 43,
140 Byte = 44,
142}
143
144impl RiscvAirId {
145 #[must_use]
147 pub fn core() -> Vec<RiscvAirId> {
148 vec![
149 RiscvAirId::Cpu,
150 RiscvAirId::AddSub,
151 RiscvAirId::Mul,
152 RiscvAirId::Bitwise,
153 RiscvAirId::ShiftLeft,
154 RiscvAirId::ShiftRight,
155 RiscvAirId::DivRem,
156 RiscvAirId::Lt,
157 RiscvAirId::Auipc,
158 RiscvAirId::MemoryLocal,
159 RiscvAirId::MemoryInstrs,
160 RiscvAirId::Branch,
161 RiscvAirId::Jump,
162 RiscvAirId::SyscallCore,
163 RiscvAirId::SyscallInstrs,
164 RiscvAirId::Global,
165 ]
166 }
167
168 #[must_use]
171 pub fn is_core(self) -> bool {
172 CoreAirId::try_from(self).is_ok()
173 }
174
175 #[must_use]
177 pub fn is_memory(self) -> bool {
178 matches!(
179 self,
180 RiscvAirId::MemoryGlobalInit | RiscvAirId::MemoryGlobalFinalize | RiscvAirId::Global
181 )
182 }
183
184 #[must_use]
186 pub fn is_precompile(self) -> bool {
187 matches!(
188 self,
189 RiscvAirId::ShaExtend |
190 RiscvAirId::ShaCompress |
191 RiscvAirId::EdAddAssign |
192 RiscvAirId::EdDecompress |
193 RiscvAirId::Secp256k1Decompress |
194 RiscvAirId::Secp256k1AddAssign |
195 RiscvAirId::Secp256k1DoubleAssign |
196 RiscvAirId::Secp256r1Decompress |
197 RiscvAirId::Secp256r1AddAssign |
198 RiscvAirId::Secp256r1DoubleAssign |
199 RiscvAirId::KeccakPermute |
200 RiscvAirId::Bn254AddAssign |
201 RiscvAirId::Bn254DoubleAssign |
202 RiscvAirId::Bls12381AddAssign |
203 RiscvAirId::Bls12381DoubleAssign |
204 RiscvAirId::Uint256MulMod |
205 RiscvAirId::U256XU2048Mul |
206 RiscvAirId::Bls12381FpOpAssign |
207 RiscvAirId::Bls12381Fp2AddSubAssign |
208 RiscvAirId::Bls12381Fp2MulAssign |
209 RiscvAirId::Bn254FpOpAssign |
210 RiscvAirId::Bn254Fp2AddSubAssign |
211 RiscvAirId::Bn254Fp2MulAssign |
212 RiscvAirId::Bls12381Decompress
213 )
214 }
215
216 #[must_use]
218 pub fn rows_per_event(&self) -> usize {
219 match self {
220 Self::ShaCompress => 80,
221 Self::ShaExtend => 48,
222 Self::KeccakPermute => 24,
223 _ => 1,
224 }
225 }
226
227 #[must_use]
229 pub fn as_str(&self) -> &'static str {
230 self.into()
231 }
232}
233
234impl FromStr for RiscvAirId {
235 type Err = String;
236
237 fn from_str(s: &str) -> Result<Self, Self::Err> {
238 let air = Self::iter().find(|chip| chip.as_str() == s);
239 match air {
240 Some(air) => Ok(air),
241 None => Err(format!("Invalid RV32IMAir: {s}")),
242 }
243 }
244}
245
246impl Display for RiscvAirId {
247 fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
248 write!(f, "{}", self.as_str())
249 }
250}
251
252#[derive(Debug, Clone, Serialize, Deserialize)]
254pub struct MaximalShapes {
255 inner: Vec<EnumMap<CoreAirId, u32>>,
256}
257
258impl FromIterator<Shape<RiscvAirId>> for MaximalShapes {
259 fn from_iter<T: IntoIterator<Item = Shape<RiscvAirId>>>(iter: T) -> Self {
260 let mut maximal_shapes = Vec::new();
261 for shape in iter {
262 let mut maximal_shape = EnumMap::<CoreAirId, u32>::default();
263 for (air, height) in shape {
264 if let Ok(core_air) = CoreAirId::try_from(air) {
265 maximal_shape[core_air] = height as u32;
266 } else if air != RiscvAirId::Program && air != RiscvAirId::Byte {
267 tracing::warn!("Invalid core air: {air}");
268 }
269 }
270 maximal_shapes.push(maximal_shape);
271 }
272 Self { inner: maximal_shapes }
273 }
274}
275
276impl MaximalShapes {
277 pub fn iter(&self) -> impl Iterator<Item = &EnumMap<CoreAirId, u32>> {
279 self.inner.iter()
280 }
281}