From 02ce315ee6f0570f7b961b0877d933ea0d6eefb1 Mon Sep 17 00:00:00 2001
From: Noam Spiegelstein <noam.s@starkware.co>
Date: Wed, 5 Feb 2025 18:21:45 +0200
Subject: [PATCH] feat(apollo_reverts): add revert blocks functionality

---
 Cargo.lock                       | 10 +++++
 Cargo.toml                       |  2 +
 commitlint.config.js             |  1 +
 crates/apollo_reverts/Cargo.toml | 15 +++++++
 crates/apollo_reverts/src/lib.rs | 73 ++++++++++++++++++++++++++++++++
 5 files changed, 101 insertions(+)
 create mode 100644 crates/apollo_reverts/Cargo.toml
 create mode 100644 crates/apollo_reverts/src/lib.rs

diff --git a/Cargo.lock b/Cargo.lock
index 7950d4389e4..aeeba44513f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -640,6 +640,16 @@ version = "1.0.95"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "34ac096ce696dc2fcabef30516bb13c0a68a11d30131d3df6f04711467681b04"
 
+[[package]]
+name = "apollo_reverts"
+version = "0.0.0"
+dependencies = [
+ "futures",
+ "papyrus_storage",
+ "starknet_api",
+ "tracing",
+]
+
 [[package]]
 name = "aquamarine"
 version = "0.6.0"
diff --git a/Cargo.toml b/Cargo.toml
index 65d7f366534..d0408966988 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -4,6 +4,7 @@
 resolver = "2"
 
 members = [
+  "crates/apollo_reverts",
   "crates/blockifier",
   "crates/blockifier_reexecution",
   "crates/mempool_test_utils",
@@ -83,6 +84,7 @@ alloy-sol-types = "0.8.3"
 alloy-transport = "0.3.5"
 alloy-transport-http = "0.3.5"
 anyhow = "1.0.44"
+apollo_reverts = { path = "crates/apollo_reverts", version = "0.0.0" }
 ark-ec = "0.4.2"
 ark-ff = "0.4.0-alpha.7"
 ark-secp256k1 = "0.4.0"
diff --git a/commitlint.config.js b/commitlint.config.js
index 316f6668e78..feef9816264 100644
--- a/commitlint.config.js
+++ b/commitlint.config.js
@@ -20,6 +20,7 @@ const Configuration = {
     rules: {
         'scope-empty': [2, 'never'],
         'scope-enum': [2, 'always', [
+            "apollo_reverts",
             'blockifier',
             'blockifier_reexecution',
             'cairo_native',
diff --git a/crates/apollo_reverts/Cargo.toml b/crates/apollo_reverts/Cargo.toml
new file mode 100644
index 00000000000..b504571436a
--- /dev/null
+++ b/crates/apollo_reverts/Cargo.toml
@@ -0,0 +1,15 @@
+[package]
+name = "apollo_reverts"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+futures.workspace = true
+papyrus_storage.workspace = true
+starknet_api.workspace = true
+tracing.workspace = true
diff --git a/crates/apollo_reverts/src/lib.rs b/crates/apollo_reverts/src/lib.rs
new file mode 100644
index 00000000000..17895fd0341
--- /dev/null
+++ b/crates/apollo_reverts/src/lib.rs
@@ -0,0 +1,73 @@
+use std::future::Future;
+
+use futures::future::pending;
+use futures::never::Never;
+use papyrus_storage::base_layer::BaseLayerStorageWriter;
+use papyrus_storage::body::BodyStorageWriter;
+use papyrus_storage::class_manager::ClassManagerStorageWriter;
+use papyrus_storage::header::HeaderStorageWriter;
+use papyrus_storage::state::StateStorageWriter;
+use papyrus_storage::StorageWriter;
+use starknet_api::block::BlockNumber;
+use tracing::info;
+
+pub async fn revert_blocks_and_eternal_pending<Fut>(
+    mut current_block_number: BlockNumber,
+    revert_up_to_and_including: BlockNumber,
+    mut revert_block_fn: impl FnMut(BlockNumber) -> Fut,
+    component_name: &str,
+) -> Never
+where
+    Fut: Future<Output = ()>,
+{
+    if current_block_number <= revert_up_to_and_including {
+        panic!(
+            "{component_name} current block {current_block_number} is not larger than the target \
+             block {revert_up_to_and_including}. No reverts are needed."
+        );
+    }
+
+    info!(
+        "Reverting {component_name} from block {current_block_number} to block \
+         {revert_up_to_and_including}"
+    );
+
+    while current_block_number > revert_up_to_and_including {
+        current_block_number = current_block_number.prev().expect(
+            "A block number that's greater than another block number should return Some on prev",
+        );
+        info!("Reverting {component_name} block number {current_block_number}.");
+        revert_block_fn(current_block_number).await;
+    }
+
+    info!(
+        "Done reverting {component_name} up to block {revert_up_to_and_including} including. \
+         Starting eternal pending."
+    );
+    pending().await
+}
+
+pub fn revert_block(
+    storage_writer: &mut StorageWriter,
+    current_block_number: BlockNumber,
+) -> impl Future<Output = ()> {
+    storage_writer
+        .begin_rw_txn()
+        .unwrap()
+        .revert_header(current_block_number)
+        .unwrap()
+        .0
+        .revert_body(current_block_number)
+        .unwrap()
+        .0
+        .revert_state_diff(current_block_number)
+        .unwrap()
+        .0
+        .try_revert_class_manager_marker(current_block_number)
+        .unwrap()
+        .try_revert_base_layer_marker(current_block_number)
+        .unwrap()
+        .commit()
+        .unwrap();
+    async {}
+}