feat: add associated types as tx hooks args

examples/demo-rollup/stf/src/hooks_impl.rs

@@ -13,24 +13,34 @@ use crate::runtime::Runtime;
 
 impl<C: Context, Da: DaSpec> TxHooks for Runtime<C, Da> {
     type Context = C;
+    type PreArg = u64;
+    type PreResult = C;
+    type PostArg = ();
+    type PostResult = ();
 
     fn pre_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
         working_set: &mut WorkingSet<C>,
-    ) -> anyhow::Result<<Self::Context as Spec>::Address> {
+        height: u64,
+    ) -> anyhow::Result<C> {
         // Before executing a transaction, retrieve the sender's address from the accounts module
         // and check the nonce
-        self.accounts.pre_dispatch_tx_hook(tx, working_set)
+        let sender_address = self.accounts.pre_dispatch_tx_hook(tx, working_set, ())?;
+        let ctx = C::new(sender_address, height);
+
+        Ok(ctx)
     }
 
     fn post_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
         working_set: &mut WorkingSet<C>,
+        _arg: (),
     ) -> anyhow::Result<()> {
         // After executing each transaction, update the nonce
-        self.accounts.post_dispatch_tx_hook(tx, working_set)
+        self.accounts.post_dispatch_tx_hook(tx, working_set, ())?;
+        Ok(())
     }
 }
 

module-system/module-implementations/integration-tests/src/chain_state/helpers.rs

@@ -20,19 +20,27 @@ pub(crate) struct TestRuntime<C: Context, Da: DaSpec> {
 
 impl<C: Context, Da: DaSpec> TxHooks for TestRuntime<C, Da> {
     type Context = C;
+    type PreArg = u64;
+    type PreResult = C;
+    type PostArg = ();
+    type PostResult = ();
 
     fn pre_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
         _working_set: &mut sov_modules_api::WorkingSet<C>,
-    ) -> anyhow::Result<<Self::Context as Spec>::Address> {
-        Ok(tx.pub_key().to_address())
+        height: u64,
+    ) -> anyhow::Result<C> {
+        let sender = tx.pub_key().to_address();
+        let ctx = C::new(sender, height);
+        Ok(ctx)
     }
 
     fn post_dispatch_tx_hook(
         &self,
         _tx: &Transaction<Self::Context>,
         _working_set: &mut sov_modules_api::WorkingSet<C>,
+        _arg: (),
     ) -> anyhow::Result<()> {
         Ok(())
     }

module-system/module-implementations/sov-accounts/src/hooks.rs

@@ -1,17 +1,22 @@
 use sov_modules_api::hooks::TxHooks;
 use sov_modules_api::transaction::Transaction;
-use sov_modules_api::{Context, Spec, StateMapAccessor, WorkingSet};
+use sov_modules_api::{Context, StateMapAccessor, WorkingSet};
 
 use crate::Accounts;
 
 impl<C: Context> TxHooks for Accounts<C> {
     type Context = C;
+    type PreArg = ();
+    type PreResult = C::Address;
+    type PostArg = ();
+    type PostResult = ();
 
     fn pre_dispatch_tx_hook(
         &self,
         tx: &Transaction<C>,
         working_set: &mut WorkingSet<C>,
-    ) -> anyhow::Result<<Self::Context as Spec>::Address> {
+        _arg: (),
+    ) -> anyhow::Result<C::Address> {
         let pub_key = tx.pub_key();
 
         let account = match self.accounts.get(pub_key, working_set) {
@@ -32,6 +37,7 @@ impl<C: Context> TxHooks for Accounts<C> {
         &self,
         tx: &Transaction<Self::Context>,
         working_set: &mut WorkingSet<C>,
+        _arg: (),
     ) -> anyhow::Result<()> {
         let mut account = self.accounts.get_or_err(tx.pub_key(), working_set)?;
         account.nonce += 1;

module-system/sov-modules-api/src/hooks.rs

@@ -4,26 +4,33 @@ use sov_rollup_interface::da::{BlobReaderTrait, DaSpec};
 use crate::transaction::Transaction;
 
 /// Hooks that execute within the `StateTransitionFunction::apply_blob` function for each processed transaction.
+///
+/// The arguments consist of expected concretely implemented associated types for the hooks. At
+/// runtime, compatible implementations are selected and utilized by the system to construct its
+/// setup procedures and define post-execution routines.
 pub trait TxHooks {
     type Context: Context;
+    type PreArg;
+    type PreResult;
+    type PostArg;
+    type PostResult;
 
     /// Runs just before a transaction is dispatched to an appropriate module.
-    /// TODO: Why does it return address?
-    /// Does it implies that it should do signature verification.
-    /// Can other code rely on that assumption?
     fn pre_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
         working_set: &mut WorkingSet<Self::Context>,
-    ) -> anyhow::Result<<Self::Context as Spec>::Address>;
+        arg: Self::PreArg,
+    ) -> anyhow::Result<Self::PreResult>;
 
     /// Runs after the tx is dispatched to an appropriate module.
     /// IF this hook returns error rollup panics
     fn post_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
         working_set: &mut WorkingSet<Self::Context>,
-    ) -> anyhow::Result<()>;
+        arg: Self::PostArg,
+    ) -> anyhow::Result<Self::PostResult>;
 }
 
 /// Hooks related to the Sequencer functionality.

module-system/sov-modules-stf-blueprint/src/lib.rs

@@ -22,10 +22,13 @@ use tracing::info;
 pub use tx_verifier::RawTx;
 
 /// This trait has to be implemented by a runtime in order to be used in `StfBlueprint`.
+///
+/// The `TxHooks` implementation sets up a transaction context based on the height at which it is
+/// to be executed.
 pub trait Runtime<C: Context, Da: DaSpec>:
     DispatchCall<Context = C>
     + Genesis<Context = C, Config = Self::GenesisConfig>
-    + TxHooks<Context = C>
+    + TxHooks<Context = C, PreArg = u64, PreResult = C, PostArg = ()>
     + SlotHooks<Da, Context = C>
     + FinalizeHook<Da, Context = C>
     + ApplyBlobHooks<

module-system/sov-modules-stf-blueprint/src/stf_blueprint.rs

@@ -179,7 +179,10 @@ where
             txs.into_iter().zip(messages.into_iter())
         {
             // Pre dispatch hook
-            let sender_address = match self.runtime.pre_dispatch_tx_hook(&tx, &mut batch_workspace)
+            let height = 1;
+            let ctx = match self
+                .runtime
+                .pre_dispatch_tx_hook(&tx, &mut batch_workspace, height)
             {
                 Ok(verified_tx) => verified_tx,
                 Err(e) => {
@@ -200,7 +203,6 @@ where
             // Commit changes after pre_dispatch_tx_hook
             batch_workspace = batch_workspace.checkpoint().to_revertable();
 
-            let ctx = C::new(sender_address.clone(), 1);
             let tx_result = self.runtime.dispatch_call(msg, &mut batch_workspace, &ctx);
 
             let events = batch_workspace.take_events();
@@ -233,7 +235,7 @@ where
 
             // TODO: `panic` will be covered in https://github.com/Sovereign-Labs/sovereign-sdk/issues/421
             self.runtime
-                .post_dispatch_tx_hook(&tx, &mut batch_workspace)
+                .post_dispatch_tx_hook(&tx, &mut batch_workspace, ())
                 .expect("Impossible happened: error in post_dispatch_tx_hook");
         }