zec
/
incrementalmerkletree
mirror of https://github.com/zcash/incrementalmerkletree.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Address comments from code review. 

This also provides a minor performance and correctness improvement to
`remove_mark`.
This commit is contained in:
Kris Nuttycombe 2023-05-24 09:30:14 -06:00
parent 4efe39eaa6
commit 2b8e2d62fa
1 changed files with 34 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
shardtree/src/lib.rs
View File

@ -1797,7 +1797,7 @@ where
pub fn get_marked_leaf(&self, position: Position) -> Result<Option<H>, S::Error> { pub fn get_marked_leaf(&self, position: Position) -> Result<Option<H>, S::Error> {
Ok(self Ok(self
.store .store
.get_shard(Address::above_position(Self::subtree_level(), position))? .get_shard(Self::subtree_addr(position))?
.and_then(|t| t.value_at_position(position).cloned()) .and_then(|t| t.value_at_position(position).cloned())
.and_then(|(v, r)| if r.is_marked() { Some(v) } else { None })) .and_then(|(v, r)| if r.is_marked() { Some(v) } else { None }))
} }
@ -1919,7 +1919,7 @@ where
values: I, values: I,
) -> Result<Option<(Position, Vec<IncompleteAt>)>, S::Error> { ) -> Result<Option<(Position, Vec<IncompleteAt>)>, S::Error> {
let mut values = values.peekable(); let mut values = values.peekable();
let mut subtree_root_addr = Address::above_position(Self::subtree_level(), start); let mut subtree_root_addr = Self::subtree_addr(start);
let mut max_insert_position = None; let mut max_insert_position = None;
let mut all_incomplete = vec![]; let mut all_incomplete = vec![];
loop { loop {
@ -2065,7 +2065,7 @@ where
checkpoints_to_delete.push(cid.clone()); checkpoints_to_delete.push(cid.clone());
let mut clear_at = |pos, flags_to_clear| { let mut clear_at = |pos, flags_to_clear| {
let subtree_addr = Address::above_position(Self::subtree_level(), pos); let subtree_addr = Self::subtree_addr(pos);
clear_positions clear_positions
.entry(subtree_addr) .entry(subtree_addr)
.and_modify(|to_clear| { .and_modify(|to_clear| {
@ -2133,7 +2133,7 @@ where
true true
} }
TreeState::AtPosition(position) => { TreeState::AtPosition(position) => {
let subtree_addr = Address::above_position(Self::subtree_level(), position); let subtree_addr = Self::subtree_addr(position);
let replacement = self let replacement = self
.store .store
.get_shard(subtree_addr)? .get_shard(subtree_addr)?
@ -2318,7 +2318,7 @@ where
Address::from_parts(Level::from(0), position.into()), Address::from_parts(Level::from(0), position.into()),
))) )))
} else { } else {
let subtree_addr = Address::above_position(Self::subtree_level(), position); let subtree_addr = Self::subtree_addr(position);
// compute the witness for the specified position up to the subtree root // compute the witness for the specified position up to the subtree root
let mut witness = self.store.get_shard(subtree_addr)?.map_or_else( let mut witness = self.store.get_shard(subtree_addr)?.map_or_else(
@ -2341,48 +2341,43 @@ where
/// Make a marked leaf at a position eligible to be pruned. /// Make a marked leaf at a position eligible to be pruned.
/// ///
/// If the checkpoint associated with the specified identifier does not exist because the /// If the checkpoint associated with the specified identifier does not exist because the
/// corresponding checkpoint would have been more than `max_checkpoints` deep, the removal /// corresponding checkpoint would have been more than `max_checkpoints` deep, the removal is
/// is recorded as of the first existing checkpoint and the associated leaves will be pruned /// recorded as of the first existing checkpoint and the associated leaves will be pruned when
/// when that checkpoint is subsequently removed. /// that checkpoint is subsequently removed.
///
/// Returns `Ok(true)` if a mark was successfully removed from the leaf at the specified
/// position, `Ok(false)` if the tree does not contain a leaf at the specified position or is
/// not marked, or an error if one is produced by the underlying data store.
pub fn remove_mark( pub fn remove_mark(
&mut self, &mut self,
position: Position, position: Position,
as_of_checkpoint: Option<&C>, as_of_checkpoint: Option<&C>,
) -> Result<bool, S::Error> ) -> Result<bool, S::Error> {
where match self.store.get_shard(Self::subtree_addr(position))? {
C: std::fmt::Debug, Some(shard)
{ if shard
#[allow(clippy::blocks_in_if_conditions)] .value_at_position(position)
if self.get_marked_leaf(position)?.is_some() { .iter()
match as_of_checkpoint { .any(|(_, r)| r.is_marked()) =>
Some(cid) if Some(cid) >= self.store.min_checkpoint_id()?.as_ref() => { {
self.store.update_checkpoint_with(cid, |checkpoint| { match as_of_checkpoint {
checkpoint.marks_removed.insert(position); Some(cid) if Some(cid) >= self.store.min_checkpoint_id()?.as_ref() => {
Ok(()) self.store.update_checkpoint_with(cid, |checkpoint| {
}) checkpoint.marks_removed.insert(position);
} Ok(())
_ => { })
// if no checkpoint was provided, or if the checkpoint is too far in the past, }
// remove the mark directly. _ => {
let cleared = // if no checkpoint was provided, or if the checkpoint is too far in the past,
self.store // remove the mark directly.
.get_shard(Self::subtree_addr(position))? self.store.put_shard(
.map(|subtree| { shard.clear_flags(BTreeMap::from([(position, RetentionFlags::MARKED)])),
subtree.clear_flags(BTreeMap::from([( )?;
position,
RetentionFlags::MARKED,
)]))
});
if let Some(cleared) = cleared {
self.store.put_shard(cleared)?;
Ok(true) Ok(true)
} else {
Ok(false)
} }
} }
} }
} else { _ => Ok(false),
Ok(false)
} }
} }
} }