zcash_client_sqlite: Support Orchard scanning

zcash_client_sqlite/src/lib.rs

@@ -545,6 +545,10 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
                     for spend in tx.sapling_spends() {
                         wallet::sapling::mark_sapling_note_spent(wdb.conn.0, tx_row, spend.nf())?;
                     }
+                    #[cfg(feature = "orchard")]
+                    for spend in tx.orchard_spends() {
+                        wallet::orchard::mark_orchard_note_spent(wdb.conn.0, tx_row, spend.nf())?;
+                    }
 
                     for output in tx.sapling_outputs() {
                         // Check whether this note was spent in a later block range that
@@ -563,6 +567,24 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
 
                         wallet::sapling::put_received_note(wdb.conn.0, output, tx_row, spent_in)?;
                     }
+                    #[cfg(feature = "orchard")]
+                    for output in tx.orchard_outputs() {
+                        // Check whether this note was spent in a later block range that
+                        // we previously scanned.
+                        let spent_in = output
+                            .nf()
+                            .map(|nf| {
+                                wallet::query_nullifier_map::<_, Scope>(
+                                    wdb.conn.0,
+                                    ShieldedProtocol::Orchard,
+                                    &nf.to_bytes(),
+                                )
+                            })
+                            .transpose()?
+                            .flatten();
+
+                        wallet::orchard::put_received_note(wdb.conn.0, output, tx_row, spent_in)?;
+                    }
                 }
 
                 // Insert the new nullifiers from this block into the nullifier map.
@@ -572,14 +594,37 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
                     ShieldedProtocol::Sapling,
                     block.sapling().nullifier_map(),
                 )?;
+                #[cfg(feature = "orchard")]
+                wallet::insert_nullifier_map(
+                    wdb.conn.0,
+                    block.height(),
+                    ShieldedProtocol::Orchard,
+                    &block
+                        .orchard()
+                        .nullifier_map()
+                        .iter()
+                        .map(|(txid, idx, nfs)| {
+                            (*txid, *idx, nfs.iter().map(|nf| nf.to_bytes()).collect())
+                        })
+                        .collect::<Vec<_>>(),
+                )?;
 
                 note_positions.extend(block.transactions().iter().flat_map(|wtx| {
-                    wtx.sapling_outputs().iter().map(|out| {
+                    let iter = wtx.sapling_outputs().iter().map(|out| {
                         (
                             ShieldedProtocol::Sapling,
                             out.note_commitment_tree_position(),
                         )
-                    })
+                    });
+                    #[cfg(feature = "orchard")]
+                    let iter = iter.chain(wtx.orchard_outputs().iter().map(|out| {
+                        (
+                            ShieldedProtocol::Orchard,
+                            out.note_commitment_tree_position(),
+                        )
+                    }));
+
+                    iter
                 }));
 
                 last_scanned_height = Some(block.height());

zcash_client_sqlite/src/wallet.rs

@@ -73,6 +73,7 @@ use std::io::{self, Cursor};
 use std::num::NonZeroU32;
 use std::ops::RangeInclusive;
 use tracing::debug;
