From 3a8b1da26be30540e1c0ad6576514b08a1915b86 Mon Sep 17 00:00:00 2001 From: Eric Tu Date: Wed, 21 Aug 2024 17:56:05 -0400 Subject: [PATCH] fix put block and implement get orchard nullifiers --- zcash_client_memory/src/mem_wallet/mod.rs | 2 + .../src/mem_wallet/wallet_read.rs | 35 +++- .../src/mem_wallet/wallet_write.rs | 159 ++++++++---------- 3 files changed, 109 insertions(+), 87 deletions(-) diff --git a/zcash_client_memory/src/mem_wallet/mod.rs b/zcash_client_memory/src/mem_wallet/mod.rs index 1ebe2319e8..f25041aa9d 100644 --- a/zcash_client_memory/src/mem_wallet/mod.rs +++ b/zcash_client_memory/src/mem_wallet/mod.rs @@ -101,6 +101,7 @@ pub struct MemoryWalletDb { sapling_spends: BTreeMap, #[cfg(feature = "orchard")] orchard_spends: BTreeMap, + sapling_tree: ShardTree< MemoryShardStore, { SAPLING_SHARD_HEIGHT * 2 }, @@ -192,6 +193,7 @@ pub struct Account { viewing_key: ViewingKey, birthday: AccountBirthday, purpose: AccountPurpose, + notes: HashSet, } impl Account { diff --git a/zcash_client_memory/src/mem_wallet/wallet_read.rs b/zcash_client_memory/src/mem_wallet/wallet_read.rs index 154975003f..62cc05b526 100644 --- a/zcash_client_memory/src/mem_wallet/wallet_read.rs +++ b/zcash_client_memory/src/mem_wallet/wallet_read.rs @@ -219,9 +219,40 @@ impl WalletRead for MemoryWalletDb { #[cfg(feature = "orchard")] fn get_orchard_nullifiers( &self, - _query: NullifierQuery, + query: NullifierQuery, ) -> Result, Self::Error> { - Ok(Vec::new()) + Ok(self + .orchard_spends + .iter() + .filter_map(|(nf, (txid, spent))| match query { + NullifierQuery::Unspent => { + if !spent { + Some((txid, self.tx_idx.get(txid).unwrap(), *nf)) + } else { + None + } + } + NullifierQuery::All => Some((txid, self.tx_idx.get(txid).unwrap(), *nf)), + }) + .map(|(txid, height, nf)| { + self.blocks + .get(height) + .and_then(|block| block.transactions.get(txid)) + .and_then(|tx| { + tx.orchard_outputs() + .iter() + .find(|o| o.nf() == Some(&nf)) + .map(|o| (*o.account_id(), *o.nf().unwrap())) + .or_else(|| { + tx.orchard_spends() + .iter() + .find(|s| s.nf() == &nf) + .map(|s| (*s.account_id(), *s.nf())) + }) + }) + }) + .flatten() + .collect()) } #[cfg(feature = "transparent-inputs")] diff --git a/zcash_client_memory/src/mem_wallet/wallet_write.rs b/zcash_client_memory/src/mem_wallet/wallet_write.rs index d4fd3349a9..a7256ae3c1 100644 --- a/zcash_client_memory/src/mem_wallet/wallet_write.rs +++ b/zcash_client_memory/src/mem_wallet/wallet_write.rs @@ -71,6 +71,7 @@ impl WalletWrite for MemoryWalletDb { viewing_key: ViewingKey::Full(Box::new(ufvk)), birthday: birthday.clone(), purpose: AccountPurpose::Spending, + notes: HashSet::new(), }; let id = account.id(); self.accounts.push(account); @@ -107,97 +108,83 @@ impl WalletWrite for MemoryWalletDb { let mut memos = HashMap::new(); for transaction in block.transactions().iter() { let txid = transaction.txid(); - for account_id in self.get_account_ids()? { - let ufvk = self - .get_account(account_id)? - .ok_or(Error::AccountUnknown(account_id))? - .ufvk() - .ok_or(Error::ViewingKeyNotFound(account_id))? - .clone(); - let dfvk = ufvk - .sapling() - .ok_or(Error::ViewingKeyNotFound(account_id))?; - let nk = dfvk.to_nk(Scope::External); - - transaction.sapling_outputs().iter().map(|o| { - // Insert the Sapling nullifiers of the spent notes into the `sapling_spends` map. - let nullifier = o.note().nf(&nk, o.note_commitment_tree_position().into()); + transaction.sapling_outputs().iter().map(|o| { + // Insert the Sapling nullifiers of the spent notes into the `sapling_spends` map. + if let Some(nullifier) = o.nf() { self.sapling_spends - .entry(nullifier) + .entry(*nullifier) .or_insert((txid, false)); + } - // Insert the memo into the `memos` map. - let note_id = NoteId::new( - txid, - Sapling, - u16::try_from(o.index()) - .expect("output indices are representable as u16"), - ); - if let Ok(Some(memo)) = self.get_memo(note_id) { - memos.insert(note_id, memo.encode()); - } - }); - - #[cfg(feature = "orchard")] - transaction.orchard_outputs().iter().map(|o| { - // Insert the Orchard nullifiers of the spent notes into the `orchard_spends` map. - if let Some(nullifier) = o.nf() { - self.orchard_spends - .entry(*nullifier) - .or_insert((txid, false)); - } - - // Insert the memo into the `memos` map. - let note_id = NoteId::new( - txid, - Orchard, - u16::try_from(o.index()) - .expect("output indices are representable as u16"), - ); - if let Ok(Some(memo)) = self.get_memo(note_id) { - memos.insert(note_id, memo.encode()); - } - }); - - // Add frontier to the sapling tree - self.sapling_tree.insert_frontier( - from_state.final_sapling_tree().clone(), - Retention::Checkpoint { - id: from_state.block_height(), - marking: Marking::Reference, - }, + // Insert the memo into the `memos` map. + let note_id = NoteId::new( + txid, + Sapling, + u16::try_from(o.index()).expect("output indices are representable as u16"), ); + if let Ok(Some(memo)) = self.get_memo(note_id) { + memos.insert(note_id, memo.encode()); + } + }); + + #[cfg(feature = "orchard")] + transaction.orchard_outputs().iter().map(|o| { + // Insert the Orchard nullifiers of the spent notes into the `orchard_spends` map. + if let Some(nullifier) = o.nf() { + self.orchard_spends + .entry(*nullifier) + .or_insert((txid, 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(), - marking: Marking::Reference, - }, + // Insert the memo into the `memos` map. + let note_id = NoteId::new( + txid, + Orchard, + u16::try_from(o.index()).expect("output indices are representable as u16"), ); - - // Mark the Sapling nullifiers of the spent notes as spent in the `sapling_spends` map. - transaction.sapling_spends().iter().map(|s| { - let nullifier = s.nf(); - if let Some((txid, spent)) = self.sapling_spends.get_mut(nullifier) { - *spent = true; - } - }); - - #[cfg(feature = "orchard")] - // Mark the Orchard nullifiers of the spent notes as spent in the `orchard_spends` map. - transaction.orchard_spends().iter().map(|s| { - let nullifier = s.nf(); - if let Some((txid, spent)) = self.orchard_spends.get_mut(nullifier) { - *spent = true; - } - }); - - self.tx_idx.insert(txid, block.height()); - transactions.insert(txid, transaction.clone()); - } + if let Ok(Some(memo)) = self.get_memo(note_id) { + memos.insert(note_id, memo.encode()); + } + }); + + // Add frontier to the sapling tree + self.sapling_tree.insert_frontier( + from_state.final_sapling_tree().clone(), + Retention::Checkpoint { + id: from_state.block_height(), + marking: Marking::Reference, + }, + ); + + #[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(), + marking: Marking::Reference, + }, + ); + + // Mark the Sapling nullifiers of the spent notes as spent in the `sapling_spends` map. + transaction.sapling_spends().iter().map(|s| { + let nullifier = s.nf(); + if let Some((txid, spent)) = self.sapling_spends.get_mut(nullifier) { + *spent = true; + } + }); + + #[cfg(feature = "orchard")] + // Mark the Orchard nullifiers of the spent notes as spent in the `orchard_spends` map. + transaction.orchard_spends().iter().map(|s| { + let nullifier = s.nf(); + if let Some((txid, spent)) = self.orchard_spends.get_mut(nullifier) { + *spent = true; + } + }); + + self.tx_idx.insert(txid, block.height()); + transactions.insert(txid, transaction.clone()); } let memory_block = MemoryWalletBlock { @@ -283,6 +270,7 @@ impl WalletWrite for MemoryWalletDb { viewing_key: ViewingKey::Full(Box::new(ufvk)), birthday: birthday.clone(), purpose: AccountPurpose::Spending, + notes: HashSet::new(), }; // TODO: Do we need to check if duplicate? self.accounts.push(account.clone()); @@ -301,6 +289,7 @@ impl WalletWrite for MemoryWalletDb { viewing_key: ViewingKey::Full(Box::new(unified_key.to_owned())), birthday: birthday.clone(), purpose, + notes: HashSet::new(), }; Ok(account) }