WIP: Add Orchard note commitment tree.

zcash_client_sqlite/Cargo.toml

@@ -24,6 +24,7 @@ zcash_client_backend = { workspace = true, features = ["unstable-serialization",
 zcash_encoding.workspace = true
 zcash_keys = { workspace = true, features = ["orchard", "sapling"] }
 zcash_primitives.workspace = true
+zcash_protocol.workspace = true
 zip32.workspace = true
 
 # Dependencies exposed in a public API:

zcash_client_sqlite/src/lib.rs

@@ -115,6 +115,8 @@ pub(crate) const PRUNING_DEPTH: u32 = 100;
 pub(crate) const VERIFY_LOOKAHEAD: u32 = 10;
 
 pub(crate) const SAPLING_TABLES_PREFIX: &str = "sapling";
+#[cfg(feature = "orchard")]
+pub(crate) const ORCHARD_TABLES_PREFIX: &str = "orchard";
 
 #[cfg(not(feature = "transparent-inputs"))]
 pub(crate) const UA_TRANSPARENT: bool = false;
@@ -576,9 +578,12 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
                 )?;
 
                 note_positions.extend(block.transactions().iter().flat_map(|wtx| {
-                    wtx.sapling_outputs()
-                        .iter()
-                        .map(|out| out.note_commitment_tree_position())
+                    wtx.sapling_outputs().iter().map(|out| {
+                        (
+                            ShieldedProtocol::Sapling,
+                            out.note_commitment_tree_position(),
+                        )
+                    })
                 }));
 
                 last_scanned_height = Some(block.height());
@@ -902,7 +907,7 @@ impl<P: consensus::Parameters> WalletCommitmentTrees for WalletDb<rusqlite::Conn
     >;
 
     #[cfg(feature = "orchard")]
-    fn with_orchard_tree_mut<F, A, E>(&mut self, _callback: F) -> Result<A, E>
+    fn with_orchard_tree_mut<F, A, E>(&mut self, mut callback: F) -> Result<A, E>
     where
         for<'a> F: FnMut(
             &'a mut ShardTree<
@@ -913,16 +918,41 @@ impl<P: consensus::Parameters> WalletCommitmentTrees for WalletDb<rusqlite::Conn
         ) -> Result<A, E>,
         E: From<ShardTreeError<Self::Error>>,
     {
-        todo!()
+        let tx = self
+            .conn
+            .transaction()
+            .map_err(|e| ShardTreeError::Storage(commitment_tree::Error::Query(e)))?;
+        let shard_store = SqliteShardStore::from_connection(&tx, ORCHARD_TABLES_PREFIX)
+            .map_err(|e| ShardTreeError::Storage(commitment_tree::Error::Query(e)))?;
+        let result = {
+            let mut shardtree = ShardTree::new(shard_store, PRUNING_DEPTH.try_into().unwrap());
+            callback(&mut shardtree)?
+        };
+
+        tx.commit()
+            .map_err(|e| ShardTreeError::Storage(commitment_tree::Error::Query(e)))?;
+        Ok(result)
     }
 
     #[cfg(feature = "orchard")]
     fn put_orchard_subtree_roots(
         &mut self,
-        _start_index: u64,
-        _roots: &[CommitmentTreeRoot<orchard::tree::MerkleHashOrchard>],
+        start_index: u64,
+        roots: &[CommitmentTreeRoot<orchard::tree::MerkleHashOrchard>],
     ) -> Result<(), ShardTreeError<Self::Error>> {
-        todo!()
+        let tx = self
+            .conn
+            .transaction()
+            .map_err(|e| ShardTreeError::Storage(commitment_tree::Error::Query(e)))?;
+        put_shard_roots::<_, { ORCHARD_SHARD_HEIGHT * 2 }, ORCHARD_SHARD_HEIGHT>(
+            &tx,
+            ORCHARD_TABLES_PREFIX,
+            start_index,
+            roots,
+        )?;
+        tx.commit()
+            .map_err(|e| ShardTreeError::Storage(commitment_tree::Error::Query(e)))?;
+        Ok(())
     }
 }
 

