feat: add associated types as tx hooks args

examples/demo-rollup/src/celestia_rollup.rs

@@ -4,7 +4,7 @@ use demo_stf::runtime::Runtime;
 use sov_celestia_adapter::verifier::{CelestiaSpec, CelestiaVerifier, RollupParams};
 use sov_celestia_adapter::{CelestiaConfig, CelestiaService};
 use sov_modules_api::default_context::{DefaultContext, ZkDefaultContext};
-use sov_modules_api::Spec;
+use sov_modules_api::{Address, Spec};
 use sov_modules_rollup_blueprint::{RollupBlueprint, WalletBlueprint};
 use sov_modules_stf_blueprint::kernels::basic::BasicKernel;
 use sov_modules_stf_blueprint::StfBlueprint;
@@ -57,12 +57,15 @@ impl RollupBlueprint for CelestiaDemoRollup {
         ledger_db: &sov_db::ledger_db::LedgerDB,
         da_service: &Self::DaService,
     ) -> Result<jsonrpsee::RpcModule<()>, anyhow::Error> {
+        // TODO set the sequencer address
+        let sequencer = Address::new([0; 32]);
+
         #[allow(unused_mut)]
         let mut rpc_methods = sov_modules_rollup_blueprint::register_rpc::<
             Self::NativeRuntime,
             Self::NativeContext,
             Self::DaService,
-        >(storage, ledger_db, da_service)?;
+        >(storage, ledger_db, da_service, sequencer)?;
 
         #[cfg(feature = "experimental")]
         crate::eth::register_ethereum::<Self::DaService>(

examples/demo-rollup/src/mock_rollup.rs

@@ -4,7 +4,7 @@ use demo_stf::runtime::Runtime;
 use sov_db::ledger_db::LedgerDB;
 use sov_mock_da::{MockDaConfig, MockDaService, MockDaSpec};
 use sov_modules_api::default_context::{DefaultContext, ZkDefaultContext};
-use sov_modules_api::Spec;
+use sov_modules_api::{Address, Spec};
 use sov_modules_rollup_blueprint::RollupBlueprint;
 use sov_modules_stf_blueprint::kernels::basic::BasicKernel;
 use sov_modules_stf_blueprint::StfBlueprint;
@@ -55,12 +55,15 @@ impl RollupBlueprint for MockDemoRollup {
         ledger_db: &LedgerDB,
         da_service: &Self::DaService,
     ) -> Result<jsonrpsee::RpcModule<()>, anyhow::Error> {
+        // TODO set the sequencer address
+        let sequencer = Address::new([0; 32]);
+
         #[allow(unused_mut)]
         let mut rpc_methods = sov_modules_rollup_blueprint::register_rpc::<
             Self::NativeRuntime,
             Self::NativeContext,
             Self::DaService,
-        >(storage, ledger_db, da_service)?;
+        >(storage, ledger_db, da_service, sequencer)?;
 
         #[cfg(feature = "experimental")]
         crate::eth::register_ethereum::<Self::DaService>(

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

@@ -1,7 +1,8 @@
+use sov_accounts::AccountsTxHook;
 use sov_modules_api::hooks::{ApplyBlobHooks, FinalizeHook, SlotHooks, TxHooks};
 use sov_modules_api::transaction::Transaction;
 use sov_modules_api::{AccessoryWorkingSet, Context, Spec, WorkingSet};
-use sov_modules_stf_blueprint::SequencerOutcome;
+use sov_modules_stf_blueprint::{RuntimeTxHook, SequencerOutcome};
 #[cfg(feature = "experimental")]
 use sov_rollup_interface::da::BlockHeaderTrait;
 use sov_rollup_interface::da::{BlobReaderTrait, DaSpec};
@@ -13,24 +14,30 @@ use crate::runtime::Runtime;
 
 impl<C: Context, Da: DaSpec> TxHooks for Runtime<C, Da> {
     type Context = C;
+    type PreArg = RuntimeTxHook<C>;
+    type PreResult = C;
 
     fn pre_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
         working_set: &mut WorkingSet<C>,
-    ) -> anyhow::Result<<Self::Context as Spec>::Address> {
-        // 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)
+        arg: RuntimeTxHook<C>,
+    ) -> anyhow::Result<C> {
+        let RuntimeTxHook { height, sequencer } = arg;
+        let AccountsTxHook { sender, sequencer } =
+            self.accounts
+                .pre_dispatch_tx_hook(tx, working_set, sequencer)?;
+        Ok(C::new(sender, sequencer, height))
     }
 
     fn post_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
+        ctx: &C,
         working_set: &mut WorkingSet<C>,
     ) -> 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, ctx, working_set)?;
+        Ok(())
     }
 }
 

