From 134fbc037ef104d04790e1709d6032d44da85137 Mon Sep 17 00:00:00 2001
From: Tiago Carvalho <sugoiuguu@tfwno.gf>
Date: Fri, 24 May 2024 10:47:01 +0100
Subject: [PATCH] Pass mutable refs to wasm host env

---
 crates/namada/src/vm/host_env.rs      | 119 ++++++++++++++------------
 crates/namada/src/vm/wasm/host_env.rs |  60 +++++++------
 crates/tests/src/vm_host_env/tx.rs    |  56 ++++++------
 crates/tests/src/vm_host_env/vp.rs    |   8 +-
 4 files changed, 135 insertions(+), 108 deletions(-)

diff --git a/crates/namada/src/vm/host_env.rs b/crates/namada/src/vm/host_env.rs
index af5adc4ce3..5a31ec09fe 100644
--- a/crates/namada/src/vm/host_env.rs
+++ b/crates/namada/src/vm/host_env.rs
@@ -587,7 +587,7 @@ where
 
 /// Add a gas cost incured in a transaction
 pub fn tx_charge_gas<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     used_gas: u64,
 ) -> TxResult<()>
 where
@@ -611,7 +611,7 @@ where
 
 /// Called from VP wasm to request to use the given gas amount
 pub fn vp_charge_gas<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     used_gas: u64,
 ) -> vp_host_fns::EnvResult<()>
 where
@@ -628,7 +628,7 @@ where
 /// Storage `has_key` function exposed to the wasm VM Tx environment. It will
 /// try to check the write log first and if no entry found then the storage.
 pub fn tx_has_key<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     key_ptr: u64,
     key_len: u64,
 ) -> TxResult<i64>
@@ -660,7 +660,7 @@ where
 /// Returns `-1` when the key is not present, or the length of the data when
 /// the key is present (the length may be `0`).
 pub fn tx_read<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     key_ptr: u64,
     key_len: u64,
 ) -> TxResult<i64>
@@ -703,7 +703,7 @@ where
 /// Returns `-1` when the key is not present, or the length of the data when
 /// the key is present (the length may be `0`).
 pub fn tx_read_temp<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     key_ptr: u64,
     key_len: u64,
 ) -> TxResult<i64>
@@ -749,7 +749,7 @@ where
 /// any) back to the guest, the second step reads the value from cache into a
 /// pre-allocated buffer with the obtained size.
 pub fn tx_result_buffer<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     result_ptr: u64,
 ) -> TxResult<()>
 where
@@ -773,7 +773,7 @@ where
 /// It will try to get an iterator from the storage and return the corresponding
 /// ID of the iterator, ordered by storage keys.
 pub fn tx_iter_prefix<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     prefix_ptr: u64,
     prefix_len: u64,
 ) -> TxResult<u64>
@@ -813,7 +813,7 @@ where
 /// Returns `-1` when the key is not present, or the length of the data when
 /// the key is present (the length may be `0`).
 pub fn tx_iter_next<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     iter_id: u64,
 ) -> TxResult<i64>
 where
@@ -824,7 +824,11 @@ where
 {
     tracing::debug!("tx_iter_next iter_id {}", iter_id,);
 
-    let state = env.state();
+    // NB: an env clone is required to avoid an immutable
+    // borrow while a mutable borrow is taking place
+    let env2 = env.clone();
+    let state = env2.state();
+
     let iterators = unsafe { env.ctx.iterators.get_mut() };
     let iter_id = PrefixIteratorId::new(iter_id);
     while let Some((key, val, iter_gas)) = iterators.next(iter_id) {
@@ -878,7 +882,7 @@ where
 /// Storage write function exposed to the wasm VM Tx environment. The given
 /// key/value will be written to the write log.
 pub fn tx_write<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     key_ptr: u64,
     key_len: u64,
     val_ptr: u64,
@@ -920,7 +924,7 @@ where
 /// given key/value will be written only to the write log. It will be never
 /// written to the storage.
 pub fn tx_write_temp<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     key_ptr: u64,
     key_len: u64,
     val_ptr: u64,
@@ -958,7 +962,7 @@ where
 }
 
 fn check_address_existence<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     key: &Key,
 ) -> TxResult<()>
 where
