feat: add context to working set

examples/demo-nft-module/README.md

@@ -182,7 +182,7 @@ impl<C: sov_modules_api::Context> Module for NonFungibleToken<C> {
     fn genesis(
         &self,
         _config: &Self::Config,
-        _working_set: &mut WorkingSet<C::Storage>,
+        _working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<(), Error> {
         Ok(())
     }
@@ -191,7 +191,7 @@ impl<C: sov_modules_api::Context> Module for NonFungibleToken<C> {
         &self,
         _msg: Self::CallMessage,
         _context: &Self::Context,
-        _working_set: &mut WorkingSet<C::Storage>,
+        _working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<sov_modules_api::CallResponse, Error> {
         Ok(sov_modules_api::CallResponse::default())
     }
@@ -219,7 +219,7 @@ impl<C: sov_modules_api::Context> sov_modules_api::Module for NonFungibleToken<C
     fn genesis(
         &self,
         config: &Self::Config,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> Result<(), Error> {
         Ok(self.init_module(config, working_set)?)
     }
@@ -230,7 +230,7 @@ impl<C: sov_modules_api::Context> NonFungibleToken<C> {
     pub(crate) fn init_module(
         &self,
         config: &<Self as sov_modules_api::Module>::Config,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<()> {
         self.admin.set(&config.admin, working_set);
         for (id, owner) in config.owners.iter() {
@@ -256,7 +256,7 @@ impl<C: sov_modules_api::Context> NonFungibleToken<C> {
         &self,
         id: u64,
         context: &C,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<sov_modules_api::CallResponse> {
         if self.owners.get(&id, working_set).is_some() {
             bail!("Token with id {} already exists", id);
@@ -273,7 +273,7 @@ impl<C: sov_modules_api::Context> NonFungibleToken<C> {
         id: u64,
         to: C::Address,
         context: &C,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<sov_modules_api::CallResponse> {
         let token_owner = match self.owners.get(&id, working_set) {
             None => {
@@ -296,7 +296,7 @@ impl<C: sov_modules_api::Context> NonFungibleToken<C> {
         &self,
         id: u64,
         context: &C,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<sov_modules_api::CallResponse> {
         let token_owner = match self.owners.get(&id, working_set) {
             None => {
@@ -326,7 +326,7 @@ impl<C: sov_modules_api::Context> sov_modules_api::Module for NonFungibleToken<C
         &self,
         msg: Self::CallMessage,
         context: &Self::Context,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> Result<sov_modules_api::CallResponse, Error> {
         let call_result = match msg {
             CallMessage::Mint { id } => self.mint(id, context, working_set),
@@ -361,7 +361,7 @@ impl<C: sov_modules_api::Context> NonFungibleToken<C> {
     pub fn get_owner(
         &self,
         token_id: u64,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> RpcResult<OwnerResponse<C>> {
         Ok(OwnerResponse {
             owner: self.owners.get(&token_id, working_set),

examples/demo-nft-module/src/call.rs

@@ -35,7 +35,7 @@ impl<C: Context> NonFungibleToken<C> {
         &self,
         id: u64,
         context: &C,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> Result<CallResponse> {
         if self.owners.get(&id, working_set).is_some() {
             bail!("Token with id {} already exists", id);
@@ -52,7 +52,7 @@ impl<C: Context> NonFungibleToken<C> {
         id: u64,
         to: C::Address,
         context: &C,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> Result<CallResponse> {
         let token_owner = match self.owners.get(&id, working_set) {
             None => {
@@ -75,7 +75,7 @@ impl<C: Context> NonFungibleToken<C> {
         &self,
         id: u64,
         context: &C,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> Result<CallResponse> {
         let token_owner = match self.owners.get(&id, working_set) {
             None => {

examples/demo-nft-module/src/genesis.rs

@@ -7,7 +7,7 @@ impl<C: Context> NonFungibleToken<C> {
     pub(crate) fn init_module(
         &self,
         config: &<Self as sov_modules_api::Module>::Config,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> Result<()> {
         self.admin.set(&config.admin, working_set);
         for (id, owner) in config.owners.iter() {

examples/demo-nft-module/src/lib.rs

@@ -43,19 +43,15 @@ impl<C: Context> Module for NonFungibleToken<C> {
 
     type CallMessage = CallMessage<C>;
 
-    fn genesis(
-        &self,
-        config: &Self::Config,
-        working_set: &mut WorkingSet<C::Storage>,
-    ) -> Result<(), Error> {
+    fn genesis(&self, config: &Self::Config, working_set: &mut WorkingSet<C>) -> Result<(), Error> {
         Ok(self.init_module(config, working_set)?)
     }
 
     fn call(
         &self,
         msg: Self::CallMessage,
         context: &Self::Context,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> Result<CallResponse, Error> {
         let call_result = match msg {
             CallMessage::Mint { id } => self.mint(id, context, working_set),

examples/demo-nft-module/src/query.rs

@@ -18,7 +18,7 @@ impl<C: Context> NonFungibleToken<C> {
     pub fn get_owner(
         &self,
         token_id: u64,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> RpcResult<OwnerResponse<C>> {
         Ok(OwnerResponse {
             owner: self.owners.get(&token_id, working_set),

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

@@ -70,7 +70,8 @@ fn transfer() {
         owners: vec![(0, admin), (1, owner1), (2, owner2)],
     };
     let tmpdir = tempfile::tempdir().unwrap();
-    let mut working_set = WorkingSet::new(ProverStorage::with_path(tmpdir.path()).unwrap());
+    let mut working_set: WorkingSet<C> =
+        WorkingSet::new(ProverStorage::with_path(tmpdir.path()).unwrap());
     let nft = NonFungibleToken::default();
     nft.genesis(&config, &mut working_set).unwrap();
 
@@ -83,11 +84,10 @@ fn transfer() {
     let error_message = transfer_attempt.err().unwrap().to_string();
     assert_eq!("Only token owner can transfer token", error_message);
 
-    let query_token_owner =
-        |token_id: u64, working_set: &mut WorkingSet<Storage>| -> Option<Address> {
-            let query: OwnerResponse<C> = nft.get_owner(token_id, working_set).unwrap();
-            query.owner
-        };
+    let query_token_owner = |token_id: u64, working_set: &mut WorkingSet<C>| -> Option<Address> {
+        let query: OwnerResponse<C> = nft.get_owner(token_id, working_set).unwrap();
+        query.owner
+    };
 
     // Normal transfer
     let token1_owner = query_token_owner(1, &mut working_set);
@@ -126,7 +126,8 @@ fn burn() {
     };
 
     let tmpdir = tempfile::tempdir().unwrap();
-    let mut working_set = WorkingSet::new(ProverStorage::with_path(tmpdir.path()).unwrap());
+    let mut working_set: WorkingSet<C> =
+        WorkingSet::new(ProverStorage::with_path(tmpdir.path()).unwrap());
     let nft = NonFungibleToken::default();
     nft.genesis(&config, &mut working_set).unwrap();
 

examples/demo-stf/README.md

@@ -115,15 +115,15 @@ impl<C: Context> TxHooks for Runtime<C> {
     fn pre_dispatch_tx_hook(
         &self,
         tx: Transaction<Self::Context>,
-        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<<Self::Context as Spec>::Address> {
         self.accounts.pre_dispatch_tx_hook(tx, working_set)
     }
 
     fn post_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
-        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<()> {
         self.accounts.post_dispatch_tx_hook(tx, working_set)
     }
@@ -137,15 +137,15 @@ impl<C: Context> ApplyBlobHooks for Runtime<C> {
     fn lock_sequencer_bond(
         &self,
         sequencer: &[u8],
-        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<()> {
         self.sequencer.lock_sequencer_bond(sequencer, working_set)
     }
 
     fn reward_sequencer(
         &self,
         amount: u64,
-        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<()> {
         self.sequencer.reward_sequencer(amount, working_set)
     }

examples/demo-stf/src/hooks_impl.rs

@@ -16,15 +16,15 @@ impl<C: Context, Da: DaSpec> TxHooks for Runtime<C, Da> {
     fn pre_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
-        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<<Self::Context as Spec>::Address> {
         self.accounts.pre_dispatch_tx_hook(tx, working_set)
     }
 
     fn post_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
-        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<()> {
         self.accounts.post_dispatch_tx_hook(tx, working_set)
     }
@@ -38,15 +38,15 @@ impl<C: Context, Da: DaSpec> ApplyBlobHooks<Da::BlobTransaction> for Runtime<C,
     fn begin_blob_hook(
         &self,
         blob: &mut Da::BlobTransaction,
-        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<()> {
         self.sequencer_registry.begin_blob_hook(blob, working_set)
     }
 
     fn end_blob_hook(
         &self,
         result: Self::BlobResult,
-        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<()> {
         match result {
             SequencerOutcome::Rewarded(_reward) => {
@@ -82,9 +82,7 @@ impl<C: Context, Da: DaSpec> SlotHooks<Da> for Runtime<C, Da> {
         &self,
         #[allow(unused_variables)] slot_header: &Da::BlockHeader,
         #[allow(unused_variables)] validity_condition: &Da::ValidityCondition,
-        #[allow(unused_variables)] working_set: &mut sov_modules_api::WorkingSet<
-            <Self::Context as Spec>::Storage,
-        >,
+        #[allow(unused_variables)] working_set: &mut sov_modules_api::WorkingSet<C>,
     ) {
         #[cfg(feature = "experimental")]
         self.evm
@@ -93,9 +91,7 @@ impl<C: Context, Da: DaSpec> SlotHooks<Da> for Runtime<C, Da> {
 
     fn end_slot_hook(
         &self,
-        #[allow(unused_variables)] working_set: &mut sov_modules_api::WorkingSet<
-            <Self::Context as Spec>::Storage,
-        >,
+        #[allow(unused_variables)] working_set: &mut sov_modules_api::WorkingSet<C>,
     ) {
         #[cfg(feature = "experimental")]
         self.evm.end_slot_hook(working_set);
@@ -108,9 +104,7 @@ impl<C: Context, Da: sov_modules_api::DaSpec> FinalizeHook<Da> for Runtime<C, Da
     fn finalize_slot_hook(
         &self,
         #[allow(unused_variables)] root_hash: [u8; 32],
-        #[allow(unused_variables)] accesorry_working_set: &mut AccessoryWorkingSet<
-            <Self::Context as Spec>::Storage,
-        >,
+        #[allow(unused_variables)] accesorry_working_set: &mut AccessoryWorkingSet<C>,
     ) {
         #[cfg(feature = "experimental")]
         self.evm

examples/demo-stf/src/runtime.rs

@@ -15,7 +15,7 @@ pub use sov_modules_api::default_context::DefaultContext;
 use sov_modules_api::macros::DefaultRuntime;
 #[cfg(feature = "native")]
 use sov_modules_api::macros::{expose_rpc, CliWallet};
-use sov_modules_api::{Context, DispatchCall, Genesis, MessageCodec, Spec};
+use sov_modules_api::{Context, DispatchCall, Genesis, MessageCodec};
 use sov_rollup_interface::da::DaSpec;
 #[cfg(feature = "native")]
 use sov_sequencer_registry::{SequencerRegistryRpcImpl, SequencerRegistryRpcServer};
@@ -110,7 +110,7 @@ impl<C: Context, Da: DaSpec> BlobSelector<Da> for Runtime<C, Da> {
     fn get_blobs_for_this_slot<'a, I>(
         &self,
         current_blobs: I,
-        working_set: &mut sov_modules_api::WorkingSet<<Self::Context as Spec>::Storage>,
+        working_set: &mut sov_modules_api::WorkingSet<C>,
     ) -> anyhow::Result<Vec<BlobRefOrOwned<'a, Da::BlobTransaction>>>
     where
         I: IntoIterator<Item = &'a mut Da::BlobTransaction>,

module-system/README.md

@@ -51,7 +51,7 @@ accessible to other modules, but cannot be directly invoked by other users. A go
 
 ```rust
 impl<C: Context> Bank<C> {
-    pub fn transfer_from(&self, from: &C::Address, to: &C::Address, coins: Coins, working_set: &mut WorkingSet<C::Storage>) {
+    pub fn transfer_from(&self, from: &C::Address, to: &C::Address, coins: Coins, working_set: &mut WorkingSet<C>) {
         // Implementation elided...
     }
 }
@@ -75,7 +75,7 @@ tells the `call` function which inner method of the module to invoke. So a typic
 ```rust
 impl<C: sov_modules_api::Context> sov_modules_api::Module for Bank<C> {
 	// Several definitions elided here ...
-    fn call(&self, msg: Self::CallMessage, context: &Self::Context, working_set: &mut WorkingSet<C::Storage>) {
+    fn call(&self, msg: Self::CallMessage, context: &Self::Context, working_set: &mut WorkingSet<C>) {
         match msg {
             CallMessage::CreateToken {
                 token_name,
@@ -101,7 +101,7 @@ impl<C: sov_modules_api::Context> Bank<C> {
         &self,
         user_address: C::Address,
         token_address: C::Address,
-        working_set: &mut WorkingSet<C::Storage>,
+        working_set: &mut WorkingSet<C>,
     ) -> RpcResult<BalanceResponse> {
         Ok(BalanceResponse {
             amount: self.get_balance_of(user_address, token_address, working_set),