examples/simple-nft-module/tests/nft_test.rs

@@ -17,6 +17,7 @@ fn genesis_and_mint() {
     let admin = generate_address("admin");
     let owner1 = generate_address("owner2");
     let owner2 = generate_address("owner2");
+    let sequencer = generate_address("sequencer");
     let config: NonFungibleTokenConfig<C> = NonFungibleTokenConfig {
         admin,
         owners: vec![(0, owner1)],
@@ -38,7 +39,7 @@ fn genesis_and_mint() {
 
     // Mint, anybody can mint
     let mint_message = CallMessage::Mint { id: 1 };
-    let owner2_context = C::new(owner2, 1);
+    let owner2_context = C::new(owner2, sequencer, 1);
     nft.call(mint_message.clone(), &owner2_context, &mut working_set)
         .expect("Minting failed");
 
@@ -61,9 +62,10 @@ fn genesis_and_mint() {
 fn transfer() {
     // Preparation
     let admin = generate_address("admin");
-    let admin_context = C::new(admin, 1);
+    let sequencer = generate_address("sequencer");
+    let admin_context = C::new(admin, sequencer, 1);
     let owner1 = generate_address("owner2");
-    let owner1_context = C::new(owner1, 1);
+    let owner1_context = C::new(owner1, sequencer, 1);
     let owner2 = generate_address("owner2");
     let config: NonFungibleTokenConfig<C> = NonFungibleTokenConfig {
         admin,
@@ -117,9 +119,10 @@ fn transfer() {
 fn burn() {
     // Preparation
     let admin = generate_address("admin");
-    let admin_context = C::new(admin, 1);
+    let sequencer = generate_address("sequencer");
+    let admin_context = C::new(admin, sequencer, 1);
     let owner1 = generate_address("owner2");
-    let owner1_context = C::new(owner1, 1);
+    let owner1_context = C::new(owner1, sequencer, 1);
     let config: NonFungibleTokenConfig<C> = NonFungibleTokenConfig {
         admin,
         owners: vec![(0, owner1)],

full-node/sov-sequencer/src/batch_builder.rs

@@ -50,6 +50,7 @@ pub struct FiFoStrictBatchBuilder<C: Context, R: DispatchCall<Context = C>> {
     runtime: R,
     max_batch_size_bytes: usize,
     current_storage: C::Storage,
+    sequencer: C::Address,
 }
 
 impl<C, R> FiFoStrictBatchBuilder<C, R>
@@ -63,13 +64,15 @@ where
         mempool_max_txs_count: usize,
         runtime: R,
         current_storage: C::Storage,
+        sequencer: C::Address,
     ) -> Self {
         Self {
             mempool: VecDeque::new(),
             mempool_max_txs_count,
             max_batch_size_bytes,
             runtime,
             current_storage,
+            sequencer,
         }
     }
 }
@@ -139,7 +142,7 @@ where
                 // TODO: Bug(!), because potential discrepancy. Should be resolved by https://github.com/Sovereign-Labs/sovereign-sdk/issues/434
                 let sender_address: C::Address = pooled.tx.pub_key().to_address();
                 // FIXME! This should use the correct height
-                let ctx = C::new(sender_address, 0);
+                let ctx = C::new(sender_address, self.sequencer.clone(), 0);
 
                 if let Err(error) = self.runtime.dispatch_call(msg, &mut working_set, &ctx) {
                     warn!(%error, tx = hex::encode(&pooled.raw), "Error during transaction dispatch");
@@ -182,7 +185,9 @@ mod tests {
     use sov_modules_api::default_signature::DefaultPublicKey;
     use sov_modules_api::macros::DefaultRuntime;
     use sov_modules_api::transaction::Transaction;
-    use sov_modules_api::{Context, DispatchCall, EncodeCall, Genesis, MessageCodec, PrivateKey};
+    use sov_modules_api::{
+        Address, Context, DispatchCall, EncodeCall, Genesis, MessageCodec, PrivateKey,
+    };
     use sov_rollup_interface::services::batch_builder::BatchBuilder;
     use sov_state::{DefaultStorageSpec, ProverStorage, Storage};
     use sov_value_setter::{CallMessage, ValueSetter, ValueSetterConfig};
@@ -246,11 +251,13 @@ mod tests {
     ) {
         let storage = ProverStorage::<DefaultStorageSpec>::with_path(tmpdir.path()).unwrap();
 
+        let sequencer = Address::from([0; 32]);
         let batch_builder = FiFoStrictBatchBuilder::new(
             batch_size_bytes,
             MAX_TX_POOL_SIZE,
             TestRuntime::<C>::default(),
             storage.clone(),
+            sequencer,
         );
         (batch_builder, storage)
     }

fuzz/fuzz_targets/accounts_call.rs

@@ -15,8 +15,8 @@ use sov_modules_api::{Context, Module, PrivateKey, Spec, WorkingSet};
 type C = DefaultContext;
 
 // Check well-formed calls
-fuzz_target!(|input: (u16, [u8; 32], Vec<DefaultPrivateKey>)| -> Corpus {
-    let (iterations, seed, keys) = input;
+fuzz_target!(|input: (u16, [u8; 32], [u8; 32], Vec<DefaultPrivateKey>)| -> Corpus {
+    let (iterations, seed, sequencer, keys) = input;
     if iterations < 1024 {
         // pointless to setup & run a small iterations count
         return Corpus::Reject;
@@ -42,6 +42,7 @@ fuzz_target!(|input: (u16, [u8; 32], Vec<DefaultPrivateKey>)| -> Corpus {
     let storage = <C as Spec>::Storage::with_path(tmpdir.path()).unwrap();
     let working_set = &mut WorkingSet::new(storage);
 
+    let sequencer = <C as Spec>::Address::from(sequencer);
     let config: AccountConfig<C> = keys.iter().map(|k| k.pub_key()).collect();
     let accounts: Accounts<C> = Accounts::default();
     accounts.genesis(&config, working_set).unwrap();
@@ -54,7 +55,7 @@ fuzz_target!(|input: (u16, [u8; 32], Vec<DefaultPrivateKey>)| -> Corpus {
     for i in 0..iterations {
         // we use slices for better select performance
         let sender = addresses.choose(rng).unwrap();
-        let context = C::new(*sender, i as u64);
+        let context = C::new(*sender, sequencer, i as u64);
 
         // clear previous state
         let previous = state.get(sender).unwrap().as_hex();

fuzz/fuzz_targets/bank_call.rs

@@ -9,12 +9,12 @@ use sov_state::ProverStorage;
 
 type C = DefaultContext;
 
-fuzz_target!(|input: (&[u8], [u8; 32])| {
-    let (data, sender) = input;
+fuzz_target!(|input: (&[u8], [u8; 32], [u8; 32])| {
+    let (data, sender, sequencer) = input;
     if let Ok(msgs) = serde_json::from_slice::<Vec<CallMessage<C>>>(data) {
         let tmpdir = tempfile::tempdir().unwrap();
         let mut working_set = WorkingSet::new(ProverStorage::with_path(tmpdir.path()).unwrap());
-        let ctx = C::new(sender.into(), 1);
+        let ctx = C::new(sender.into(), sequencer.into(), 1);
         let bank = Bank::default();
         for msg in msgs {
             bank.call(msg, &ctx, &mut working_set).ok();

module-system/module-implementations/examples/sov-accessory-state/tests/test.rs

@@ -18,7 +18,8 @@ fn test_accessory_value_setter() {
     let mut working_set_for_check: WorkingSet<DefaultContext> = WorkingSet::new(storage.clone());
 
     let admin = Address::from([1; 32]);
-    let context = DefaultContext::new(admin, 1);
+    let sequencer = Address::from([2; 32]);
+    let context = DefaultContext::new(admin, sequencer, 1);
 
     let module = AccessorySetter::<DefaultContext>::default();
 

module-system/module-implementations/examples/sov-value-setter/src/tests.rs

@@ -10,11 +10,12 @@ fn test_value_setter() {
     let tmpdir = tempfile::tempdir().unwrap();
     let mut working_set = WorkingSet::new(ProverStorage::with_path(tmpdir.path()).unwrap());
     let admin = Address::from([1; 32]);
+    let sequencer = Address::from([2; 32]);
     // Test Native-Context
     #[cfg(feature = "native")]
     {
         let config = ValueSetterConfig { admin };
-        let context = DefaultContext::new(admin, 1);
+        let context = DefaultContext::new(admin, sequencer, 1);
         test_value_setter_helper(context, &config, &mut working_set);
     }
 
@@ -23,7 +24,7 @@ fn test_value_setter() {
     // Test Zk-Context
     {
         let config = ValueSetterConfig { admin };
-        let zk_context = ZkDefaultContext::new(admin, 1);
+        let zk_context = ZkDefaultContext::new(admin, sequencer, 1);
         let mut zk_working_set = WorkingSet::with_witness(ZkStorage::new(), witness);
         test_value_setter_helper(zk_context, &config, &mut zk_working_set);
     }
@@ -63,6 +64,7 @@ fn test_value_setter_helper<C: Context>(
 #[test]
 fn test_err_on_sender_is_not_admin() {
     let sender = Address::from([1; 32]);
+    let sequencer = Address::from([2; 32]);
 
     let tmpdir = tempfile::tempdir().unwrap();
     let backing_store = ProverStorage::with_path(tmpdir.path()).unwrap();
@@ -75,7 +77,7 @@ fn test_err_on_sender_is_not_admin() {
         let config = ValueSetterConfig {
             admin: sender_not_admin,
         };
-        let context = DefaultContext::new(sender, 1);
+        let context = DefaultContext::new(sender, sequencer, 1);
         test_err_on_sender_is_not_admin_helper(context, &config, &mut native_working_set);
     }
     let (_, witness) = native_working_set.checkpoint().freeze();
@@ -86,7 +88,7 @@ fn test_err_on_sender_is_not_admin() {
             admin: sender_not_admin,
         };
         let zk_backing_store = ZkStorage::new();
-        let zk_context = ZkDefaultContext::new(sender, 1);
+        let zk_context = ZkDefaultContext::new(sender, sequencer, 1);
         let zk_working_set = &mut WorkingSet::with_witness(zk_backing_store, witness);
         test_err_on_sender_is_not_admin_helper(zk_context, &config, zk_working_set);
     }

module-system/module-implementations/examples/sov-vec-setter/tests/tests.rs

@@ -5,24 +5,25 @@ use sov_vec_setter::{CallMessage, VecSetter, VecSetterConfig};
 
 // rustfmt doesn't like long lines, but it's easier to read in this case.
 #[rustfmt::skip]
-fn test_cases() -> Vec<(Address, CallMessage, Option<Vec<u32>>)> {
+fn test_cases() -> Vec<(Address, Address, CallMessage, Option<Vec<u32>>)> {
     let admin = Address::from([1; 32]);
     let not_admin = Address::from([2; 32]);
+    let sequencer = Address::from([3; 32]);
 
     // (sender, call, expected vec contents or None if call should fail)
     vec![
-        (admin, CallMessage::PushValue(1), Some(vec![1])),
-        (admin, CallMessage::PushValue(2), Some(vec![1, 2])),
-        (admin, CallMessage::PopValue, Some(vec![1])),
-        (not_admin, CallMessage::PopValue, None),
-        (admin, CallMessage::PopValue, Some(vec![])),
-        (not_admin, CallMessage::SetValue { index: 0, value: 10 }, None),
-        (admin, CallMessage::SetValue { index: 0, value: 10 }, None),
-        (admin, CallMessage::PushValue(8), Some(vec![8])),
-        (admin, CallMessage::SetValue { index: 0, value: 10 }, Some(vec![10])),
-        (admin, CallMessage::PushValue(0), Some(vec![10, 0])),
-        (admin, CallMessage::SetAllValues(vec![11, 12]), Some(vec![11, 12])),
-        (not_admin, CallMessage::SetAllValues(vec![]), None),
+        (admin, sequencer, CallMessage::PushValue(1), Some(vec![1])),
+        (admin, sequencer, CallMessage::PushValue(2), Some(vec![1, 2])),
+        (admin, sequencer, CallMessage::PopValue, Some(vec![1])),
+        (not_admin, sequencer, CallMessage::PopValue, None),
+        (admin, sequencer, CallMessage::PopValue, Some(vec![])),
+        (not_admin, sequencer, CallMessage::SetValue { index: 0, value: 10 }, None),
+        (admin, sequencer, CallMessage::SetValue { index: 0, value: 10 }, None),
+        (admin, sequencer, CallMessage::PushValue(8), Some(vec![8])),
+        (admin, sequencer, CallMessage::SetValue { index: 0, value: 10 }, Some(vec![10])),
+        (admin, sequencer, CallMessage::PushValue(0), Some(vec![10, 0])),
+        (admin, sequencer, CallMessage::SetAllValues(vec![11, 12]), Some(vec![11, 12])),
+        (not_admin, sequencer, CallMessage::SetAllValues(vec![]), None),
     ]
 }
 
@@ -39,8 +40,8 @@ fn test_vec_setter_calls() {
     let vec_setter = VecSetter::default();
     vec_setter.genesis(&config, &mut working_set).unwrap();
 
-    for (sender, call, expected_contents) in test_cases().iter().cloned() {
-        let context = DefaultContext::new(sender, 1);
+    for (sender, sequencer, call, expected_contents) in test_cases().iter().cloned() {
+        let context = DefaultContext::new(sender, sequencer, 1);
 
         let call_result = vec_setter.call(call, &context, &mut working_set);
 

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

@@ -7,7 +7,7 @@ use sov_modules_api::{
     AccessoryWorkingSet, BlobReaderTrait, Context, DaSpec, DispatchCall, Genesis, MessageCodec,
     PublicKey, Spec,
 };
-use sov_modules_stf_blueprint::{Runtime, SequencerOutcome};
+use sov_modules_stf_blueprint::{Runtime, RuntimeTxHook, SequencerOutcome};
 use sov_state::Storage;
 use sov_value_setter::{ValueSetter, ValueSetterConfig};
 
@@ -20,18 +20,25 @@ pub(crate) struct TestRuntime<C: Context, Da: DaSpec> {
 
 impl<C: Context, Da: DaSpec> TxHooks for TestRuntime<C, Da> {
     type Context = C;
+    type PreArg = RuntimeTxHook<C>;
+    type PreResult = C;
 
     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())
+        arg: RuntimeTxHook<C>,
+    ) -> anyhow::Result<C> {
+        let RuntimeTxHook { height, sequencer } = arg;
+        let sender = tx.pub_key().to_address();
+        let sequencer = sequencer.to_address();
+        Ok(C::new(sender, sequencer, height))
     }
 
     fn post_dispatch_tx_hook(
         &self,
         _tx: &Transaction<Self::Context>,
+        _ctx: &C,
         _working_set: &mut sov_modules_api::WorkingSet<C>,
     ) -> anyhow::Result<()> {
         Ok(())