Merge pull request #10 from ChainSafe/willem/store-shard-heights

Add a mapping between shard index and the max block height that is updated on `put_{pool}_subtree_roots`
ChainSafe · Aug 29, 2024 · 3d9cc60 · 3d9cc60
2 parents 9d2b9b1 + b5266b2
commit 3d9cc60
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/zcash_client_memory/src/lib.rs b/zcash_client_memory/src/lib.rs
@@ -1,8 +1,9 @@
+use incrementalmerkletree::Address;
 use scanning::ScanQueue;
 
 use shardtree::{store::memory::MemoryShardStore, ShardTree};
 use std::{
-    collections::{hash_map::Entry, BTreeMap},
+    collections::{hash_map::Entry, BTreeMap, HashMap},
     hash::Hash,
     ops::Deref,
 };
@@ -58,7 +59,6 @@ pub struct MemoryWalletDb {
     network: Network,
     accounts: Vec<Account>,
     blocks: BTreeMap<BlockHeight, MemoryWalletBlock>,
-
     tx_table: TransactionTable,
 
     received_notes: ReceivedNoteTable,
@@ -74,12 +74,18 @@ pub struct MemoryWalletDb {
         { SAPLING_SHARD_HEIGHT * 2 },
         SAPLING_SHARD_HEIGHT,
     >,
+    /// Stores the block height corresponding to the last note commitment in a shard
+    sapling_tree_shard_end_heights: BTreeMap<Address, BlockHeight>,
+
     #[cfg(feature = "orchard")]
     orchard_tree: ShardTree<
         MemoryShardStore<orchard::tree::MerkleHashOrchard, BlockHeight>,
         { ORCHARD_SHARD_HEIGHT * 2 },
         ORCHARD_SHARD_HEIGHT,
     >,
+    #[cfg(feature = "orchard")]
+    /// Stores the block height corresponding to the last note commitment in a shard
+    orchard_tree_shard_end_heights: BTreeMap<Address, BlockHeight>,
 }
 
 impl MemoryWalletDb {
@@ -89,8 +95,11 @@ impl MemoryWalletDb {
             accounts: Vec::new(),
             blocks: BTreeMap::new(),
             sapling_tree: ShardTree::new(MemoryShardStore::empty(), max_checkpoints),
+            sapling_tree_shard_end_heights: BTreeMap::new(),
             #[cfg(feature = "orchard")]
             orchard_tree: ShardTree::new(MemoryShardStore::empty(), max_checkpoints),
+            #[cfg(feature = "orchard")]
+            orchard_tree_shard_end_heights: BTreeMap::new(),
             tx_table: TransactionTable::new(),
             received_notes: ReceivedNoteTable::new(),
             nullifiers: NullifierMap::new(),

diff --git a/zcash_client_memory/src/wallet_commitment_trees.rs b/zcash_client_memory/src/wallet_commitment_trees.rs
@@ -45,6 +45,13 @@ impl WalletCommitmentTrees for MemoryWalletDb {
             Ok::<_, ShardTreeError<Self::Error>>(())
         })?;
 
+        // store the end block heights for each shard as well
+        for (root, i) in roots.iter().zip(0u64..) {
+            let root_addr = Address::from_parts(SAPLING_SHARD_HEIGHT.into(), start_index + i);
+            self.sapling_tree_shard_end_heights
+                .insert(root_addr, root.subtree_end_height());
+        }
+
         Ok(())
     }
 
@@ -81,6 +88,13 @@ impl WalletCommitmentTrees for MemoryWalletDb {
             Ok::<_, ShardTreeError<Self::Error>>(())
         })?;
 
+        // store the end block heights for each shard as well
+        for (root, i) in roots.iter().zip(0u64..) {
+            let root_addr = Address::from_parts(SAPLING_SHARD_HEIGHT.into(), start_index + i);
+            self.orchard_tree_shard_end_heights
+                .insert(root_addr, root.subtree_end_height());
+        }
+
         Ok(())
     }
 }