From e2de5ecd91a70bd05fa64b0905f8e0718a220aae Mon Sep 17 00:00:00 2001
From: YaelD <70628564+Yael-Starkware@users.noreply.github.com>
Date: Wed, 21 Feb 2024 16:34:46 +0200
Subject: [PATCH] refactor: add BlockWeights for the future bouncer (#1444)

---
 crates/blockifier/src/bouncer.rs      | 49 ++++++++++++++++++++
 crates/blockifier/src/bouncer_test.rs | 64 +++++++++++++++++++++++++++
 crates/blockifier/src/lib.rs          |  1 +
 3 files changed, 114 insertions(+)
 create mode 100644 crates/blockifier/src/bouncer.rs
 create mode 100644 crates/blockifier/src/bouncer_test.rs

diff --git a/crates/blockifier/src/bouncer.rs b/crates/blockifier/src/bouncer.rs
new file mode 100644
index 0000000000..0bcb4b990c
--- /dev/null
+++ b/crates/blockifier/src/bouncer.rs
@@ -0,0 +1,49 @@
+use serde::Deserialize;
+
+#[cfg(test)]
+#[path = "bouncer_test.rs"]
+mod test;
+
+macro_rules! impl_checked_sub {
+    ($($field:ident),+) => {
+        pub fn checked_sub(self: Self, other: Self) -> Option<Self> {
+            Some(
+                Self {
+                    $(
+                        $field: self.$field.checked_sub(other.$field)?,
+                    )+
+                }
+            )
+        }
+    };
+}
+
+#[derive(Clone, Copy, Debug, Default, derive_more::Sub, Deserialize, PartialEq)]
+/// Represents the execution resources counted throughout block creation.
+pub struct BouncerWeights {
+    gas: u64,
+    n_steps: u64,
+    message_segment_length: u64,
+    state_diff_size: u64,
+    builtin_count: BuiltinCount,
+}
+
+impl BouncerWeights {
+    impl_checked_sub!(gas, n_steps, message_segment_length, state_diff_size, builtin_count);
+}
+
+#[derive(Clone, Copy, Debug, Default, derive_more::Sub, Deserialize, PartialEq)]
+pub struct BuiltinCount {
+    bitwise: u64,
+    ecdsa: u64,
+    ec_op: u64,
+    keccak: u64,
+    output: u64,
+    pedersen: u64,
+    poseidon: u64,
+    range_check: u64,
+}
+
+impl BuiltinCount {
+    impl_checked_sub!(bitwise, ecdsa, ec_op, keccak, output, pedersen, poseidon, range_check);
+}
diff --git a/crates/blockifier/src/bouncer_test.rs b/crates/blockifier/src/bouncer_test.rs
new file mode 100644
index 0000000000..4785fe56b2
--- /dev/null
+++ b/crates/blockifier/src/bouncer_test.rs
@@ -0,0 +1,64 @@
+use std::ops::Sub;
+
+use crate::bouncer::{BouncerWeights, BuiltinCount};
+
+#[test]
+fn test_block_weights_sub_checked() {
+    let max_bouncer_weights = BouncerWeights {
+        gas: 10,
+        n_steps: 10,
+        message_segment_length: 10,
+        state_diff_size: 10,
+        builtin_count: BuiltinCount {
+            bitwise: 10,
+            ecdsa: 10,
+            ec_op: 10,
+            keccak: 10,
+            output: 10,
+            pedersen: 10,
+            poseidon: 10,
+            range_check: 10,
+        },
+    };
+
+    let bouncer_weights = BouncerWeights {
+        gas: 7,
+        n_steps: 0,
+        message_segment_length: 10,
+        state_diff_size: 7,
+        builtin_count: BuiltinCount {
+            bitwise: 6,
+            ecdsa: 7,
+            ec_op: 7,
+            keccak: 8,
+            output: 7,
+            pedersen: 7,
+            poseidon: 9,
+            range_check: 10,
+        },
+    };
+
+    let result = max_bouncer_weights.checked_sub(bouncer_weights).unwrap();
+    let difference_bouncer_weights = max_bouncer_weights.sub(bouncer_weights);
+    assert_eq!(result, difference_bouncer_weights);
+
+    let bouncer_weights_exceeds_max = BouncerWeights {
+        gas: 5,
+        n_steps: 5,
+        message_segment_length: 5,
+        state_diff_size: 5,
+        builtin_count: BuiltinCount {
+            bitwise: 11,
+            ecdsa: 5,
+            ec_op: 5,
+            keccak: 5,
+            output: 5,
+            pedersen: 5,
+            poseidon: 5,
+            range_check: 5,
+        },
+    };
+
+    let result = max_bouncer_weights.checked_sub(bouncer_weights_exceeds_max);
+    assert!(result.is_none());
+}
diff --git a/crates/blockifier/src/lib.rs b/crates/blockifier/src/lib.rs
index 6ddbc8b77f..1783f6fede 100644
--- a/crates/blockifier/src/lib.rs
+++ b/crates/blockifier/src/lib.rs
@@ -11,6 +11,7 @@
 
 pub mod abi;
 pub mod blockifier;
+pub mod bouncer;
 pub mod context;
 pub mod execution;
 pub mod fee;