@@ -1005,7 +1009,7 @@ where
 /// Storage delete function exposed to the wasm VM Tx environment. The given
 /// key/value will be written as deleted to the write log.
 pub fn tx_delete<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     key_ptr: u64,
     key_len: u64,
 ) -> TxResult<()>
@@ -1035,7 +1039,7 @@ where
 /// Expose the functionality to emit events to the wasm VM's Tx environment.
 /// An emitted event will land in the write log.
 pub fn tx_emit_event<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     event_ptr: u64,
     event_len: u64,
 ) -> TxResult<()>
@@ -1062,7 +1066,7 @@ where
 
 /// Expose the functionality to query events from the wasm VM's Tx environment.
 pub fn tx_get_events<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     event_type_ptr: u64,
     event_type_len: u64,
 ) -> TxResult<i64>
@@ -1103,7 +1107,7 @@ where
 /// Returns `-1` when the key is not present, or the length of the data when
 /// the key is present (the length may be `0`).
 pub fn vp_read_pre<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     key_ptr: u64,
     key_len: u64,
 ) -> vp_host_fns::EnvResult<i64>
@@ -1153,7 +1157,7 @@ where
 /// Returns `-1` when the key is not present, or the length of the data when
 /// the key is present (the length may be `0`).
 pub fn vp_read_post<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     key_ptr: u64,
     key_len: u64,
 ) -> vp_host_fns::EnvResult<i64>
@@ -1198,7 +1202,7 @@ where
 /// Returns `-1` when the key is not present, or the length of the data when
 /// the key is present (the length may be `0`).
 pub fn vp_read_temp<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     key_ptr: u64,
     key_len: u64,
 ) -> vp_host_fns::EnvResult<i64>
@@ -1246,7 +1250,7 @@ where
 /// any) back to the guest, the second step reads the value from cache into a
 /// pre-allocated buffer with the obtained size.
 pub fn vp_result_buffer<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     result_ptr: u64,
 ) -> vp_host_fns::EnvResult<()>
 where
@@ -1271,7 +1275,7 @@ where
 /// Storage `has_key` in prior state (before tx execution) function exposed to
 /// the wasm VM VP environment. It will try to read from the storage.
 pub fn vp_has_key_pre<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     key_ptr: u64,
     key_len: u64,
 ) -> vp_host_fns::EnvResult<i64>
@@ -1302,7 +1306,7 @@ where
 /// to the wasm VM VP environment. It will try to check the write log first and
 /// if no entry found then the storage.
 pub fn vp_has_key_post<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     key_ptr: u64,
     key_len: u64,
 ) -> vp_host_fns::EnvResult<i64>
@@ -1334,7 +1338,7 @@ where
 /// the storage and return the corresponding ID of the iterator, ordered by
 /// storage keys.
 pub fn vp_iter_prefix_pre<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     prefix_ptr: u64,
     prefix_len: u64,
 ) -> vp_host_fns::EnvResult<u64>
@@ -1373,7 +1377,7 @@ where
 /// the storage and return the corresponding ID of the iterator, ordered by
 /// storage keys.
 pub fn vp_iter_prefix_post<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     prefix_ptr: u64,
     prefix_len: u64,
 ) -> vp_host_fns::EnvResult<u64>
@@ -1414,7 +1418,7 @@ where
 /// Returns `-1` when the key is not present, or the length of the data when
 /// the key is present (the length may be `0`).
 pub fn vp_iter_next<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     iter_id: u64,
 ) -> vp_host_fns::EnvResult<i64>
 where
@@ -1447,7 +1451,7 @@ where
 
 /// Verifier insertion function exposed to the wasm VM Tx environment.
 pub fn tx_insert_verifier<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     addr_ptr: u64,
     addr_len: u64,
 ) -> TxResult<()>
@@ -1481,7 +1485,7 @@ where
 
 /// Update a validity predicate function exposed to the wasm VM Tx environment
 pub fn tx_update_validity_predicate<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     addr_ptr: u64,
     addr_len: u64,
     code_hash_ptr: u64,
