Skip to content

Commit

Permalink
Merge pull request #3 from ChainSafe/ec2/fix-put-block
Browse files Browse the repository at this point in the history 
fix put block and implement get orchard nullifiers
  • Loading branch information
@ec2
ec2 authored Aug 21, 2024
2 parents 8a46ea9 + 3a8b1da commit 2c27c59
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 87 deletions.
2 changes: 2 additions & 0 deletions zcash_client_memory/src/mem_wallet/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@ pub struct MemoryWalletDb {
sapling_spends: BTreeMap<sapling::Nullifier, (TxId, bool)>,
#[cfg(feature = "orchard")]
orchard_spends: BTreeMap<orchard::note::Nullifier, (TxId, bool)>,

sapling_tree: ShardTree<
MemoryShardStore<sapling::Node, BlockHeight>,
{ SAPLING_SHARD_HEIGHT * 2 },
Expand Down Expand Up @@ -192,6 +193,7 @@ pub struct Account {
viewing_key: ViewingKey,
birthday: AccountBirthday,
purpose: AccountPurpose,
notes: HashSet<NoteId>,
}

impl Account {
Expand Down
35 changes: 33 additions & 2 deletions zcash_client_memory/src/mem_wallet/wallet_read.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -219,9 +219,40 @@ impl WalletRead for MemoryWalletDb {
#[cfg(feature = "orchard")]
fn get_orchard_nullifiers(
&self,
_query: NullifierQuery,
query: NullifierQuery,
) -> Result<Vec<(Self::AccountId, orchard::note::Nullifier)>, 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")]
Expand Down
159 changes: 74 additions & 85 deletions zcash_client_memory/src/mem_wallet/wallet_write.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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());
Expand All @@ -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)
}
Expand Down

0 comments on commit 2c27c59

Please sign in to comment.