zcash_client_sqlite/src/wallet/init.rs

@@ -11,9 +11,8 @@ use uuid::Uuid;
 use zcash_client_backend::keys::AddressGenerationError;
 use zcash_primitives::{consensus, transaction::components::amount::BalanceError};
 
-use crate::WalletDb;
-
 use super::commitment_tree;
+use crate::WalletDb;
 
 mod migrations;
 

zcash_client_sqlite/src/wallet/init/migrations.rs

@@ -5,6 +5,7 @@ mod addresses_table;
 mod full_account_ids;
 mod initial_setup;
 mod nullifier_map;
+mod orchard_shardtree;
 mod received_notes_nullable_nf;
 mod receiving_key_scopes;
 mod sapling_memo_consistency;
@@ -24,7 +25,7 @@ use std::rc::Rc;
 
 use schemer_rusqlite::RusqliteMigration;
 use secrecy::SecretVec;
-use zcash_primitives::consensus;
+use zcash_protocol::consensus;
 
 use super::WalletMigrationError;
 
@@ -98,5 +99,8 @@ pub(super) fn all_migrations<P: consensus::Parameters + 'static>(
             params: params.clone(),
         }),
         Box::new(full_account_ids::Migration { seed }),
+        Box::new(orchard_shardtree::Migration {
+            params: params.clone(),
+        }),
     ]
 }

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

