zcash_client_sqlite: Add Orchard note commitments in put_blocks and…

… `put_block`

zcash_client_sqlite/src/lib.rs

@@ -540,6 +540,8 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
                 )
             });
             let mut sapling_commitments = vec![];
+            #[cfg(feature = "orchard")]
+            let mut orchard_commitments = vec![];
             let mut last_scanned_height = None;
             let mut note_positions = vec![];
             for block in blocks.into_iter() {
@@ -558,6 +560,10 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
                     block.block_time(),
                     block.sapling().final_tree_size(),
                     block.sapling().commitments().len().try_into().unwrap(),
+                    #[cfg(feature = "orchard")]
+                    block.orchard().final_tree_size(),
+                    #[cfg(feature = "orchard")]
+                    block.orchard().commitments().len().try_into().unwrap(),
                 )?;
 
                 for tx in block.transactions() {
@@ -652,6 +658,8 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
                 last_scanned_height = Some(block.height());
                 let block_commitments = block.into_commitments();
                 sapling_commitments.extend(block_commitments.sapling.into_iter().map(Some));
+                #[cfg(feature = "orchard")]
+                orchard_commitments.extend(block_commitments.orchard.into_iter().map(Some));
             }
 
             // Prune the nullifier map of entries we no longer need.
@@ -669,31 +677,63 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
             {
                 // Create subtrees from the note commitments in parallel.
                 const CHUNK_SIZE: usize = 1024;
-                let subtrees = sapling_commitments
-                    .par_chunks_mut(CHUNK_SIZE)
-                    .enumerate()
-                    .filter_map(|(i, chunk)| {
-                        let start = start_position + (i * CHUNK_SIZE) as u64;
-                        let end = start + chunk.len() as u64;
-
-                        shardtree::LocatedTree::from_iter(
-                            start..end,
-                            SAPLING_SHARD_HEIGHT.into(),
-                            chunk.iter_mut().map(|n| n.take().expect("always Some")),
-                        )
-                    })
-                    .map(|res| (res.subtree, res.checkpoints))
-                    .collect::<Vec<_>>();
-
-                // Update the Sapling note commitment tree with all newly read note commitments
-                let mut subtrees = subtrees.into_iter();
-                wdb.with_sapling_tree_mut::<_, _, Self::Error>(move |sapling_tree| {
-                    for (tree, checkpoints) in &mut subtrees {
-                        sapling_tree.insert_tree(tree, checkpoints)?;
-                    }
+                {
+                    let sapling_subtrees = sapling_commitments
+                        .par_chunks_mut(CHUNK_SIZE)
+                        .enumerate()
+                        .filter_map(|(i, chunk)| {
+                            let start = start_position + (i * CHUNK_SIZE) as u64;
+                            let end = start + chunk.len() as u64;
+
+                            shardtree::LocatedTree::from_iter(
+                                start..end,
+                                SAPLING_SHARD_HEIGHT.into(),
+                                chunk.iter_mut().map(|n| n.take().expect("always Some")),
+                            )
+                        })
+                        .map(|res| (res.subtree, res.checkpoints))
+                        .collect::<Vec<_>>();
+
+                    // Update the Sapling note commitment tree with all newly read note commitments
+                    let mut sapling_subtrees = sapling_subtrees.into_iter();
+                    wdb.with_sapling_tree_mut::<_, _, Self::Error>(move |sapling_tree| {
+                        for (tree, checkpoints) in &mut sapling_subtrees {
+                            sapling_tree.insert_tree(tree, checkpoints)?;
+                        }
 
-                    Ok(())
-                })?;
+                        Ok(())
+                    })?;
+                }
+
+                // Create subtrees from the note commitments in parallel.
+                #[cfg(feature = "orchard")]
+                {
+                    let orchard_subtrees = orchard_commitments
+                        .par_chunks_mut(CHUNK_SIZE)
+                        .enumerate()
+                        .filter_map(|(i, chunk)| {
+                            let start = start_position + (i * CHUNK_SIZE) as u64;
+                            let end = start + chunk.len() as u64;
+
+                            shardtree::LocatedTree::from_iter(
+                                start..end,
+                                ORCHARD_SHARD_HEIGHT.into(),
+                                chunk.iter_mut().map(|n| n.take().expect("always Some")),
+                            )
+                        })
+                        .map(|res| (res.subtree, res.checkpoints))
+                        .collect::<Vec<_>>();
+
+                    // Update the Sapling note commitment tree with all newly read note commitments
+                    let mut orchard_subtrees = orchard_subtrees.into_iter();
+                    wdb.with_orchard_tree_mut::<_, _, Self::Error>(move |orchard_tree| {
+                        for (tree, checkpoints) in &mut orchard_subtrees {
+                            orchard_tree.insert_tree(tree, checkpoints)?;
+                        }
+
+                        Ok(())
+                    })?;
+                }
 
                 // Update now-expired transactions that didn't get mined.
                 wallet::update_expired_notes(wdb.conn.0, last_scanned_height)?;

zcash_client_sqlite/src/wallet.rs

@@ -1665,6 +1665,8 @@ pub(crate) fn put_block(
     block_time: u32,
     sapling_commitment_tree_size: u32,
     sapling_output_count: u32,
+    #[cfg(feature = "orchard")] orchard_commitment_tree_size: u32,
+    #[cfg(feature = "orchard")] orchard_action_count: u32,
 ) -> Result<(), SqliteClientError> {
     let block_hash_data = conn
         .query_row(
@@ -1695,29 +1697,42 @@ pub(crate) fn put_block(
             time,
             sapling_commitment_tree_size,
             sapling_output_count,
-            sapling_tree
+            sapling_tree,
+            orchard_commitment_tree_size,
+            orchard_action_count
         )
         VALUES (
             :height,
             :hash,
             :block_time,
             :sapling_commitment_tree_size,
             :sapling_output_count,
-            x'00'
+            x'00',
+            :orchard_commitment_tree_size,
+            :orchard_action_count
         )
         ON CONFLICT (height) DO UPDATE
         SET hash = :hash,
             time = :block_time,
             sapling_commitment_tree_size = :sapling_commitment_tree_size,
-            sapling_output_count = :sapling_output_count",
+            sapling_output_count = :sapling_output_count,
+            orchard_commitment_tree_size = :orchard_commitment_tree_size,
+            orchard_action_count = :orchard_action_count",
     )?;
 
+    #[cfg(not(feature = "orchard"))]
+    let orchard_commitment_tree_size: Option<u32> = None;
+    #[cfg(not(feature = "orchard"))]
+    let orchard_action_count: Option<u32> = None;
+
     stmt_upsert_block.execute(named_params![
         ":height": u32::from(block_height),
         ":hash": &block_hash.0[..],
         ":block_time": block_time,
         ":sapling_commitment_tree_size": sapling_commitment_tree_size,
         ":sapling_output_count": sapling_output_count,
+        ":orchard_commitment_tree_size": orchard_commitment_tree_size,
+        ":orchard_action_count": orchard_action_count,
     ])?;
 
     Ok(())

zcash_client_sqlite/src/wallet/init/migrations/receiving_key_scopes.rs

@@ -669,6 +669,10 @@ mod tests {
                         block.block_time(),
                         block.sapling().final_tree_size(),
                         block.sapling().commitments().len().try_into().unwrap(),
+                        #[cfg(feature = "orchard")]
+                        block.orchard().final_tree_size(),
+                        #[cfg(feature = "orchard")]
+                        block.orchard().commitments().len().try_into().unwrap(),
                     )?;
 
                     for tx in block.transactions() {