@@ -1532,7 +1536,7 @@ where
 /// Initialize a new account established address.
 #[allow(clippy::too_many_arguments)]
 pub fn tx_init_account<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     code_hash_ptr: u64,
     code_hash_len: u64,
     code_tag_ptr: u64,
@@ -1588,7 +1592,7 @@ where
 
 /// Getting the chain ID function exposed to the wasm VM Tx environment.
 pub fn tx_get_chain_id<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     result_ptr: u64,
 ) -> TxResult<()>
 where
@@ -1611,7 +1615,7 @@ where
 /// environment. The height is that of the block to which the current
 /// transaction is being applied.
 pub fn tx_get_block_height<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
 ) -> TxResult<u64>
 where
     MEM: VmMemory,
@@ -1629,7 +1633,7 @@ where
 /// environment. The index is that of the transaction being applied
 /// in the current block.
 pub fn tx_get_tx_index<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
 ) -> TxResult<u32>
 where
     MEM: VmMemory,
@@ -1651,7 +1655,7 @@ where
 /// environment. The height is that of the block to which the current
 /// transaction is being applied.
 pub fn vp_get_tx_index<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
 ) -> vp_host_fns::EnvResult<u32>
 where
     MEM: VmMemory,
@@ -1670,7 +1674,7 @@ where
 /// environment. The epoch is that of the block to which the current
 /// transaction is being applied.
 pub fn tx_get_block_epoch<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
 ) -> TxResult<u64>
 where
     MEM: VmMemory,
@@ -1686,7 +1690,7 @@ where
 
 /// Get predecessor epochs function exposed to the wasm VM Tx environment.
 pub fn tx_get_pred_epochs<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
 ) -> TxResult<i64>
 where
     MEM: VmMemory,
@@ -1713,7 +1717,7 @@ where
 
 /// Get the native token's address
 pub fn tx_get_native_token<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     result_ptr: u64,
 ) -> TxResult<()>
 where
@@ -1741,7 +1745,7 @@ where
 
 /// Getting the block header function exposed to the wasm VM Tx environment.
 pub fn tx_get_block_header<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     height: u64,
 ) -> TxResult<i64>
 where
@@ -1773,7 +1777,7 @@ where
 
 /// Getting the chain ID function exposed to the wasm VM VP environment.
 pub fn vp_get_chain_id<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     result_ptr: u64,
 ) -> vp_host_fns::EnvResult<()>
 where
@@ -1797,7 +1801,7 @@ where
 /// environment. The height is that of the block to which the current
 /// transaction is being applied.
 pub fn vp_get_block_height<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
 ) -> vp_host_fns::EnvResult<u64>
 where
     MEM: VmMemory,
@@ -1814,7 +1818,7 @@ where
 
 /// Getting the block header function exposed to the wasm VM VP environment.
 pub fn vp_get_block_header<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     height: u64,
 ) -> vp_host_fns::EnvResult<i64>
 where
@@ -1847,7 +1851,7 @@ where
 
 /// Getting the transaction hash function exposed to the wasm VM VP environment.
 pub fn vp_get_tx_code_hash<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     result_ptr: u64,
 ) -> vp_host_fns::EnvResult<()>
 where
@@ -1880,7 +1884,7 @@ where
 /// environment. The epoch is that of the block to which the current
 /// transaction is being applied.
 pub fn vp_get_block_epoch<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
 ) -> vp_host_fns::EnvResult<u64>
 where
     MEM: VmMemory,
@@ -1897,7 +1901,7 @@ where
 
 /// Get predecessor epochs function exposed to the wasm VM VP environment.
 pub fn vp_get_pred_epochs<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
 ) -> vp_host_fns::EnvResult<i64>
 where
     MEM: VmMemory,
@@ -1921,7 +1925,7 @@ where
 
 /// Expose the functionality to query events from the wasm VM's VP environment.
 pub fn vp_get_events<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     event_type_ptr: u64,
     event_type_len: u64,
 ) -> vp_host_fns::EnvResult<i64>
