From 2ff9ce9e9335158b4a620e35b38e59879e8d1ae9 Mon Sep 17 00:00:00 2001
From: Yael Doweck <yael@starkware.co>
Date: Tue, 17 Oct 2023 15:00:12 +0300
Subject: [PATCH] feat: dump_declared_classes can dump for a specific block
 range

---
 Cargo.lock                               |  1 +
 crates/dump_declared_classes/Cargo.toml  |  1 +
 crates/dump_declared_classes/src/main.rs | 52 +++++++++++++++++++++---
 crates/papyrus_storage/src/utils.rs      | 35 ++++++++++++++++
 4 files changed, 84 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 583b0ee602..5324d066b5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1906,6 +1906,7 @@ checksum = "9ea835d29036a4087793836fa931b08837ad5e957da9e23886b29586fb9b6650"
 name = "dump_declared_classes"
 version = "0.0.5"
 dependencies = [
+ "clap",
  "papyrus_storage",
 ]
 
diff --git a/crates/dump_declared_classes/Cargo.toml b/crates/dump_declared_classes/Cargo.toml
index 20c6383af1..b53f5a9990 100644
--- a/crates/dump_declared_classes/Cargo.toml
+++ b/crates/dump_declared_classes/Cargo.toml
@@ -7,3 +7,4 @@ license-file.workspace = true
 
 [dependencies]
 papyrus_storage = { path = "../papyrus_storage", version = "0.0.5" }
+clap.workspace = true
diff --git a/crates/dump_declared_classes/src/main.rs b/crates/dump_declared_classes/src/main.rs
index bf50f226f0..9288f02456 100644
--- a/crates/dump_declared_classes/src/main.rs
+++ b/crates/dump_declared_classes/src/main.rs
@@ -1,14 +1,56 @@
-use papyrus_storage::utils::dump_declared_classes_table_to_file;
+use clap::{Arg, Command};
+use papyrus_storage::utils::{
+    dump_declared_classes_table_by_block_range,
+    dump_declared_classes_table_to_file,
+};
+use papyrus_storage::StorageResult;
 
 /// This executable dumps the declared_classes table from the storage to a file.
 /// The file path can be passed as an argument, otherwise it will be dumped to
 /// "dump_declared_classes.json".
+/// A starting block and an ending block can also be passed as optional arguments, otherwise the
+/// entire table will be dumped.
 fn main() {
-    let args = std::env::args().collect::<Vec<_>>();
-    let default_file_path = "dump_declared_classes.json".to_string();
-    let file_path = args.get(1).unwrap_or(&default_file_path);
+    let matches = Command::new("Dump declared classes")
+        .arg(
+            Arg::new("file_path")
+                .short('f')
+                .long("file_path")
+                .default_value("dump_declared_classes.json")
+                .help("The file path to dump the declared classes table to."),
+        )
+        .arg(
+            Arg::new("start_block")
+                .short('s')
+                .long("start_block")
+                .help("The block number to start dumping from."),
+        )
+        .arg(
+            Arg::new("end_block")
+                .short('e')
+                .long("end_block")
+                .help("The block number to end dumping at."),
+        )
+        .get_matches();
 
-    match dump_declared_classes_table_to_file(file_path) {
+    let file_path = matches.get_one::<String>("file_path").unwrap().as_str();
+    let res: StorageResult<()> =
+        if matches.contains_id("start_block") && matches.contains_id("end_block") {
+            let start_block = matches
+                .get_one::<String>("start_block")
+                .unwrap()
+                .parse::<u64>()
+                .expect("Failed parsing start_block");
+            let end_block = matches
+                .get_one::<String>("end_block")
+                .unwrap()
+                .parse::<u64>()
+                .expect("Failed parsing end_block");
+            dump_declared_classes_table_by_block_range(start_block, end_block, file_path)
+        } else {
+            dump_declared_classes_table_to_file(file_path)
+        };
+    match res {
         Ok(_) => println!("Dumped declared_classes table to file: {}", file_path),
         Err(e) => println!("Failed dumping declared_classes table with error: {}", e),
     }
diff --git a/crates/papyrus_storage/src/utils.rs b/crates/papyrus_storage/src/utils.rs
index 29d6e69206..3e6f4b438c 100644
--- a/crates/papyrus_storage/src/utils.rs
+++ b/crates/papyrus_storage/src/utils.rs
@@ -6,8 +6,11 @@ mod utils_test;
 use std::fs::File;
 use std::io::{BufWriter, Write};
 
+use starknet_api::block::BlockNumber;
+
 use crate::db::serialization::StorageSerde;
 use crate::db::{DbIter, TableIdentifier, RO};
+use crate::state::StateStorageReader;
 use crate::{open_storage, StorageConfig, StorageResult, StorageTxn};
 
 /// Dumps a table from the storage to a file in JSON format.
@@ -46,3 +49,35 @@ pub fn dump_declared_classes_table_to_file(file_path: &str) -> StorageResult<()>
     dump_table_to_file(&txn, &txn.tables.declared_classes, file_path)?;
     Ok(())
 }
+
+/// Dumps the declared_classes at a given block range from the storage to a file.
+pub fn dump_declared_classes_table_by_block_range(
+    start_block: u64,
+    end_block: u64,
+    file_path: &str,
+) -> StorageResult<()> {
+    let storage_config = StorageConfig::default();
+    let (storage_reader, _) = open_storage(storage_config)?;
+    let txn = storage_reader.begin_ro_txn()?;
+    let table_handle = txn.txn.open_table(&txn.tables.declared_classes)?;
+    let file = File::create(file_path)?;
+    let mut writer = BufWriter::new(file);
+    writer.write_all(b"[")?;
+    let mut first = true;
+    for block_number in start_block..=end_block {
+        println!("block_number: {}", block_number);
+        if let Some(thin_state_diff) = txn.get_state_diff(BlockNumber(block_number))? {
+            for (class_hash, _) in thin_state_diff.declared_classes.iter() {
+                if let Some(contract_class) = table_handle.get(&txn.txn, class_hash)? {
+                    if !first {
+                        writer.write_all(b",")?;
+                    }
+                    serde_json::to_writer(&mut writer, &(class_hash, &contract_class))?;
+                    first = false;
+                }
+            }
+        };
+    }
+    writer.write_all(b"]")?;
+    Ok(())
+}