@@ -0,0 +1,154 @@
+//! This migration adds tables to the wallet database that are needed to persist Orchard note
+//! commitment tree data using the `shardtree` crate.
+
+use std::collections::HashSet;
+
+use rusqlite::{self, named_params, OptionalExtension};
+use schemer;
+use schemer_rusqlite::RusqliteMigration;
+
+use tracing::debug;
+use uuid::Uuid;
+
+use zcash_client_backend::data_api::scanning::ScanPriority;
+use zcash_protocol::consensus::{self, BlockHeight, NetworkUpgrade};
+
+use crate::wallet::{
+    init::{migrations::received_notes_nullable_nf, WalletMigrationError},
+    scan_queue_extrema,
+    scanning::priority_code,
+};
+
+pub(super) const MIGRATION_ID: Uuid = Uuid::from_u128(0x3a6487f7_e068_42bb_9d12_6bb8dbe6da00);
+
+pub(super) struct Migration<P> {
+    pub(super) params: P,
+}
+
+impl<P> schemer::Migration for Migration<P> {
+    fn id(&self) -> Uuid {
+        MIGRATION_ID
+    }
+
+    fn dependencies(&self) -> HashSet<Uuid> {
+        [received_notes_nullable_nf::MIGRATION_ID]
+            .into_iter()
+            .collect()
+    }
+
+    fn description(&self) -> &'static str {
+        "Add support for storage of Orchard note commitment tree data using the `shardtree` crate."
+    }
+}
+
+impl<P: consensus::Parameters> RusqliteMigration for Migration<P> {
+    type Error = WalletMigrationError;
+
+    fn up(&self, transaction: &rusqlite::Transaction) -> Result<(), WalletMigrationError> {
+        // Add shard persistence
+        debug!("Creating tables for Orchard shard persistence");
+        transaction.execute_batch(
+            "CREATE TABLE orchard_tree_shards (
+                shard_index INTEGER PRIMARY KEY,
+                subtree_end_height INTEGER,
+                root_hash BLOB,
+                shard_data BLOB,
+                contains_marked INTEGER,
+                CONSTRAINT root_unique UNIQUE (root_hash)
+            );
+            CREATE TABLE orchard_tree_cap (
+                -- cap_id exists only to be able to take advantage of `ON CONFLICT`
+                -- upsert functionality; the table will only ever contain one row
+                cap_id INTEGER PRIMARY KEY,
+                cap_data BLOB NOT NULL
+            );",
+        )?;
+
+        // Add checkpoint persistence
+        debug!("Creating tables for checkpoint persistence");
+        transaction.execute_batch(
+            "CREATE TABLE orchard_tree_checkpoints (
+                checkpoint_id INTEGER PRIMARY KEY,
+                position INTEGER
+            );
+            CREATE TABLE orchard_tree_checkpoint_marks_removed (
+                checkpoint_id INTEGER NOT NULL,
+                mark_removed_position INTEGER NOT NULL,
+                FOREIGN KEY (checkpoint_id) REFERENCES orchard_tree_checkpoints(checkpoint_id)
+                ON DELETE CASCADE,
+                CONSTRAINT spend_position_unique UNIQUE (checkpoint_id, mark_removed_position)
+            );",
+        )?;
+
+        // Treat the current best-known chain tip height as the height to use for Orchard
+        // initialization, bounded below by NU5 activation.
+        if let Some(orchard_init_height) = scan_queue_extrema(transaction)?.and_then(|r| {
+            self.params
+                .activation_height(NetworkUpgrade::Nu5)
+                .map(|orchard_activation| std::cmp::max(orchard_activation, *r.end()))
+        }) {
+            // If a scan range exists that contains the Orchard init height, split it in two at the
+            // init height.
+            if let Some((start, end, range_priority)) = transaction
+                .query_row_and_then(
+                    "SELECT block_range_start, block_range_end, priority 
+                     FROM scan_queue
+                     WHERE block_range_start <= :orchard_init_height
+                     AND block_range_end > :orchard_init_height",
+                    named_params![":orchard_init_height": u32::from(orchard_init_height)],
+                    |row| {
+                        let start = BlockHeight::from(row.get::<_, u32>(0)?);
+                        let end = BlockHeight::from(row.get::<_, u32>(1)?);
+                        let range_priority: i64 = row.get(2)?;
+                        Ok((start, end, range_priority))
+                    },
+                )
+                .optional()?
+            {
+                transaction.execute(
+                    "DELETE from scan_queue WHERE block_range_start = :start",
+                    named_params![":start": u32::from(start)],
+                )?;
+                // Rewrite the start of the scan range to be exactly what it was prior to the
+                // change.
+                transaction.execute(
+                    "INSERT INTO scan_queue (block_range_start, block_range_end, priority)
+                     VALUES (:block_range_start, :block_range_end, :priority)",
+                    named_params![
+                        ":block_range_start": u32::from(start),
+                        ":block_range_end": u32::from(orchard_init_height),
+                        ":priority": range_priority,
+                    ],
+                )?;
+                // Rewrite the remainder of the range to have priority `Historic`
+                transaction.execute(
+                    "INSERT INTO scan_queue (block_range_start, block_range_end, priority)
+                     VALUES (:block_range_start, :block_range_end, :priority)",
+                    named_params![
+                        ":block_range_start": u32::from(orchard_init_height),
+                        ":block_range_end": u32::from(end),
+                        ":priority": priority_code(&ScanPriority::Historic),
+                    ],
+                )?;
+                // Rewrite any scanned ranges above the end of the first Orchard
+                // range to have priority `Historic`
+                transaction.execute(
+                    "UPDATE scan_queue SET priority = :historic
+                     WHERE :block_range_start >= :orchard_initial_range_end
+                     AND priority = :scanned",
+                    named_params![
+                        ":historic": priority_code(&ScanPriority::Historic),
+                        ":orchard_initial_range_end": u32::from(end),
+                        ":scanned": priority_code(&ScanPriority::Scanned),
+                    ],
+                )?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn down(&self, _transaction: &rusqlite::Transaction) -> Result<(), WalletMigrationError> {
+        Err(WalletMigrationError::CannotRevert(MIGRATION_ID))
+    }
+}

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

@@ -1,6 +1,6 @@
-//! This migration adds tables to the wallet database that are needed to persist note commitment
-//! tree data using the `shardtree` crate, and migrates existing witness data into these data
-//! structures.
+//! This migration adds tables to the wallet database that are needed to persist Sapling note
+//! commitment tree data using the `shardtree` crate, and migrates existing witness data into these
+//! data structures.
 
 use std::collections::{BTreeSet, HashSet};