zcash_client_sqlite: Track Orchard commitment tree sizes in TestState

zcash_client_sqlite/src/chain.rs

@@ -414,6 +414,7 @@ mod tests {
             AddressType::DefaultExternal,
             NonNegativeAmount::const_from_u64(8),
             2,
+            2,
         );
         st.generate_next_block(
             &dfvk,

zcash_client_sqlite/src/testing.rs

@@ -166,10 +166,54 @@
     }
 }
 
+pub(crate) struct CachedBlock {
+    height: BlockHeight,
+    hash: BlockHash,
+    sapling_end_size: u32,
+    orchard_end_size: u32,
+}
+
+impl CachedBlock {
+    fn none(sapling_activation_height: BlockHeight) -> Self {
+        Self {
+            height: sapling_activation_height,
+            hash: BlockHash([0; 32]),
+            sapling_end_size: 0,
+            orchard_end_size: 0,
+        }
+    }
+
+    fn at(
+        height: BlockHeight,
+        hash: BlockHash,
+        sapling_tree_size: u32,
+        orchard_tree_size: u32,
+    ) -> Self {
+        Self {
+            height,
+            hash,
+            sapling_end_size: sapling_tree_size,
+            orchard_end_size: orchard_tree_size,
+        }
+    }
+
+    fn roll_forward(self, cb: &CompactBlock) -> Self {
+        assert_eq!(self.height + 1, cb.height());
+        Self {
+            height: cb.height(),
+            hash: cb.hash(),
+            sapling_end_size: self.sapling_end_size
+                + cb.vtx.iter().map(|tx| tx.outputs.len() as u32).sum::<u32>(),
+            orchard_end_size: self.orchard_end_size
+                + cb.vtx.iter().map(|tx| tx.actions.len() as u32).sum::<u32>(),
+        }
+    }
+}
+
 /// The state for a `zcash_client_sqlite` test.
 pub(crate) struct TestState<Cache> {
     cache: Cache,
-    latest_cached_block: Option<(BlockHeight, BlockHash, u32)>,
+    latest_cached_block: Option<CachedBlock>,
     _data_file: NamedTempFile,
     db_data: WalletDb<Connection, Network>,
     test_account: Option<(
@@ -190,8 +234,8 @@
         self.cache.block_source()
     }
 
-    pub(crate) fn latest_cached_block(&self) -> &Option<(BlockHeight, BlockHash, u32)> {
-        &self.latest_cached_block
+    pub(crate) fn latest_cached_block(&self) -> Option<&CachedBlock> {
+        self.latest_cached_block.as_ref()
     }
 
     /// Creates a fake block at the expected next height containing a single output of the
@@ -202,19 +246,22 @@
         req: AddressType,
         value: NonNegativeAmount,
     ) -> (BlockHeight, Cache::InsertResult, Fvk::Nullifier) {
-        let (height, prev_hash, initial_sapling_tree_size) = self
+        let cached_block = self
             .latest_cached_block
-            .map(|(prev_height, prev_hash, end_size)| (prev_height + 1, prev_hash, end_size))
-            .unwrap_or_else(|| (self.sapling_activation_height(), BlockHash([0; 32]), 0));
+            .take()
+            .unwrap_or_else(|| CachedBlock::none(self.sapling_activation_height() - 1));
+        let height = cached_block.height + 1;
 
         let (res, nf) = self.generate_block_at(
             height,
-            prev_hash,
+            cached_block.hash,
             fvk,
             req,
             value,
-            initial_sapling_tree_size,
+            cached_block.sapling_end_size,
+            cached_block.orchard_end_size,
         );
+        assert!(self.latest_cached_block.is_some());
 
         (height, res, nf)
     }
@@ -224,6 +271,7 @@
     ///
     /// This generated block will be treated as the latest block, and subsequent calls to
     /// [`Self::generate_next_block`] will build on it.
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn generate_block_at<Fvk: TestFvk>(
         &mut self,
         height: BlockHeight,
@@ -232,6 +280,7 @@
         req: AddressType,
         value: NonNegativeAmount,
         initial_sapling_tree_size: u32,
+        initial_orchard_tree_size: u32,
     ) -> (Cache::InsertResult, Fvk::Nullifier) {
         let (cb, nf) = fake_compact_block(
             &self.network(),
@@ -241,15 +290,19 @@
             req,
             value,
             initial_sapling_tree_size,
+            initial_orchard_tree_size,
         );
         let res = self.cache.insert(&cb);
 
-        self.latest_cached_block = Some((
-            height,
-            cb.hash(),
-            initial_sapling_tree_size
-                + cb.vtx.iter().map(|tx| tx.outputs.len() as u32).sum::<u32>(),
-        ));
+        self.latest_cached_block = Some(
+            CachedBlock::at(
+                height - 1,
+                cb.hash(),
+                initial_sapling_tree_size,
+                initial_orchard_tree_size,
+            )
+            .roll_forward(&cb),
+        );
 
         (res, nf)
     }
@@ -263,29 +316,26 @@
         to: impl Into<Address>,
         value: NonNegativeAmount,
     ) -> (BlockHeight, Cache::InsertResult) {
-        let (height, prev_hash, initial_sapling_tree_size) = self
+        let cached_block = self
             .latest_cached_block
-            .map(|(prev_height, prev_hash, end_size)| (prev_height + 1, prev_hash, end_size))
-            .unwrap_or_else(|| (self.sapling_activation_height(), BlockHash([0; 32]), 0));
+            .take()
+            .unwrap_or_else(|| CachedBlock::none(self.sapling_activation_height() - 1));
+        let height = cached_block.height + 1;
 
         let cb = fake_compact_block_spending(
             &self.network(),
             height,
-            prev_hash,
+            cached_block.hash,
             note,
             fvk,
             to.into(),
             value,
-            initial_sapling_tree_size,
+            cached_block.sapling_end_size,
+            cached_block.orchard_end_size,
         );
         let res = self.cache.insert(&cb);
 
-        self.latest_cached_block = Some((
-            height,
-            cb.hash(),
-            initial_sapling_tree_size
-                + cb.vtx.iter().map(|tx| tx.outputs.len() as u32).sum::<u32>(),
-        ));
+        self.latest_cached_block = Some(cached_block.roll_forward(&cb));
 
         (height, res)
     }