+use zcash_client_backend::data_api::ORCHARD_SHARD_HEIGHT;
 use zcash_keys::keys::HdSeedFingerprint;
 
 use zcash_client_backend::{
@@ -100,6 +101,7 @@ use zcash_primitives::{
     zip32::{self, DiversifierIndex, Scope},
 };
 
+use crate::ORCHARD_TABLES_PREFIX;
 use crate::{
     error::SqliteClientError,
     wallet::commitment_tree::{get_max_checkpointed_height, SqliteShardStore},
@@ -260,7 +262,7 @@ pub(crate) fn add_account<P: consensus::Parameters>(
     // If a birthday frontier is available, insert it into the note commitment tree. If the
     // birthday frontier is the empty frontier, we don't need to do anything.
     if let Some(frontier) = birthday.sapling_frontier().value() {
-        debug!("Inserting frontier into ShardTree: {:?}", frontier);
+        debug!("Inserting Sapling frontier into ShardTree: {:?}", frontier);
         let shard_store =
             SqliteShardStore::<_, ::sapling::Node, SAPLING_SHARD_HEIGHT>::from_connection(
                 conn,
@@ -285,6 +287,33 @@ pub(crate) fn add_account<P: consensus::Parameters>(
         )?;
     }
 
+    if let Some(frontier) = birthday.orchard_frontier().value() {
+        debug!("Inserting Orchard frontier into ShardTree: {:?}", frontier);
+        let shard_store = SqliteShardStore::<
+            _,
+            ::orchard::tree::MerkleHashOrchard,
+            ORCHARD_SHARD_HEIGHT,
+        >::from_connection(conn, ORCHARD_TABLES_PREFIX)?;
+        let mut shard_tree: ShardTree<
+            _,
+            { ::orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 },
+            ORCHARD_SHARD_HEIGHT,
+        > = ShardTree::new(shard_store, PRUNING_DEPTH.try_into().unwrap());
+        shard_tree.insert_frontier_nodes(
+            frontier.clone(),
+            Retention::Checkpoint {
+                // This subtraction is safe, because all leaves in the tree appear in blocks, and
+                // the invariant that birthday.height() always corresponds to the block for which
+                // `frontier` is the tree state at the start of the block. Together, this means
+                // there exists a prior block for which frontier is the tree state at the end of
+                // the block.
+                id: birthday.height() - 1,
+                is_marked: false,
+            },
+        )?;
+    }
+
+    // The ignored range always starts at Sapling activation
     let sapling_activation_height = params
         .activation_height(NetworkUpgrade::Sapling)
         .expect("Sapling activation height must be available.");
@@ -1865,13 +1894,20 @@ pub(crate) fn update_expired_notes(
     conn: &rusqlite::Connection,
     expiry_height: BlockHeight,
 ) -> Result<(), SqliteClientError> {
-    let mut stmt_update_expired = conn.prepare_cached(
+    let mut stmt_update_sapling_expired = conn.prepare_cached(
         "UPDATE sapling_received_notes SET spent = NULL WHERE EXISTS (
             SELECT id_tx FROM transactions
             WHERE id_tx = sapling_received_notes.spent AND block IS NULL AND expiry_height < ?
         )",
     )?;
-    stmt_update_expired.execute([u32::from(expiry_height)])?;
+    stmt_update_sapling_expired.execute([u32::from(expiry_height)])?;
+    let mut stmt_update_orchard_expired = conn.prepare_cached(
+        "UPDATE orchard_received_notes SET spent = NULL WHERE EXISTS (
+            SELECT id_tx FROM transactions
+            WHERE id_tx = orchard_received_notes.spent AND block IS NULL AND expiry_height < ?
+        )",
+    )?;
+    stmt_update_orchard_expired.execute([u32::from(expiry_height)])?;
     Ok(())
 }
 

zcash_client_sqlite/src/wallet/orchard.rs

@@ -1,5 +1,5 @@
 use incrementalmerkletree::Position;
-use rusqlite::{named_params, Connection};
+use rusqlite::{named_params, params, Connection};
 
 use zcash_client_backend::{
     data_api::NullifierQuery, wallet::WalletOrchardOutput, DecryptedOutput, TransferType,
@@ -195,3 +195,23 @@ pub(crate) fn get_orchard_nullifiers(
     let res: Vec<_> = nullifiers.collect::<Result<_, _>>()?;
     Ok(res)
 }
+
+/// Marks a given nullifier as having been revealed in the construction
+/// of the specified transaction.
+///
+/// Marking a note spent in this fashion does NOT imply that the
+/// spending transaction has been mined.
+pub(crate) fn mark_orchard_note_spent(
+    conn: &Connection,
+    tx_ref: i64,
+    nf: &orchard::note::Nullifier,
+) -> Result<bool, SqliteClientError> {
+    let mut stmt_mark_orchard_note_spent =
+        conn.prepare_cached("UPDATE orchard_received_notes SET spent = ? WHERE nf = ?")?;
+
+    match stmt_mark_orchard_note_spent.execute(params![tx_ref, nf.to_bytes()])? {
+        0 => Ok(false),
+        1 => Ok(true),
+        _ => unreachable!("nf column is marked as UNIQUE"),
+    }
+}