@@ -1955,7 +1959,7 @@ where
 /// performance
 #[allow(clippy::too_many_arguments)]
 pub fn vp_verify_tx_section_signature<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     hash_list_ptr: u64,
     hash_list_len: u64,
     public_keys_map_ptr: u64,
@@ -2037,7 +2041,7 @@ where
 /// printed at the [`tracing::Level::INFO`]. This function is for development
 /// only.
 pub fn tx_log_string<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     str_ptr: u64,
     str_len: u64,
 ) -> TxResult<()>
@@ -2057,7 +2061,7 @@ where
 
 /// Validate a VP WASM code hash in a tx environment.
 fn tx_validate_vp_code_hash<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     code_hash: &[u8],
     code_tag: &Option<String>,
 ) -> TxResult<()>
@@ -2069,7 +2073,11 @@ where
 {
     let code_hash = Hash::try_from(code_hash)
         .map_err(|e| TxRuntimeError::InvalidVpCodeHash(e.to_string()))?;
-    let state = env.state();
+
+    // NB: an env clone is required to avoid an immutable
+    // borrow while a mutable borrow is taking place
+    let env2 = env.clone();
+    let state = env2.state();
 
     // First check that code hash corresponds to the code tag if it is present
     if let Some(tag) = code_tag {
@@ -2115,8 +2123,9 @@ where
 }
 
 /// Set the sentinel for an invalid tx section commitment
-pub fn tx_set_commitment_sentinel<MEM, D, H, CA>(env: &TxVmEnv<MEM, D, H, CA>)
-where
+pub fn tx_set_commitment_sentinel<MEM, D, H, CA>(
+    env: &mut TxVmEnv<MEM, D, H, CA>,
+) where
     D: 'static + DB + for<'iter> DBIter<'iter>,
     H: 'static + StorageHasher,
     MEM: VmMemory,
@@ -2129,7 +2138,7 @@ where
 /// Verify a transaction signature
 #[allow(clippy::too_many_arguments)]
 pub fn tx_verify_tx_section_signature<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     hash_list_ptr: u64,
     hash_list_len: u64,
     public_keys_map_ptr: u64,
@@ -2201,7 +2210,7 @@ where
 
 /// Appends the new note commitments to the tree in storage
 pub fn tx_update_masp_note_commitment_tree<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     transaction_ptr: u64,
     transaction_len: u64,
 ) -> TxResult<i64>
@@ -2238,7 +2247,7 @@ where
 
 /// Yield a byte array value from the guest.
 pub fn tx_yield_value<MEM, D, H, CA>(
-    env: &TxVmEnv<MEM, D, H, CA>,
+    env: &mut TxVmEnv<MEM, D, H, CA>,
     buf_ptr: u64,
     buf_len: u64,
 ) -> TxResult<()>
@@ -2260,7 +2269,7 @@ where
 
 /// Evaluate a validity predicate with the given input data.
 pub fn vp_eval<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     vp_code_hash_ptr: u64,
     vp_code_hash_len: u64,
     input_data_ptr: u64,
@@ -2312,7 +2321,7 @@ where
 
 /// Get the native token's address
 pub fn vp_get_native_token<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     result_ptr: u64,
 ) -> vp_host_fns::EnvResult<()>
 where
@@ -2337,7 +2346,7 @@ where
 /// printed at the [`tracing::Level::INFO`]. This function is for development
 /// only.
 pub fn vp_log_string<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     str_ptr: u64,
     str_len: u64,
 ) -> vp_host_fns::EnvResult<()>
@@ -2358,7 +2367,7 @@ where
 
 /// Yield a byte array value from the guest.
 pub fn vp_yield_value<MEM, D, H, EVAL, CA>(
-    env: &VpVmEnv<MEM, D, H, EVAL, CA>,
+    env: &mut VpVmEnv<MEM, D, H, EVAL, CA>,
     buf_ptr: u64,
     buf_len: u64,
 ) -> vp_host_fns::EnvResult<()>
