From 19913f47ba1f96480fed42539d64cad9382be64e Mon Sep 17 00:00:00 2001
From: Blazej Kolad <blazejkolad@gmail.com>
Date: Mon, 28 Aug 2023 18:20:19 +0200
Subject: [PATCH] Update `end_slot_hook` (#742)

* add root_hash to end_sloot_hook

* update end slot hook

* support experimental

* allow unused variables

* fix experimental flag for hooks in evm
---
 examples/demo-stf/src/runtime.rs                         | 7 ++++++-
 .../integration-tests/src/chain_state/helpers.rs         | 7 ++++++-
 .../module-implementations/sov-chain-state/src/hooks.rs  | 7 ++++++-
 .../module-implementations/sov-evm/src/hooks.rs          | 9 +++++++++
 module-system/module-implementations/sov-evm/src/lib.rs  | 2 ++
 module-system/sov-modules-api/src/hooks.rs               | 6 +++++-
 module-system/sov-modules-stf-template/src/lib.rs        | 3 +--
 7 files changed, 35 insertions(+), 6 deletions(-)
 create mode 100644 module-system/module-implementations/sov-evm/src/hooks.rs

diff --git a/examples/demo-stf/src/runtime.rs b/examples/demo-stf/src/runtime.rs
index 23141f4ca0..699464a86d 100644
--- a/examples/demo-stf/src/runtime.rs
+++ b/examples/demo-stf/src/runtime.rs
@@ -114,8 +114,13 @@ impl<C: Context, Cond: ValidityCondition> SlotHooks<Cond> for Runtime<C> {
 
     fn end_slot_hook(
         &self,
-        _working_set: &mut sov_state::WorkingSet<<Self::Context as sov_modules_api::Spec>::Storage>,
+        #[allow(unused_variables)] root_hash: [u8; 32],
+        #[allow(unused_variables)] working_set: &mut sov_state::WorkingSet<
+            <Self::Context as sov_modules_api::Spec>::Storage,
+        >,
     ) {
+        #[cfg(feature = "experimental")]
+        self.evm.end_slot_hook(root_hash, working_set);
     }
 }
 
diff --git a/module-system/module-implementations/integration-tests/src/chain_state/helpers.rs b/module-system/module-implementations/integration-tests/src/chain_state/helpers.rs
index 2987198489..65744b247e 100644
--- a/module-system/module-implementations/integration-tests/src/chain_state/helpers.rs
+++ b/module-system/module-implementations/integration-tests/src/chain_state/helpers.rs
@@ -72,7 +72,12 @@ impl<C: Context, Cond: ValidityCondition> SlotHooks<Cond> for TestRuntime<C, Con
         self.chain_state.begin_slot_hook(slot_data, working_set)
     }
 
-    fn end_slot_hook(&self, _working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>) {}
+    fn end_slot_hook(
+        &self,
+        _root_hash: [u8; 32],
+        _working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    ) {
+    }
 }
 
 impl<C: Context, Cond: ValidityCondition> BlobSelector for TestRuntime<C, Cond> {
diff --git a/module-system/module-implementations/sov-chain-state/src/hooks.rs b/module-system/module-implementations/sov-chain-state/src/hooks.rs
index 39b4429161..79820f44b9 100644
--- a/module-system/module-implementations/sov-chain-state/src/hooks.rs
+++ b/module-system/module-implementations/sov-chain-state/src/hooks.rs
@@ -63,5 +63,10 @@ impl<Ctx: Context, Cond: ValidityCondition> SlotHooks<Cond> for ChainState<Ctx,
         );
     }
 
-    fn end_slot_hook(&self, _working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>) {}
+    fn end_slot_hook(
+        &self,
+        _root_hash: [u8; 32],
+        _working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    ) {
+    }
 }
diff --git a/module-system/module-implementations/sov-evm/src/hooks.rs b/module-system/module-implementations/sov-evm/src/hooks.rs
new file mode 100644
index 0000000000..722d456bd8
--- /dev/null
+++ b/module-system/module-implementations/sov-evm/src/hooks.rs
@@ -0,0 +1,9 @@
+use sov_state::WorkingSet;
+
+use crate::Evm;
+
+impl<C: sov_modules_api::Context> Evm<C> {
+    pub fn end_slot_hook(&self, _root_hash: [u8; 32], _working_set: &mut WorkingSet<C::Storage>) {
+        // TODO implement block creation logic.
+    }
+}
diff --git a/module-system/module-implementations/sov-evm/src/lib.rs b/module-system/module-implementations/sov-evm/src/lib.rs
index 86cb57e205..726ab689aa 100644
--- a/module-system/module-implementations/sov-evm/src/lib.rs
+++ b/module-system/module-implementations/sov-evm/src/lib.rs
@@ -4,6 +4,8 @@ pub mod call;
 pub mod evm;
 #[cfg(feature = "experimental")]
 pub mod genesis;
+#[cfg(feature = "experimental")]
+pub mod hooks;
 #[cfg(feature = "native")]
 #[cfg(feature = "experimental")]
 pub mod query;
diff --git a/module-system/sov-modules-api/src/hooks.rs b/module-system/sov-modules-api/src/hooks.rs
index 460b19ac0f..d5eee73aa8 100644
--- a/module-system/sov-modules-api/src/hooks.rs
+++ b/module-system/sov-modules-api/src/hooks.rs
@@ -63,5 +63,9 @@ pub trait SlotHooks<Condition: ValidityCondition> {
         working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
     );
 
-    fn end_slot_hook(&self, working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>);
+    fn end_slot_hook(
+        &self,
+        root_hash: [u8; 32],
+        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    );
 }
diff --git a/module-system/sov-modules-stf-template/src/lib.rs b/module-system/sov-modules-stf-template/src/lib.rs
index 54bd73265a..20c164022b 100644
--- a/module-system/sov-modules-stf-template/src/lib.rs
+++ b/module-system/sov-modules-stf-template/src/lib.rs
@@ -99,8 +99,7 @@ where
             .expect("jellyfish merkle tree update must succeed");
 
         let mut working_set = WorkingSet::new(self.current_storage.clone());
-
-        self.runtime.end_slot_hook(&mut working_set);
+        self.runtime.end_slot_hook(root_hash, &mut working_set);
 
         (jmt::RootHash(root_hash), witness)
     }