From 2609f17c2ab30705e4b92c3dd738093ff70ec234 Mon Sep 17 00:00:00 2001
From: 0xLucca <95830307+0xLucca@users.noreply.github.com>
Date: Wed, 13 Nov 2024 22:26:45 +0700
Subject: [PATCH] Add snapshot functionality to ink_sandbox (#2261)

* Add Snapshot functionality

* Add Snapshot testing

* Add snapshot impl to sandbox_client
---
 crates/e2e/sandbox/src/api/system_api.rs | 22 ++++++++++++++++
 crates/e2e/sandbox/src/lib.rs            | 19 ++++++++++++++
 crates/e2e/sandbox/src/macros.rs         | 27 ++++++++++++++++++++
 crates/e2e/src/sandbox_client.rs         | 32 ++++++++++++++++++++++++
 4 files changed, 100 insertions(+)

diff --git a/crates/e2e/sandbox/src/api/system_api.rs b/crates/e2e/sandbox/src/api/system_api.rs
index e8b416c4c1f..c743afe28a5 100644
--- a/crates/e2e/sandbox/src/api/system_api.rs
+++ b/crates/e2e/sandbox/src/api/system_api.rs
@@ -197,4 +197,26 @@ mod tests {
         make_transfer(&mut sandbox, RECIPIENT, 1).expect("Failed to make transfer");
         assert!(!sandbox.events().is_empty());
     }
+
+    #[test]
+    fn snapshot_works() {
+        let mut sandbox = DefaultSandbox::default();
+
+        // Check state before
+        let block_before = sandbox.block_number();
+        let snapshot_before = sandbox.take_snapshot();
+
+        // Advance some blocks to have some state change
+        let _ = sandbox.build_blocks(5);
+        let block_after = sandbox.block_number();
+
+        // Check block number and state after
+        assert_eq!(block_before + 5, block_after);
+
+        // Restore state
+        sandbox.restore_snapshot(snapshot_before);
+
+        // Check state after restore
+        assert_eq!(block_before, sandbox.block_number());
+    }
 }
diff --git a/crates/e2e/sandbox/src/lib.rs b/crates/e2e/sandbox/src/lib.rs
index 2f05d4b8842..d4a46c4e1d9 100644
--- a/crates/e2e/sandbox/src/lib.rs
+++ b/crates/e2e/sandbox/src/lib.rs
@@ -23,6 +23,7 @@ use pallet_contracts::{
 };
 /// Export pallets that are used in [`crate::create_sandbox`]
 pub use {
+    frame_support::sp_runtime::testing::H256,
     frame_support::{
         self,
         sp_runtime::{
@@ -43,6 +44,18 @@ pub use {
     sp_io::TestExternalities,
 };
 
+/// A snapshot of the storage.
+#[derive(Clone, Debug)]
+pub struct Snapshot {
+    /// The storage raw key-value pairs.
+    pub storage: RawStorage,
+    /// The storage root hash.
+    pub storage_root: StorageRoot,
+}
+
+pub type RawStorage = Vec<(Vec<u8>, (Vec<u8>, i32))>;
+pub type StorageRoot = H256;
+
 /// Alias for the balance type.
 type BalanceOf<R> =
     <<R as pallet_contracts::Config>::Currency as Inspect<AccountIdFor<R>>>::Balance;
@@ -118,4 +131,10 @@ pub trait Sandbox {
     fn convert_account_to_origin(
         account: AccountIdFor<Self::Runtime>,
     ) -> <<Self::Runtime as frame_system::Config>::RuntimeCall as Dispatchable>::RuntimeOrigin;
+
+    /// Take a snapshot of the storage.
+    fn take_snapshot(&mut self) -> Snapshot;
+
+    /// Restore the storage from the given snapshot.
+    fn restore_snapshot(&mut self, snapshot: Snapshot);
 }
diff --git a/crates/e2e/sandbox/src/macros.rs b/crates/e2e/sandbox/src/macros.rs
index 0812aba55fe..6b6feafa93d 100644
--- a/crates/e2e/sandbox/src/macros.rs
+++ b/crates/e2e/sandbox/src/macros.rs
@@ -115,6 +115,8 @@ mod construct_runtime {
         weights::Weight,
     };
 
+    use $crate::Snapshot;
+
     // Define the runtime type as a collection of pallets
     construct_runtime!(
         pub enum $runtime {
@@ -289,6 +291,31 @@ mod construct_runtime {
         ) -> <<Self::Runtime as $crate::frame_system::Config>::RuntimeCall as $crate::frame_support::sp_runtime::traits::Dispatchable>::RuntimeOrigin {
             Some(account).into()
         }
+
+        fn take_snapshot(&mut self) -> Snapshot {
+            let mut backend = self.ext.as_backend().clone();
+            let raw_key_values = backend
+                .backend_storage_mut()
+                .drain()
+                .into_iter()
+                .filter(|(_, (_, r))| *r > 0)
+                .collect::<Vec<(Vec<u8>, (Vec<u8>, i32))>>();
+            let root = backend.root().to_owned();
+            Snapshot {
+                storage: raw_key_values,
+                storage_root: root,
+            }
+        }
+
+        fn restore_snapshot(&mut self, snapshot: Snapshot) {
+            self.ext = $crate::TestExternalities::from_raw_snapshot(
+                snapshot.storage,
+                snapshot.storage_root,
+                Default::default(),
+            );
+        }
+
+
     }
 }
 
diff --git a/crates/e2e/src/sandbox_client.rs b/crates/e2e/src/sandbox_client.rs
index 91f7f9b2d1c..fe5358e2d2c 100644
--- a/crates/e2e/src/sandbox_client.rs
+++ b/crates/e2e/src/sandbox_client.rs
@@ -469,6 +469,7 @@ pub mod preset {
             Extension,
             RuntimeMetadataPrefixed,
             Sandbox,
+            Snapshot,
         };
         pub use pallet_contracts_mock_network::*;
         use sp_runtime::traits::Dispatchable;
@@ -543,6 +544,37 @@ pub mod preset {
             {
                 Some(account).into()
             }
+
+            fn take_snapshot(&mut self) -> Snapshot {
+                EXT_PARAA.with(|v| {
+                    let v = v.borrow();
+                    let mut backend = v.as_backend().clone();
+                    let raw_key_values = backend
+                        .backend_storage_mut()
+                        .drain()
+                        .into_iter()
+                        .filter(|(_, (_, r))| *r > 0)
+                        .collect::<Vec<(Vec<u8>, (Vec<u8>, i32))>>();
+                    let root = backend.root().to_owned();
+
+                    Snapshot {
+                        storage: raw_key_values,
+                        storage_root: root,
+                    }
+                })
+            }
+
+            fn restore_snapshot(&mut self, snapshot: ink_sandbox::Snapshot) {
+                EXT_PARAA.with(|v| {
+                    let mut v = v.borrow_mut();
+
+                    *v = ink_sandbox::TestExternalities::from_raw_snapshot(
+                        snapshot.storage,
+                        snapshot.storage_root,
+                        Default::default(),
+                    );
+                })
+            }
         }
     }
 }