@@ -320,27 +370,23 @@
         tx_index: usize,
         tx: &Transaction,
     ) -> (BlockHeight, Cache::InsertResult) {
-        let (height, prev_hash, initial_sapling_tree_size) = self
+        let cached_block = self
             .latest_cached_block
-            .map(|(prev_height, prev_hash, end_size)| (prev_height + 1, prev_hash, end_size))
-            .unwrap_or_else(|| (self.sapling_activation_height(), BlockHash([0; 32]), 0));
+            .take()
+            .unwrap_or_else(|| CachedBlock::none(self.sapling_activation_height() - 1));
+        let height = cached_block.height + 1;
 
         let cb = fake_compact_block_from_tx(
             height,
-            prev_hash,
+            cached_block.hash,
             tx_index,
             tx,
-            initial_sapling_tree_size,
-            0,
+            cached_block.sapling_end_size,
+            cached_block.orchard_end_size,
         );
         let res = self.cache.insert(&cb);
 
-        self.latest_cached_block = Some((
-            height,
-            cb.hash(),
-            initial_sapling_tree_size
-                + cb.vtx.iter().map(|tx| tx.outputs.len() as u32).sum::<u32>(),
-        ));
+        self.latest_cached_block = Some(cached_block.roll_forward(&cb));
 
         (height, res)
     }
