feat: estimate prover gas while executing

crates/core/executor/Cargo.toml

@@ -65,3 +65,4 @@ profiling = [
   "dep:gecko_profile",
   "dep:indicatif",
 ]
+gas = []

crates/core/executor/src/air.rs

@@ -62,9 +62,9 @@ pub enum RiscvAirId {
     /// The bls12-381 fp op assign chip.
     Bls12381FpOpAssign = 19,
     /// The bls12-831 fp2 add sub assign chip.
-    Bls12831Fp2AddSubAssign = 20,
+    Bls12381Fp2AddSubAssign = 20,
     /// The bls12-831 fp2 mul assign chip.
-    Bls12831Fp2MulAssign = 21,
+    Bls12381Fp2MulAssign = 21,
     /// The bn254 fp2 add sub assign chip.
     Bn254FpOpAssign = 22,
     /// The bn254 fp op assign chip.
@@ -151,6 +151,84 @@ impl RiscvAirId {
             RiscvAirId::Global,
         ]
     }
+
+    /// TODO replace these three with subenums or something
+    /// Whether the ID represents a core AIR.
+    #[must_use]
+    pub fn is_core(self) -> bool {
+        matches!(
+            self,
+            RiscvAirId::Cpu
+                | RiscvAirId::AddSub
+                | RiscvAirId::Mul
+                | RiscvAirId::Bitwise
+                | RiscvAirId::ShiftLeft
+                | RiscvAirId::ShiftRight
+                | RiscvAirId::DivRem
+                | RiscvAirId::Lt
+                | RiscvAirId::Auipc
+                | RiscvAirId::MemoryLocal
+                | RiscvAirId::MemoryInstrs
+                | RiscvAirId::Branch
+                | RiscvAirId::Jump
+                | RiscvAirId::SyscallCore
+                | RiscvAirId::SyscallInstrs
+                | RiscvAirId::Global,
+        )
+    }
+
+    /// Whether the ID represents a memory AIR.
+    #[must_use]
+    pub fn is_memory(self) -> bool {
+        matches!(
+            self,
+            RiscvAirId::MemoryGlobalInit | RiscvAirId::MemoryGlobalFinalize | RiscvAirId::Global
+        )
+    }
+
+    /// Whether the ID represents a precompile AIR.
+    #[must_use]
+    pub fn is_precompile(self) -> bool {
+        matches!(
+            self,
+            RiscvAirId::ShaExtend
+                | RiscvAirId::ShaCompress
+                | RiscvAirId::EdAddAssign
+                | RiscvAirId::EdDecompress
+                | RiscvAirId::Secp256k1Decompress
+                | RiscvAirId::Secp256k1AddAssign
+                | RiscvAirId::Secp256k1DoubleAssign
+                | RiscvAirId::Secp256r1Decompress
+                | RiscvAirId::Secp256r1AddAssign
+                | RiscvAirId::Secp256r1DoubleAssign
+                | RiscvAirId::KeccakPermute
+                | RiscvAirId::Bn254AddAssign
+                | RiscvAirId::Bn254DoubleAssign
+                | RiscvAirId::Bls12381AddAssign
+                | RiscvAirId::Bls12381DoubleAssign
+                | RiscvAirId::Uint256MulMod
+                | RiscvAirId::U256XU2048Mul
+                | RiscvAirId::Bls12381FpOpAssign
+                | RiscvAirId::Bls12381Fp2AddSubAssign
+                | RiscvAirId::Bls12381Fp2MulAssign
+                | RiscvAirId::Bn254FpOpAssign
+                | RiscvAirId::Bn254Fp2AddSubAssign
+                | RiscvAirId::Bn254Fp2MulAssign
+                | RiscvAirId::Bls12381Decompress
+        )
+    }
+
+    /// >>>>>> FIX BEFORE MERGING make the visibility pub(crate)
+    #[must_use]
+    pub fn rows_per_event(&self) -> usize {
+        match self {
+            Self::ShaCompress => 80,
+            Self::ShaExtend => 48,
+            Self::KeccakPermute => 24,
+            _ => 1,
+        }
+    }
+
     /// Returns the string representation of the AIR.
     #[must_use]
     pub fn as_str(&self) -> &str {
@@ -175,8 +253,8 @@ impl RiscvAirId {
             Self::Uint256MulMod => "Uint256MulMod",
             Self::U256XU2048Mul => "U256XU2048Mul",
             Self::Bls12381FpOpAssign => "Bls12381FpOpAssign",
-            Self::Bls12831Fp2AddSubAssign => "Bls12831Fp2AddSubAssign",
-            Self::Bls12831Fp2MulAssign => "Bls12831Fp2MulAssign",
+            Self::Bls12381Fp2AddSubAssign => "Bls12381Fp2AddSubAssign",
+            Self::Bls12381Fp2MulAssign => "Bls12381Fp2MulAssign",
             Self::Bn254FpOpAssign => "Bn254FpOpAssign",
             Self::Bn254Fp2AddSubAssign => "Bn254Fp2AddSubAssign",
             Self::Bn254Fp2MulAssign => "Bn254Fp2MulAssign",

crates/core/executor/src/artifacts/rv32im_costs.json

@@ -1,7 +1,6 @@
 {
   "U256XU2048Mul": 5849,
   "Bn254Fp2MulAssign": 2957,
-  "Bls12831Fp2AddSubAssign": 2174,
   "SyscallInstrs": 80,
   "DivRem": 146,
   "ShiftRight": 134,
@@ -11,7 +10,7 @@
   "Bn254AddAssign": 4461,
   "ShiftLeft": 68,
   "Auipc": 41,
-  "KeccakPermute": 93720,
+  "KeccakPermute": 3905,
   "Program": 31,
   "MemoryLocal": 100,
   "Global": 428,
@@ -20,8 +19,9 @@
   "Jump": 50,
   "Bn254FpOpAssign": 738,
   "Mul": 83,
-  "ShaExtend": 15936,
+  "ShaExtend": 332,
   "Bls12381AddAssign": 6717,
+  "Bls12381Fp2AddSubAssign": 2174,
   "MemoryGlobalFinalize": 124,
   "Byte": 52,
   "EdDecompress": 3100,
@@ -31,7 +31,7 @@
   "Bn254DoubleAssign": 4564,
   "Uint256MulMod": 880,
   "Bls12381DoubleAssign": 6876,
-  "Bls12831Fp2MulAssign": 4445,
+  "Bls12381Fp2MulAssign": 4445,
   "EdAddAssign": 3709,
   "Bls12381Decompress": 4149,
   "Lt": 53,
@@ -41,7 +41,7 @@
   "Bn254Fp2AddSubAssign": 1454,
   "Bls12381FpOpAssign": 1098,
   "Cpu": 109,
-  "ShaCompress": 40480,
+  "ShaCompress": 506,
   "MemoryInstrs": 93,
   "Secp256k1DoubleAssign": 4564
 }

crates/core/executor/src/estimator.rs

@@ -0,0 +1,28 @@
+//! Data that may be collected during execution and used to estimate trace area.
+
+use std::ops::AddAssign;
+
+use enum_map::EnumMap;
+
+use crate::RiscvAirId;
+
+/// Data accumulated during execution to estimate the core trace area used to prove the execution.
+#[derive(Clone, Debug, Default)]
+pub struct TraceAreaEstimator {
+    /// Core shards, represented by the number of events per AIR.
+    pub core_shards: Vec<EnumMap<RiscvAirId, u64>>,
+    /// Deferred events, which are used to calculate trace area after execution has finished.
+    pub deferred_events: EnumMap<RiscvAirId, u64>,
+}
+
+impl AddAssign for TraceAreaEstimator {
+    fn add_assign(&mut self, rhs: Self) {
+        let TraceAreaEstimator { core_shards, deferred_events } = self;
+        core_shards.extend(rhs.core_shards);
+        deferred_events
+            .as_mut_array()
+            .iter_mut()
+            .zip(rhs.deferred_events.as_array())
+            .for_each(|(l, r)| *l += r);
+    }
+}