WIP: zcash_client_sqlite: Add Orchard support.

Fixes #1176

zcash_client_sqlite/src/lib.rs

@@ -115,6 +115,7 @@ pub(crate) const PRUNING_DEPTH: u32 = 100;
 pub(crate) const VERIFY_LOOKAHEAD: u32 = 10;
 
 pub(crate) const SAPLING_TABLES_PREFIX: &str = "sapling";
+pub(crate) const ORCHARD_TABLES_PREFIX: &str = "orchard";
 
 #[cfg(not(feature = "transparent-inputs"))]
 pub(crate) const UA_TRANSPARENT: bool = false;
@@ -902,7 +903,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 +914,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;

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

@@ -0,0 +1,133 @@
+//! 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_primitives::consensus::{self, BlockHeight};
+
+use crate::wallet::{
+    block_height_extrema,
+    init::{migrations::received_notes_nullable_nf, WalletMigrationError},
+};
+
+pub(super) const MIGRATION_ID: Uuid = Uuid::from_u128(0x3a6487f7_e068_42bb_9d12_6bb8dbe6da00);
+
+pub(super) struct Migration<P> {
+    pub(super) params: P,
+    pub(super) orchard_init_height: BlockHeight,
+}
+
+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)
+            );",
+        )?;
+
+        let _block_height_extrema = block_height_extrema(transaction)?;
+
+        // If a scan range exists that contains the Orchard init height, split it in two at the
+        // init height.
+        if let Some((start, end, priority_code)) = 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(self.orchard_init_height)],
+                |row| {
+                    let start = BlockHeight::from(row.get::<_, u32>(0)?);
+                    let end = BlockHeight::from(row.get::<_, u32>(1)?);
+                    let priority_code: i64 = row.get(2)?;
+                    Ok((start, end, priority_code))
+                },
+            )
+            .optional()?
+        {
+            transaction.execute(
+                "DELETE from scan_queue WHERE block_range_start = :start",
+                named_params![":start": u32::from(start)],
+            )?;
+            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(self.orchard_init_height),
+                   ":priority": priority_code,
+                ],
+            )?;
+            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(self.orchard_init_height),
+                   ":block_range_end": u32::from(end),
+                   ":priority": priority_code,
+                ],
+            )?;
+        }
+
+        Ok(())
+    }
+
+    fn down(&self, _transaction: &rusqlite::Transaction) -> Result<(), WalletMigrationError> {
+        // TODO: something better than just panic?
+        panic!("Cannot revert this migration.");
+    }
+}

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};