New storage write API

.changelog/unreleased/improvements/2355-new-storage-write.md

@@ -0,0 +1,2 @@
+- Skip writing some MASP and IBC storage keys to merkle tree and DB diffs.
+  ([\#2355](https://github.com/anoma/namada/pull/2355))

Cargo.toml

@@ -67,6 +67,7 @@ base58 = "0.2.0"
 base64 = "0.13.0"
 bech32 = "0.8.0"
 bimap = {version = "0.6.2", features = ["serde"]}
+bitflags = "2.4.2"
 bit-set = "0.5.2"
 blake2b-rs = "0.2.0"
 byte-unit = "4.0.13"

crates/apps/src/lib/node/ledger/shell/finalize_block.rs

@@ -598,15 +598,16 @@ where
 
         // Update the MASP commitment tree anchor if the tree was updated
         let tree_key = token::storage_key::masp_commitment_tree_key();
-        if let Some(StorageModification::Write { value }) =
+        if let Some(StorageModification::Write { value, action: _ }) =
             self.wl_storage.write_log.read(&tree_key).0
         {
             let updated_tree = CommitmentTree::<Node>::try_from_slice(value)
                 .into_storage_result()?;
             let anchor_key = token::storage_key::masp_commitment_anchor_key(
                 updated_tree.root(),
             );
-            self.wl_storage.write(&anchor_key, ())?;
+            self.wl_storage
+                .write_without_merkle_diffs(&anchor_key, ())?;
         }
 
         if update_for_tendermint {

crates/apps/src/lib/node/ledger/shell/init_chain.rs

@@ -141,17 +141,20 @@ where
         let note_commitment_tree_key =
             token::storage_key::masp_commitment_tree_key();
         self.wl_storage
-            .write(&note_commitment_tree_key, empty_commitment_tree)
+            .write_without_merkle_diffs(
+                &note_commitment_tree_key,
+                empty_commitment_tree,
+            )
             .unwrap();
         let commitment_tree_anchor_key =
             token::storage_key::masp_commitment_anchor_key(anchor);
         self.wl_storage
-            .write(&commitment_tree_anchor_key, ())
+            .write_without_merkle_diffs(&commitment_tree_anchor_key, ())
             .unwrap();
 
         // Init masp convert anchor
         let convert_anchor_key = token::storage_key::masp_convert_anchor_key();
-        self.wl_storage.write(
+        self.wl_storage.write_without_merkle_diffs(
             &convert_anchor_key,
             namada::types::hash::Hash(
                 bls12_381::Scalar::from(
@@ -215,9 +218,9 @@ where
             self.update_eth_oracle(&Default::default());
         } else {
             self.wl_storage
-                .write_bytes(
+                .write(
                     &namada::eth_bridge::storage::active_key(),
-                    EthBridgeStatus::Disabled.serialize_to_vec(),
+                    EthBridgeStatus::Disabled,
                 )
                 .unwrap();
         }

crates/apps/src/lib/node/ledger/shell/mod.rs

@@ -2041,10 +2041,7 @@ mod test_utils {
         use namada::eth_bridge::storage::eth_bridge_queries::EthBridgeStatus;
         shell
             .wl_storage
-            .write_bytes(
-                &active_key(),
-                EthBridgeStatus::Disabled.serialize_to_vec(),
-            )
+            .write(&active_key(), EthBridgeStatus::Disabled)
             .expect("Test failed");
     }
 

crates/apps/src/lib/node/ledger/storage/mod.rs

@@ -52,14 +52,17 @@ fn new_blake2b() -> Blake2b {
 mod tests {
     use std::collections::HashMap;
 
+    use borsh::BorshDeserialize;
     use itertools::Itertools;
     use namada::eth_bridge::storage::proof::BridgePoolRootProof;
     use namada::ledger::eth_bridge::storage::bridge_pool;
     use namada::ledger::gas::STORAGE_ACCESS_GAS_PER_BYTE;
     use namada::ledger::ibc::storage::ibc_key;
     use namada::ledger::parameters::{EpochDuration, Parameters};
     use namada::state::write_log::WriteLog;
-    use namada::state::{self, StorageWrite, StoreType, WlStorage};
+    use namada::state::{
+        self, StorageRead, StorageWrite, StoreType, WlStorage, WriteOpts, DB,
+    };
     use namada::token::conversion::update_allowed_conversions;
     use namada::types::chain::ChainId;
     use namada::types::ethereum_events::Uint;
@@ -530,14 +533,19 @@ mod tests {
                     storage.write(&key, value_bytes)?;
                 }
                 3 => {
-                    storage.batch_delete_subspace_val(&mut batch, &key)?;
+                    storage.batch_delete_subspace_val(
+                        &mut batch,
+                        &key,
+                        WriteOpts::ALL,
+                    )?;
                 }
                 _ => {
                     let value_bytes = types::encode(&storage.block.height);
                     storage.batch_write_subspace_val(
                         &mut batch,
                         &key,
                         value_bytes,
+                        WriteOpts::ALL,
                     )?;
                 }
             }
@@ -783,4 +791,191 @@ mod tests {
             .map(Result::unwrap);
         itertools::assert_equal(iter, expected);
     }
+
+    #[test]
+    fn test_persistent_storage_writing_without_merklizing_or_diffs() {
+        let db_path =
+            TempDir::new().expect("Unable to create a temporary DB directory");
+        let storage = PersistentStorage::open(
+            db_path.path(),
+            ChainId::default(),
+            address::nam(),
+            None,
+            None,
+        );
+        let mut wls = WlStorage {
+            storage,
+            write_log: Default::default(),
+        };
+        // Start the first block
+        let first_height = BlockHeight::first();
+        wls.storage.block.height = first_height;
+
+        let key1 = Key::parse("testing1").unwrap();
+        let val1 = 1u64;
+        let key2 = Key::parse("testing2").unwrap();
+        let val2 = 2u64;
+
+        // Standard write of key-val-1
+        wls.write(&key1, val1).unwrap();
+
+        // Read from WlStorage should return val1
+        let res = wls.read::<u64>(&key1).unwrap().unwrap();
+        assert_eq!(res, val1);
+
+        // Read from Storage shouldn't return val1 because the block hasn't been
+        // committed
+        let (res, _) = wls.storage.read(&key1).unwrap();
+        assert!(res.is_none());
+
+        // Write key-val-2 without merklizing or diffs
+        wls.write_without_merkle_diffs(&key2, val2).unwrap();
+
+        // Read from WlStorage should return val2
+        let res = wls.read::<u64>(&key2).unwrap().unwrap();
+        assert_eq!(res, val2);
+
+        // Commit block and storage changes
+        wls.commit_block().unwrap();
+        wls.storage.block.height = wls.storage.block.height.next_height();
+        let second_height = wls.storage.block.height;
+
+        // Read key1 from Storage should return val1
+        let (res1, _) = wls.storage.read(&key1).unwrap();
+        let res1 = u64::try_from_slice(&res1.unwrap()).unwrap();
+        assert_eq!(res1, val1);
+
+        // Check merkle tree inclusion of key-val-1 explicitly
+        let is_merklized1 = wls.storage.block.tree.has_key(&key1).unwrap();
+        assert!(is_merklized1);
+
+        // Key2 should be in storage. Confirm by reading from
+        // WlStorage and also by reading Storage subspace directly
+        let res2 = wls.read::<u64>(&key2).unwrap().unwrap();
+        assert_eq!(res2, val2);
+        let res2 = wls.storage.db.read_subspace_val(&key2).unwrap().unwrap();
+        let res2 = u64::try_from_slice(&res2).unwrap();
+        assert_eq!(res2, val2);
+
+        // Check explicitly that key-val-2 is not in merkle tree
+        let is_merklized2 = wls.storage.block.tree.has_key(&key2).unwrap();
+        assert!(!is_merklized2);
+
+        // Check that the proper diffs exist for key-val-1
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, first_height, true)
+            .unwrap();
+        assert!(res1.is_none());
+
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, first_height, false)
+            .unwrap()
+            .unwrap();
+        let res1 = u64::try_from_slice(&res1).unwrap();
+        assert_eq!(res1, val1);
+
+        // Check that there are diffs for key-val-2 in block 0, since all keys
+        // need to have diffs for at least 1 block for rollback purposes
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, first_height, true)
+            .unwrap();
+        assert!(res2.is_none());
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, first_height, false)
+            .unwrap()
+            .unwrap();
+        let res2 = u64::try_from_slice(&res2).unwrap();
+        assert_eq!(res2, val2);
+
+        // Delete the data then commit the block
+        wls.delete(&key1).unwrap();
+        wls.delete_without_diffs(&key2).unwrap();
+        wls.commit_block().unwrap();
+        wls.storage.block.height = wls.storage.block.height.next_height();
+
+        // Check the key-vals are removed from the storage subspace
+        let res1 = wls.read::<u64>(&key1).unwrap();
+        let res2 = wls.read::<u64>(&key2).unwrap();
+        assert!(res1.is_none() && res2.is_none());
+        let res1 = wls.storage.db.read_subspace_val(&key1).unwrap();
+        let res2 = wls.storage.db.read_subspace_val(&key2).unwrap();
+        assert!(res1.is_none() && res2.is_none());
+
+        // Check that the key-vals don't exist in the merkle tree anymore
+        let is_merklized1 = wls.storage.block.tree.has_key(&key1).unwrap();
+        let is_merklized2 = wls.storage.block.tree.has_key(&key2).unwrap();
+        assert!(!is_merklized1 && !is_merklized2);
+
+        // Check that key-val-1 diffs are properly updated for blocks 0 and 1
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, first_height, true)
+            .unwrap();
+        assert!(res1.is_none());
+
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, first_height, false)
+            .unwrap()
+            .unwrap();
+        let res1 = u64::try_from_slice(&res1).unwrap();
+        assert_eq!(res1, val1);
+
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, second_height, true)
+            .unwrap()
+            .unwrap();
+        let res1 = u64::try_from_slice(&res1).unwrap();
+        assert_eq!(res1, val1);
+
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, second_height, false)
+            .unwrap();
+        assert!(res1.is_none());
+
+        // Check that key-val-2 diffs don't exist for block 0 anymore
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, first_height, true)
+            .unwrap();
+        assert!(res2.is_none());
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, first_height, false)
+            .unwrap();
+        assert!(res2.is_none());
+
+        // Check that the block 1 diffs for key-val-2 include an "old" value of
+        // val2 and no "new" value
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, second_height, true)
+            .unwrap()
+            .unwrap();
+        let res2 = u64::try_from_slice(&res2).unwrap();
+        assert_eq!(res2, val2);
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, second_height, false)
+            .unwrap();
+        assert!(res2.is_none());
+    }
 }