diff --git a/crates/namada/src/vm/wasm/host_env.rs b/crates/namada/src/vm/wasm/host_env.rs
index 20e968067a..fa5eea43b8 100644
--- a/crates/namada/src/vm/wasm/host_env.rs
+++ b/crates/namada/src/vm/wasm/host_env.rs
@@ -127,9 +127,9 @@ mod wrap_tx {
         D: DB + for<'iter> DBIter<'iter> + 'static,
         H: StorageHasher + 'static,
         CA: WasmCacheAccess + 'static,
-        F: Fn(&TxVmEnv<WasmMemory, D, H, CA>) -> RET,
+        F: Fn(&mut TxVmEnv<WasmMemory, D, H, CA>) -> RET,
     {
-        move |env| f(env.data())
+        move |mut env| f(env.data_mut())
     }
 
     pub(super) fn _1<F, ARG0, RET, D, H, CA>(
@@ -139,9 +139,9 @@ mod wrap_tx {
         D: DB + for<'iter> DBIter<'iter> + 'static,
         H: StorageHasher + 'static,
         CA: WasmCacheAccess + 'static,
-        F: Fn(&TxVmEnv<WasmMemory, D, H, CA>, ARG0) -> RET,
+        F: Fn(&mut TxVmEnv<WasmMemory, D, H, CA>, ARG0) -> RET,
     {
-        move |env, arg0| f(env.data(), arg0)
+        move |mut env, arg0| f(env.data_mut(), arg0)
     }
 
     pub(super) fn _2<F, ARG0, ARG1, RET, D, H, CA>(
@@ -151,9 +151,9 @@ mod wrap_tx {
         D: DB + for<'iter> DBIter<'iter> + 'static,
         H: StorageHasher + 'static,
         CA: WasmCacheAccess + 'static,
-        F: Fn(&TxVmEnv<WasmMemory, D, H, CA>, ARG0, ARG1) -> RET,
+        F: Fn(&mut TxVmEnv<WasmMemory, D, H, CA>, ARG0, ARG1) -> RET,
     {
-        move |env, arg0, arg1| f(env.data(), arg0, arg1)
+        move |mut env, arg0, arg1| f(env.data_mut(), arg0, arg1)
     }
 
     pub(super) fn _4<F, ARG0, ARG1, ARG2, ARG3, RET, D, H, CA>(
@@ -169,9 +169,17 @@ mod wrap_tx {
         D: DB + for<'iter> DBIter<'iter> + 'static,
         H: StorageHasher + 'static,
         CA: WasmCacheAccess + 'static,
-        F: Fn(&TxVmEnv<WasmMemory, D, H, CA>, ARG0, ARG1, ARG2, ARG3) -> RET,
+        F: Fn(
+            &mut TxVmEnv<WasmMemory, D, H, CA>,
+            ARG0,
+            ARG1,
+            ARG2,
+            ARG3,
+        ) -> RET,
     {
-        move |env, arg0, arg1, arg2, arg3| f(env.data(), arg0, arg1, arg2, arg3)
+        move |mut env, arg0, arg1, arg2, arg3| {
+            f(env.data_mut(), arg0, arg1, arg2, arg3)
+        }
     }
 
     pub(super) fn _6<F, ARG0, ARG1, ARG2, ARG3, ARG4, ARG5, RET, D, H, CA>(
@@ -190,7 +198,7 @@ mod wrap_tx {
         H: StorageHasher + 'static,
         CA: WasmCacheAccess + 'static,
         F: Fn(
-            &TxVmEnv<WasmMemory, D, H, CA>,
+            &mut TxVmEnv<WasmMemory, D, H, CA>,
             ARG0,
             ARG1,
             ARG2,
@@ -199,8 +207,8 @@ mod wrap_tx {
             ARG5,
         ) -> RET,
     {
-        move |env, arg0, arg1, arg2, arg3, arg4, arg5| {
-            f(env.data(), arg0, arg1, arg2, arg3, arg4, arg5)
+        move |mut env, arg0, arg1, arg2, arg3, arg4, arg5| {
+            f(env.data_mut(), arg0, arg1, arg2, arg3, arg4, arg5)
         }
     }
 
@@ -234,7 +242,7 @@ mod wrap_tx {
         H: StorageHasher + 'static,
         CA: WasmCacheAccess + 'static,
         F: Fn(
-            &TxVmEnv<WasmMemory, D, H, CA>,
+            &mut TxVmEnv<WasmMemory, D, H, CA>,
             ARG0,
             ARG1,
             ARG2,
@@ -244,8 +252,8 @@ mod wrap_tx {
             ARG6,
         ) -> RET,
     {
-        move |env, arg0, arg1, arg2, arg3, arg4, arg5, arg6| {
-            f(env.data(), arg0, arg1, arg2, arg3, arg4, arg5, arg6)
+        move |mut env, arg0, arg1, arg2, arg3, arg4, arg5, arg6| {
+            f(env.data_mut(), arg0, arg1, arg2, arg3, arg4, arg5, arg6)
         }
     }
 }
@@ -273,9 +281,9 @@ mod wrap_vp {
         H: StorageHasher + 'static,
         CA: WasmCacheAccess + 'static,
         EVAL: VpEvaluator<Db = D, H = H, Eval = EVAL, CA = CA> + 'static,
-        F: Fn(&VpVmEnv<WasmMemory, D, H, EVAL, CA>) -> RET,
+        F: Fn(&mut VpVmEnv<WasmMemory, D, H, EVAL, CA>) -> RET,
     {
-        move |env| f(env.data())
+        move |mut env| f(env.data_mut())
     }
 
     pub(super) fn _1<F, ARG0, RET, D, H, EVAL, CA>(
@@ -286,9 +294,9 @@ mod wrap_vp {
         H: StorageHasher + 'static,
         CA: WasmCacheAccess + 'static,
         EVAL: VpEvaluator<Db = D, H = H, Eval = EVAL, CA = CA> + 'static,
-        F: Fn(&VpVmEnv<WasmMemory, D, H, EVAL, CA>, ARG0) -> RET,
+        F: Fn(&mut VpVmEnv<WasmMemory, D, H, EVAL, CA>, ARG0) -> RET,
     {
-        move |env, arg0| f(env.data(), arg0)
+        move |mut env, arg0| f(env.data_mut(), arg0)
     }
 
     pub(super) fn _2<F, ARG0, ARG1, RET, D, H, EVAL, CA>(
@@ -303,9 +311,9 @@ mod wrap_vp {
         H: StorageHasher + 'static,
         CA: WasmCacheAccess + 'static,
         EVAL: VpEvaluator<Db = D, H = H, Eval = EVAL, CA = CA> + 'static,
-        F: Fn(&VpVmEnv<WasmMemory, D, H, EVAL, CA>, ARG0, ARG1) -> RET,
+        F: Fn(&mut VpVmEnv<WasmMemory, D, H, EVAL, CA>, ARG0, ARG1) -> RET,
     {
-        move |env, arg0, arg1| f(env.data(), arg0, arg1)
+        move |mut env, arg0, arg1| f(env.data_mut(), arg0, arg1)
     }
 
     pub(super) fn _4<F, ARG0, ARG1, ARG2, ARG3, RET, D, H, EVAL, CA>(
@@ -323,14 +331,16 @@ mod wrap_vp {
         CA: WasmCacheAccess + 'static,
         EVAL: VpEvaluator<Db = D, H = H, Eval = EVAL, CA = CA> + 'static,
         F: Fn(
-            &VpVmEnv<WasmMemory, D, H, EVAL, CA>,
+            &mut VpVmEnv<WasmMemory, D, H, EVAL, CA>,
             ARG0,
             ARG1,
             ARG2,
             ARG3,
         ) -> RET,
     {
-        move |env, arg0, arg1, arg2, arg3| f(env.data(), arg0, arg1, arg2, arg3)
+        move |mut env, arg0, arg1, arg2, arg3| {
+            f(env.data_mut(), arg0, arg1, arg2, arg3)
+        }
     }
 
     pub(super) fn _9<
@@ -369,7 +379,7 @@ mod wrap_vp {
         CA: WasmCacheAccess + 'static,
         EVAL: VpEvaluator<Db = D, H = H, Eval = EVAL, CA = CA> + 'static,
         F: Fn(
-            &VpVmEnv<WasmMemory, D, H, EVAL, CA>,
+            &mut VpVmEnv<WasmMemory, D, H, EVAL, CA>,
             ARG0,
             ARG1,
             ARG2,
@@ -381,9 +391,9 @@ mod wrap_vp {
             ARG8,
         ) -> RET,
     {
-        move |env, arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8| {
+        move |mut env, arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8| {
             f(
-                env.data(),
+                env.data_mut(),
                 arg0,
                 arg1,
                 arg2,
diff --git a/crates/tests/src/vm_host_env/tx.rs b/crates/tests/src/vm_host_env/tx.rs
index f6f0ff580f..7df8beb282 100644
--- a/crates/tests/src/vm_host_env/tx.rs
+++ b/crates/tests/src/vm_host_env/tx.rs
@@ -354,7 +354,7 @@ mod native_tx_host_env {
                                 batched_tx,
                             }: &mut TestTxEnv| {
 
-                            let tx_env = vm::host_env::testing::tx_env(
+                            let mut tx_env = vm::host_env::testing::tx_env(
                                 state,
                                 iterators,
                                 verifiers,
@@ -371,7 +371,7 @@ mod native_tx_host_env {
 
                             // Call the `host_env` function and unwrap any
                             // runtime errors
-                            $fn( &tx_env, $($arg),* ).unwrap()
+                            $fn( &mut tx_env, $($arg),* ).unwrap()
                         })
                     }
                 });
@@ -398,7 +398,7 @@ mod native_tx_host_env {
                                 batched_tx
                             }: &mut TestTxEnv| {
 
-                            let tx_env = vm::host_env::testing::tx_env(
+                            let mut tx_env = vm::host_env::testing::tx_env(
                                 state,
                                 iterators,
                                 verifiers,
@@ -415,7 +415,7 @@ mod native_tx_host_env {
 
                             // Call the `host_env` function and unwrap any
                             // runtime errors
-                            $fn( &tx_env, $($arg),* ).unwrap()
+                            $fn( &mut tx_env, $($arg),* ).unwrap()
                         })
                     }
                 });
@@ -442,7 +442,7 @@ mod native_tx_host_env {
                                 batched_tx,
                             }: &mut TestTxEnv| {
 
-                            let tx_env = vm::host_env::testing::tx_env(
+                            let mut tx_env = vm::host_env::testing::tx_env(
                                 state,
                                 iterators,
                                 verifiers,
@@ -458,7 +458,7 @@ mod native_tx_host_env {
                             );
 
                             // Call the `host_env` function
-                            $fn( &tx_env, $($arg),* )
+                            $fn( &mut tx_env, $($arg),* )
                         })
                     }
                 });
@@ -582,13 +582,18 @@ mod tests {
         // dbg!(&setup);
 
         let mut test_env = TestTxEnv::default();
-        let tx_env = setup_host_env(&setup, &mut test_env);
+        let mut tx_env = setup_host_env(&setup, &mut test_env);
 
         // Can fail, but must not panic
-        let _res =
-            host_env::tx_read(&tx_env, setup.key_memory_ptr, setup.key_len());
-        let _res =
-            host_env::tx_result_buffer(&tx_env, setup.read_buffer_memory_ptr);
+        let _res = host_env::tx_read(
+            &mut tx_env,
+            setup.key_memory_ptr,
+            setup.key_len(),
+        );
+        let _res = host_env::tx_result_buffer(
+            &mut tx_env,
+            setup.read_buffer_memory_ptr,
+        );
     }
 
     proptest! {
@@ -603,10 +608,10 @@ mod tests {
 
     fn test_tx_charge_gas_cannot_panic_aux(setup: TestSetup, gas: u64) {
         let mut test_env = TestTxEnv::default();
-        let tx_env = setup_host_env(&setup, &mut test_env);
+        let mut tx_env = setup_host_env(&setup, &mut test_env);
 
         // Can fail, but must not panic
-        let _res = host_env::tx_charge_gas(&tx_env, gas);
+        let _res = host_env::tx_charge_gas(&mut tx_env, gas);
     }
 
     proptest! {
@@ -622,11 +627,11 @@ mod tests {
         // dbg!(&setup);
 
         let mut test_env = TestTxEnv::default();
-        let tx_env = setup_host_env(&setup, &mut test_env);
+        let mut tx_env = setup_host_env(&setup, &mut test_env);
 
         // Can fail, but must not panic
         let _res = host_env::tx_has_key(
-            &tx_env,
+            &mut tx_env,
             setup.key_memory_ptr,
             setup.key_len(),
         );
@@ -645,18 +650,18 @@ mod tests {
         // dbg!(&setup);
 
         let mut test_env = TestTxEnv::default();
-        let tx_env = setup_host_env(&setup, &mut test_env);
+        let mut tx_env = setup_host_env(&setup, &mut test_env);
 
         // Can fail, but must not panic
         let _res = host_env::tx_write(
-            &tx_env,
+            &mut tx_env,
             setup.key_memory_ptr,
             setup.key_len(),
             setup.val_memory_ptr,
             setup.val_len(),
         );
         let _res = host_env::tx_write_temp(
-            &tx_env,
+            &mut tx_env,
             setup.key_memory_ptr,
             setup.key_len(),
             setup.val_memory_ptr,
@@ -677,11 +682,14 @@ mod tests {
         // dbg!(&setup);
 
         let mut test_env = TestTxEnv::default();
-        let tx_env = setup_host_env(&setup, &mut test_env);
+        let mut tx_env = setup_host_env(&setup, &mut test_env);
 
         // Can fail, but must not panic
-        let _res =
-            host_env::tx_delete(&tx_env, setup.key_memory_ptr, setup.key_len());
+        let _res = host_env::tx_delete(
+            &mut tx_env,
+            setup.key_memory_ptr,
+            setup.key_len(),
+        );
     }
 
     proptest! {
@@ -697,16 +705,16 @@ mod tests {
         // dbg!(&setup);
 
         let mut test_env = TestTxEnv::default();
-        let tx_env = setup_host_env(&setup, &mut test_env);
+        let mut tx_env = setup_host_env(&setup, &mut test_env);
 
         // Can fail, but must not panic
         let _res = host_env::tx_iter_prefix(
-            &tx_env,
+            &mut tx_env,
             setup.key_memory_ptr,
             setup.key_len(),
         );
         let _res = host_env::tx_iter_next(
-            &tx_env,
+            &mut tx_env,
             // This field is not used for anything else in this test
             setup.val_memory_ptr,
         );
diff --git a/crates/tests/src/vm_host_env/vp.rs b/crates/tests/src/vm_host_env/vp.rs
index 4444310a39..10a7bda496 100644
--- a/crates/tests/src/vm_host_env/vp.rs
+++ b/crates/tests/src/vm_host_env/vp.rs
@@ -267,7 +267,7 @@ mod native_vp_host_env {
                                 vp_cache_dir: _,
                             }: &mut TestVpEnv| {
 
-                            let env = vm::host_env::testing::vp_env(
+                            let mut env = vm::host_env::testing::vp_env(
                                 addr,
                                 state,
                                 iterators,
@@ -285,7 +285,7 @@ mod native_vp_host_env {
 
                             // Call the `host_env` function and unwrap any
                             // runtime errors
-                            $fn( &env, $($arg),* ).unwrap()
+                            $fn( &mut env, $($arg),* ).unwrap()
                         })
                     }
                 });
@@ -312,7 +312,7 @@ mod native_vp_host_env {
                                 vp_cache_dir: _,
                             }: &mut TestVpEnv| {
 
-                            let env = vm::host_env::testing::vp_env(
+                            let mut env = vm::host_env::testing::vp_env(
                                 addr,
                                 state,
                                 iterators,
@@ -330,7 +330,7 @@ mod native_vp_host_env {
 
                             // Call the `host_env` function and unwrap any
                             // runtime errors
-                            $fn( &env, $($arg),* ).unwrap()
+                            $fn( &mut env, $($arg),* ).unwrap()
                         })
                     }
                 });