Basic put_block function for the memory wallet

Cargo.toml

@@ -52,7 +52,7 @@ bitvec = "1"
 blake2s_simd = "1"
 bls12_381 = "0.8"
 jubjub = "0.10"
-sapling = { package = "sapling-crypto", version = "0.1.2" }
+sapling = { package = "sapling-crypto", version = "0.1.3" }
 
 # - Orchard
 nonempty = "0.7"

zcash_client_backend/src/data_api/mem_wallet.rs

@@ -1,15 +1,16 @@
 #![allow(unused)]
-use incrementalmerkletree::Address;
+use incrementalmerkletree::{Address, Retention};
 use secrecy::{ExposeSecret, SecretVec};
 use shardtree::{error::ShardTreeError, store::memory::MemoryShardStore, ShardTree};
 use std::{
     cmp::Ordering,
     collections::{BTreeMap, HashMap, HashSet},
     convert::Infallible,
-    num::NonZeroU32, hash::Hash,
+    hash::Hash,
+    num::NonZeroU32,
 };
 use zcash_keys::keys::{AddressGenerationError, DerivationError};
-use zip32::{fingerprint::SeedFingerprint, DiversifierIndex};
+use zip32::{fingerprint::SeedFingerprint, DiversifierIndex, Scope};
 
 use zcash_primitives::{
     block::BlockHash,
@@ -25,7 +26,7 @@ use zcash_protocol::{
 use crate::{
     address::UnifiedAddress,
     keys::{UnifiedAddressRequest, UnifiedFullViewingKey, UnifiedSpendingKey},
-    wallet::{NoteId, WalletTransparentOutput, WalletTx},
+    wallet::{NoteId, WalletSpend, WalletTransparentOutput, WalletTx},
 };
 
 use super::{
@@ -100,10 +101,7 @@ pub struct MemoryWalletDb {
 }
 
 impl MemoryWalletDb {
-    pub fn new(
-        network: Network,
-        max_checkpoints: usize
-    ) -> Self {
+    pub fn new(network: Network, max_checkpoints: usize) -> Self {
         Self {
             network,
             accounts: BTreeMap::new(),
@@ -122,6 +120,7 @@ impl MemoryWalletDb {
 #[derive(Debug)]
 pub enum Error {
     AccountUnknown(u32),
+    ViewingKeyNotFound(u32),
     MemoDecryption(memo::Error),
     KeyDerivation(DerivationError),
     AddressGeneration(AddressGenerationError),
@@ -379,12 +378,98 @@ impl WalletWrite for MemoryWalletDb {
         todo!()
     }
 
+    /// Adds a sequence of blocks to the data store.
+    ///
+    /// Assumes blocks will be here in order.
     fn put_blocks(
         &mut self,
-        _from_state: &super::chain::ChainState,
-        _blocks: Vec<ScannedBlock<Self::AccountId>>,
+        from_state: &super::chain::ChainState,
+        blocks: Vec<ScannedBlock<Self::AccountId>>,
     ) -> Result<(), Self::Error> {
-        todo!()
+        // TODO:
+        // - Make sure blocks are coming in order.
+        // - Make sure the first block in the sequence is tip + 1?
+        // - Add a check to make sure the blocks are not already in the data store.
+        for block in blocks.into_iter() {
+            let mut transactions = HashMap::new();
+            for transaction in block.transactions().into_iter().cloned() {
+                let txid = transaction.txid();
+                let account_id = 0; // TODO: Assuming the account is 0, handle this accordingly.
+                let ufvk = self
+                    .accounts
+                    .get(&account_id)
+                    .ok_or(Error::AccountUnknown(0))?
+                    .ufvk
+                    .sapling()
+                    .ok_or(Error::ViewingKeyNotFound(0))?;
+                let nk = ufvk.to_nk(Scope::External);
+
+                // Insert the Sapling nullifiers of the spent notes into the `sapling_spends` map.
+                transaction.sapling_outputs().iter().map(|o| {
+                    let nullifier = o.note().nf(&nk, o.note_commitment_tree_position().into());
+                    // TODO: Populate the bool field properly.
+                    self.sapling_spends.entry(nullifier).or_insert((txid, true));
+                });
+
+                #[cfg(feature = "orchard")]
+                // Insert the Orchard nullifiers of the spent notes into the `orchard_spends` map.
+                transaction.orchard_outputs().iter().map(|o| {
+                    if let Some(nullifier) = o.nf() {
+                        self.orchard_spends
+                            .entry(*nullifier)
+                            .or_insert((txid, true));
+                    }
+                });
+
+                // Add frontier to the sapling tree
+                self.sapling_tree.insert_frontier(
+                    from_state.final_sapling_tree().clone(),
+                    Retention::Checkpoint {
+                        id: from_state.block_height(),
+                        is_marked: false,
+                    },
+                );
+
+                #[cfg(feature = "orchard")]
+                // Add frontier to the orchard tree
+                self.orchard_tree.insert_frontier(
+                    from_state.final_orchard_tree().clone(),
+                    Retention::Checkpoint {
+                        id: from_state.block_height(),
+                        is_marked: false,
+                    },
+                );
+
+                // TODO: Is `self.tx_idx` field filled with all the transaction ids from the scanned blocks ?
+                self.tx_idx.insert(txid, block.block_height);
+                transactions.insert(txid, transaction);
+            }
+
+            let memory_block = MemoryWalletBlock {
+                height: block.block_height,
+                hash: block.block_hash,
+                block_time: block.block_time,
+                transactions,
+                // TODO: Add memos
+                memos: HashMap::new(),
+            };
+
+            self.blocks.insert(block.block_height, memory_block);
+
+            // Add the sapling commitments to the sapling tree.
+            let sapling_block_commitments = block.into_commitments().sapling;
+            sapling_block_commitments.iter().map(|(node, height)| {
+                self.sapling_tree.append(*node, *height);
+            });
+
+            // TODO: Add orchard commitments to the orchard tree.
+
+            // TODO: Received notes need to be made available for note selection & balance calculation
+
+            // TODO: Spent notes need to be made unavailable for note selection & balance calculation
+        }
+
+        Ok(())
     }
 
     /// Adds a transparent UTXO received by the wallet to the data store.

zcash_client_backend/src/wallet.rs

@@ -98,6 +98,7 @@ impl<AccountId, N> Recipient<AccountId, Option<N>> {
 /// The shielded subset of a [`Transaction`]'s data that is relevant to a particular wallet.
 ///
 /// [`Transaction`]: zcash_primitives::transaction::Transaction
+#[derive(Clone)]
 pub struct WalletTx<AccountId> {
     txid: TxId,
     block_index: usize,
@@ -227,6 +228,7 @@ impl transparent_fees::InputView for WalletTransparentOutput {
 }
 
 /// A reference to a spent note belonging to the wallet within a transaction.
+#[derive(Clone)]
 pub struct WalletSpend<Nf, AccountId> {
     index: usize,
     nf: Nf,
@@ -266,6 +268,7 @@ pub type WalletSaplingSpend<AccountId> = WalletSpend<sapling::Nullifier, Account
 pub type WalletOrchardSpend<AccountId> = WalletSpend<orchard::note::Nullifier, AccountId>;
 
 /// An output that was successfully decrypted in the process of wallet scanning.
+#[derive(Clone)]
 pub struct WalletOutput<Note, Nullifier, AccountId> {
     index: usize,
     ephemeral_key: EphemeralKeyBytes,