@@ -399,10 +445,12 @@
         self.cache
             .block_source()
             .with_blocks::<_, Infallible>(None, None, |block: CompactBlock| {
-                self.latest_cached_block = Some((
+                let chain_metadata = block.chain_metadata.unwrap();
+                self.latest_cached_block = Some(CachedBlock::at(
                     BlockHeight::from_u32(block.height.try_into().unwrap()),
                     BlockHash::from_slice(block.hash.as_slice()),
-                    block.chain_metadata.unwrap().sapling_commitment_tree_size,
+                    chain_metadata.sapling_commitment_tree_size,
+                    chain_metadata.orchard_commitment_tree_size,
                 ));
                 Ok(())
             })
@@ -1070,6 +1118,7 @@
 
 /// Create a fake CompactBlock at the given height, containing a single output paying
 /// an address. Returns the CompactBlock and the nullifier for the new note.
+#[allow(clippy::too_many_arguments)]
 fn fake_compact_block<P: consensus::Parameters, Fvk: TestFvk>(
     params: &P,
     height: BlockHeight,
@@ -1078,6 +1127,7 @@
     req: AddressType,
     value: NonNegativeAmount,
     initial_sapling_tree_size: u32,
+    initial_orchard_tree_size: u32,
 ) -> (CompactBlock, Fvk::Nullifier) {
     // Create a fake Note for the account
     let mut rng = OsRng;
@@ -1094,8 +1144,13 @@
         &mut rng,
     );
 
-    let cb =
-        fake_compact_block_from_compact_tx(ctx, height, prev_hash, initial_sapling_tree_size, 0);
+    let cb = fake_compact_block_from_compact_tx(
+        ctx,
+        height,
+        prev_hash,
+        initial_sapling_tree_size,
+        initial_orchard_tree_size,
+    );
     (cb, nf)
 }
 
@@ -1152,6 +1207,7 @@
     to: Address,
     value: NonNegativeAmount,
     initial_sapling_tree_size: u32,
+    initial_orchard_tree_size: u32,
 ) -> CompactBlock {
     let mut rng = OsRng;
     let mut ctx = fake_compact_tx(&mut rng);
@@ -1229,7 +1285,13 @@
         }
     }
 
-    fake_compact_block_from_compact_tx(ctx, height, prev_hash, initial_sapling_tree_size, 0)
+    fake_compact_block_from_compact_tx(
+        ctx,
+        height,
+        prev_hash,
+        initial_sapling_tree_size,
+        initial_orchard_tree_size,
+    )
 }
 
 fn fake_compact_block_from_compact_tx(

zcash_client_sqlite/src/testing/pool.rs

@@ -1248,10 +1248,23 @@ pub(crate) fn birthday_in_anchor_shard<T: ShieldedPoolTester>() {
 
     let received_tx_height = birthday.height() + 10;
 
-    let initial_sapling_tree_size =
-        u64::from(birthday.sapling_frontier().value().unwrap().position() + 1)
-            .try_into()
-            .unwrap();
+    let initial_sapling_tree_size = birthday
+        .sapling_frontier()
+        .value()
+        .map(|f| u64::from(f.position() + 1))
+        .unwrap_or(0)
+        .try_into()
+        .unwrap();
+    #[cfg(feature = "orchard")]
+    let initial_orchard_tree_size = birthday
+        .orchard_frontier()
+        .value()
+        .map(|f| u64::from(f.position() + 1))
+        .unwrap_or(0)
+        .try_into()
+        .unwrap();
+    #[cfg(not(feature = "orchard"))]
+    let initial_orchard_tree_size = 0;
 
     // Generate 9 blocks that have no value for us, starting at the birthday height.
     let not_our_key = T::sk_to_fvk(&T::sk(&[]));
@@ -1263,6 +1276,7 @@ pub(crate) fn birthday_in_anchor_shard<T: ShieldedPoolTester>() {
         AddressType::DefaultExternal,
         not_our_value,
         initial_sapling_tree_size,
+        initial_orchard_tree_size,
     );
     for _ in 1..9 {
         st.generate_next_block(&not_our_key, AddressType::DefaultExternal, not_our_value);
@@ -1340,7 +1354,8 @@ pub(crate) fn checkpoint_gaps<T: ShieldedPoolTester>() {
         &not_our_key,
         AddressType::DefaultExternal,
         not_our_value,
-        st.latest_cached_block().unwrap().2,
+        st.latest_cached_block().unwrap().sapling_end_size,
+        st.latest_cached_block().unwrap().orchard_end_size,
     );
 
     // Scan the block

zcash_client_sqlite/src/wallet.rs

@@ -2645,6 +2645,7 @@ mod tests {
             AddressType::DefaultExternal,
             not_our_value,
             17,
+            17,
         );
         st.scan_cached_blocks(end_height, 1);
 
@@ -2661,6 +2662,7 @@ mod tests {
             AddressType::DefaultExternal,
             not_our_value,
             0,
+            0,
         );
         st.scan_cached_blocks(start_height, 1);