diff --git a/bridgetree/src/lib.rs b/bridgetree/src/lib.rs index 4614ec0..7cd1c0e 100644 --- a/bridgetree/src/lib.rs +++ b/bridgetree/src/lib.rs @@ -470,13 +470,13 @@ impl BridgeTree { /// /// Panics if `max_checkpoints < 1` because mark/rewind logic depends upon the presence /// of checkpoints to function. - pub fn new(max_checkpoints: usize, initial_checkpoint_id: C) -> Self { + pub fn new(max_checkpoints: usize) -> Self { assert!(max_checkpoints >= 1); Self { prior_bridges: vec![], current_bridge: None, saved: BTreeMap::new(), - checkpoints: VecDeque::from(vec![Checkpoint::at_length(0, initial_checkpoint_id)]), + checkpoints: VecDeque::new(), max_checkpoints, } } @@ -535,12 +535,8 @@ impl BridgeTree { impl BridgeTree { /// Construct a new BridgeTree that will start recording changes from the state of /// the specified frontier. - pub fn from_frontier( - max_checkpoints: usize, - frontier: NonEmptyFrontier, - checkpoint_id: C, - ) -> Self { - let mut bridge = Self { + pub fn from_frontier(max_checkpoints: usize, frontier: NonEmptyFrontier) -> Self { + Self { prior_bridges: vec![], current_bridge: Some(MerkleBridge::from_parts( None, @@ -551,9 +547,7 @@ impl BridgeTree BridgeTree BridgeTree bool { if self.saved.contains_key(&position) { - let c = self - .checkpoints - .back_mut() - .expect("Checkpoints deque must never be empty."); - c.forgotten.insert(position); + if let Some(c) = self.checkpoints.back_mut() { + c.forgotten.insert(position); + } else { + self.saved.remove(&position); + } true } else { false @@ -805,12 +797,7 @@ impl BridgeTree bool { - if self.checkpoints.len() > 1 { - let c = self - .checkpoints - .pop_back() - .expect("Checkpoints deque is known to be non-empty."); - + if let Some(c) = self.checkpoints.pop_back() { // Remove marks for positions that were marked during the lifetime of this checkpoint. for pos in c.marked { self.saved.remove(&pos); @@ -1063,7 +1050,7 @@ mod tests { #[test] fn tree_depth() { - let mut tree = BridgeTree::::new(100, 0); + let mut tree = BridgeTree::::new(100); for c in 'a'..'i' { assert!(tree.append(c.to_string())) } @@ -1095,7 +1082,7 @@ mod tests { { let pos_gen = (0..max_count).prop_map(|p| Position::try_from(p).unwrap()); proptest::collection::vec(arb_operation(item_gen, pos_gen), 0..max_count).prop_map(|ops| { - let mut tree: BridgeTree = BridgeTree::new(10, 0); + let mut tree: BridgeTree = BridgeTree::new(10); for (i, op) in ops.into_iter().enumerate() { apply_operation(&mut tree, op.map_checkpoint_id(|_| i)); } @@ -1130,33 +1117,31 @@ mod tests { #[test] fn root_hashes() { - check_root_hashes(|max_checkpoints| { - BridgeTree::::new(max_checkpoints, 0) - }); + check_root_hashes(BridgeTree::::new); } #[test] fn witness() { - check_witnesses(|max_checkpoints| BridgeTree::::new(max_checkpoints, 0)); + check_witnesses(BridgeTree::::new); } #[test] fn checkpoint_rewind() { check_checkpoint_rewind(|max_checkpoints| { - BridgeTree::::new(max_checkpoints, 0) + BridgeTree::::new(max_checkpoints) }); } #[test] fn rewind_remove_mark() { check_rewind_remove_mark(|max_checkpoints| { - BridgeTree::::new(max_checkpoints, 0) + BridgeTree::::new(max_checkpoints) }); } #[test] fn garbage_collect() { - let mut tree: BridgeTree = BridgeTree::new(1000, 0); + let mut tree: BridgeTree = BridgeTree::new(1000); let empty_root = tree.root(0); tree.append("a".to_string()); for i in 0..100 { @@ -1167,7 +1152,7 @@ mod tests { tree.rewind(); assert!(tree.root(0) != empty_root); - let mut t = BridgeTree::::new(10, 0); + let mut t = BridgeTree::::new(10); let mut to_unmark = vec![]; let mut has_witness = vec![]; for i in 0u64..100 { @@ -1213,8 +1198,8 @@ mod tests { max_checkpoints: usize, ) -> CombinedTree, BridgeTree> { CombinedTree::new( - CompleteTree::::new(max_checkpoints, 0), - BridgeTree::::new(max_checkpoints, 0), + CompleteTree::::new(max_checkpoints), + BridgeTree::::new(max_checkpoints), ) } diff --git a/incrementalmerkletree/src/testing/complete_tree.rs b/incrementalmerkletree/src/testing/complete_tree.rs index f0f375f..e6984ae 100644 --- a/incrementalmerkletree/src/testing/complete_tree.rs +++ b/incrementalmerkletree/src/testing/complete_tree.rs @@ -71,11 +71,11 @@ pub struct CompleteTree { impl CompleteTree { /// Creates a new, empty binary tree - pub fn new(max_checkpoints: usize, initial_checkpoint_id: C) -> Self { + pub fn new(max_checkpoints: usize) -> Self { Self { leaves: vec![], marks: BTreeSet::new(), - checkpoints: BTreeMap::from([(initial_checkpoint_id, Checkpoint::at_length(0))]), + checkpoints: BTreeMap::new(), max_checkpoints, } } @@ -140,22 +140,18 @@ impl CompleteTr /// Marks the current tree state leaf as a value that we're interested in /// marking. Returns the current position if the tree is non-empty. fn mark(&mut self) -> Option { - match self.current_position() { - Some(pos) => { - if !self.marks.contains(&pos) { - self.marks.insert(pos); - self.checkpoints - .iter_mut() - .rev() - .next() - .unwrap() - .1 - .marked - .insert(pos); + if let Some(pos) = self.current_position() { + if !self.marks.contains(&pos) { + self.marks.insert(pos); + + if let Some(checkpoint) = self.checkpoints.values_mut().rev().next() { + checkpoint.marked.insert(pos); } - Some(pos) } - None => None, + + Some(pos) + } else { + None } } @@ -282,14 +278,11 @@ impl bool { if self.marks.contains(&position) { - self.checkpoints - .iter_mut() - .rev() - .next() - .unwrap() - .1 - .forgotten - .insert(position); + if let Some(c) = self.checkpoints.values_mut().rev().next() { + c.forgotten.insert(position); + } else { + self.marks.remove(&position); + } true } else { false @@ -297,7 +290,7 @@ impl bool { - if Some(&id) > self.checkpoints.iter().rev().next().map(|(id, _)| id) { + if Some(&id) > self.checkpoints.keys().rev().next() { Self::checkpoint(self, id, self.current_position()); true } else { @@ -306,8 +299,7 @@ impl bool { - if self.checkpoints.len() > 1 { - let (id, c) = self.checkpoints.iter().rev().next().unwrap(); + if let Some((id, c)) = self.checkpoints.iter().rev().next() { self.leaves.truncate(c.leaves_len); for pos in c.marked.iter() { self.marks.remove(pos); @@ -342,7 +334,7 @@ mod tests { expected = SipHashable::combine(lvl.into(), &expected, &expected); } - let tree = CompleteTree::::new(100, ()); + let tree = CompleteTree::::new(100); assert_eq!(tree.root(0).unwrap(), expected); } @@ -351,7 +343,7 @@ mod tests { const DEPTH: u8 = 3; let values = (0..(1 << DEPTH)).map(SipHashable); - let mut tree = CompleteTree::::new(100, ()); + let mut tree = CompleteTree::::new(100); for value in values { assert!(tree.append(value, Retention::Ephemeral).is_ok()); } @@ -376,21 +368,17 @@ mod tests { #[test] fn append() { - check_append(|max_checkpoints| CompleteTree::::new(max_checkpoints, 0)); + check_append(CompleteTree::::new); } #[test] fn root_hashes() { - check_root_hashes(|max_checkpoints| { - CompleteTree::::new(max_checkpoints, 0) - }); + check_root_hashes(CompleteTree::::new); } #[test] fn witnesses() { - check_witnesses(|max_checkpoints| { - CompleteTree::::new(max_checkpoints, 0) - }); + check_witnesses(CompleteTree::::new); } #[test] @@ -400,7 +388,7 @@ mod tests { const DEPTH: u8 = 3; let values = (0..(1 << DEPTH)).map(SipHashable); - let mut tree = CompleteTree::::new(100, ()); + let mut tree = CompleteTree::::new(100); for value in values { assert!(Tree::append(&mut tree, value, Retention::Marked)); } @@ -435,14 +423,14 @@ mod tests { #[test] fn checkpoint_rewind() { check_checkpoint_rewind(|max_checkpoints| { - CompleteTree::::new(max_checkpoints, 0) + CompleteTree::::new(max_checkpoints) }); } #[test] fn rewind_remove_mark() { check_rewind_remove_mark(|max_checkpoints| { - CompleteTree::::new(max_checkpoints, 0) + CompleteTree::::new(max_checkpoints) }); } } diff --git a/shardtree/proptest-regressions/lib.txt b/shardtree/proptest-regressions/lib.txt index 4054830..ae68750 100644 --- a/shardtree/proptest-regressions/lib.txt +++ b/shardtree/proptest-regressions/lib.txt @@ -16,3 +16,7 @@ cc cf1a33ef6df58bbf7cc199b8b1879c3a078e7784aa3edc0aba9ca03772bea5f2 # shrinks to cc 544e027d994eaf7f97b1c8d9ee7b35522a64a610b1430d56d74ec947018b759d # shrinks to ops = [Append("a", Marked), Append("a", Ephemeral), Append("a", Ephemeral), Append("a", Ephemeral), Append("a", Ephemeral), Append("a", Ephemeral), Append("a", Ephemeral), Checkpoint(()), Append("a", Marked), Checkpoint(()), Checkpoint(()), Append("a", Checkpoint { id: (), is_marked: false }), Rewind, Rewind, Witness(Position(7), 2)] cc 55d00b68a0f0a02f83ab53f18a29d16d0233153b69a01414a1622104e0eead31 # shrinks to ops = [Append("a", Marked), Append("a", Checkpoint { id: (), is_marked: false }), Append("a", Marked), Checkpoint(()), Checkpoint(()), Checkpoint(()), Append("a", Checkpoint { id: (), is_marked: false }), Append("a", Checkpoint { id: (), is_marked: false }), Witness(Position(0), 7)] cc 9dd966ff1ab66965c5b84153ae13f684258560cdd5e84c7deb24f724cb12aba7 # shrinks to ops = [Append("a", Marked), Append("a", Ephemeral), Append("a", Checkpoint { id: (), is_marked: true }), Checkpoint(()), Append("a", Checkpoint { id: (), is_marked: false }), Rewind, Rewind, Append("a", Checkpoint { id: (), is_marked: false }), Append("a", Checkpoint { id: (), is_marked: false }), Checkpoint(()), Witness(Position(2), 4)] +cc d53a73021238de143764ee1d48b944abb93bd4bc54f35d16e514261220d3eb78 # shrinks to ops = [Append(SipHashable(0), Marked), Unmark(Position(0))] +cc d9460b8acbc5b4d112cae5d9e2296fcd793999b2b2e1d5405722f2bd8d176c31 # shrinks to ops = [Append("a", Checkpoint { id: (), is_marked: true }), Rewind, Append("a", Ephemeral), Rewind, Unmark(Position(0))] +cc 644c7763bc7bdc65bd9e6eb156b3b1a9b0632571a571c462bd44f3e04a389ca0 # shrinks to ops = [Append("a", Ephemeral), Append("a", Checkpoint { id: (), is_marked: true }), Append("a", Ephemeral), Append("a", Ephemeral), Unmark(Position(1)), Witness(Position(1), 0)] +cc 12790169d3df4280dd155d9cdfa76719318b8ec97a80bd562b7cb182d4f9bc79 # shrinks to ops = [CurrentPosition, CurrentPosition, Append(SipHashable(0), Ephemeral), Append(SipHashable(0), Marked), Append(SipHashable(0), Ephemeral), Checkpoint(()), Checkpoint(()), Checkpoint(()), Unmark(Position(1)), Checkpoint(()), Witness(Position(1), 0)] diff --git a/shardtree/src/lib.rs b/shardtree/src/lib.rs index 74420df..7a5c3cf 100644 --- a/shardtree/src/lib.rs +++ b/shardtree/src/lib.rs @@ -1761,21 +1761,12 @@ impl< where S::Error: From + From, { - /// Creates a new empty tree and establishes a checkpoint for the empty tree at the given - /// checkpoint identifier. - pub fn empty( - store: S, - max_checkpoints: usize, - initial_checkpoint_id: C, - ) -> Result { - let mut result = Self { + /// Creates a new empty tree. + pub fn empty(store: S, max_checkpoints: usize) -> Self { + Self { store, max_checkpoints, - }; - result - .store - .add_checkpoint(initial_checkpoint_id, Checkpoint::tree_empty())?; - Ok(result) + } } /// Constructs a wrapper around the provided shard store without initialization. @@ -1797,11 +1788,16 @@ where Level::from(SHARD_HEIGHT - 1) } + /// Returns the root address of the subtree that contains the specified position. + pub fn subtree_addr(pos: Position) -> Address { + Address::above_position(Self::subtree_level(), pos) + } + /// Returns the leaf value at the specified position, if it is a marked leaf. pub fn get_marked_leaf(&self, position: Position) -> Result, S::Error> { Ok(self .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(|(v, r)| if r.is_marked() { Some(v) } else { None })) } @@ -1923,7 +1919,7 @@ where values: I, ) -> Result)>, S::Error> { 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 all_incomplete = vec![]; loop { @@ -2031,19 +2027,13 @@ where // Update the rightmost subtree to add the `CHECKPOINT` flag to the right-most leaf (which // need not be a level-0 leaf; it's fine to rewind to a pruned state). if let Some(subtree) = self.store.last_shard()? { - if let Some((replacement, checkpoint_position)) = go(subtree.root_addr, &subtree.root) { - if self - .store - .put_shard(LocatedTree { - root_addr: subtree.root_addr, - root: replacement, - }) - .is_err() - { - return Ok(false); - } + if let Some((replacement, pos)) = go(subtree.root_addr, &subtree.root) { + self.store.put_shard(LocatedTree { + root_addr: subtree.root_addr, + root: replacement, + })?; self.store - .add_checkpoint(checkpoint_id, Checkpoint::at_position(checkpoint_position))?; + .add_checkpoint(checkpoint_id, Checkpoint::at_position(pos))?; // early return once we've updated the tree state self.prune_excess_checkpoints()?; @@ -2075,7 +2065,7 @@ where checkpoints_to_delete.push(cid.clone()); 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 .entry(subtree_addr) .and_modify(|to_clear| { @@ -2132,36 +2122,32 @@ where ) -> Result { Ok(if checkpoint_depth == 0 { true - } else if self.store.checkpoint_count()? > 1 { - if let Some((checkpoint_id, c)) = - self.store.get_checkpoint_at_depth(checkpoint_depth)? - { - match c.tree_state { - TreeState::Empty => { - self.store - .truncate(Address::from_parts(Self::subtree_level(), 0))?; + } else if let Some((checkpoint_id, c)) = + self.store.get_checkpoint_at_depth(checkpoint_depth)? + { + match c.tree_state { + TreeState::Empty => { + self.store + .truncate(Address::from_parts(Self::subtree_level(), 0))?; + self.store.truncate_checkpoints(&checkpoint_id)?; + true + } + TreeState::AtPosition(position) => { + let subtree_addr = Self::subtree_addr(position); + let replacement = self + .store + .get_shard(subtree_addr)? + .and_then(|s| s.truncate_to_position(position)); + + if let Some(truncated) = replacement { + self.store.truncate(subtree_addr)?; + self.store.put_shard(truncated)?; self.store.truncate_checkpoints(&checkpoint_id)?; true - } - TreeState::AtPosition(position) => { - let subtree_addr = Address::above_position(Self::subtree_level(), position); - let replacement = self - .store - .get_shard(subtree_addr)? - .and_then(|s| s.truncate_to_position(position)); - - if let Some(truncated) = replacement { - self.store.truncate(subtree_addr)?; - self.store.put_shard(truncated)?; - self.store.truncate_checkpoints(&checkpoint_id)?; - true - } else { - false - } + } else { + false } } - } else { - false } } else { false @@ -2332,7 +2318,7 @@ where Address::from_parts(Level::from(0), position.into()), ))) } 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 let mut witness = self.store.get_shard(subtree_addr)?.map_or_else( @@ -2355,37 +2341,44 @@ where /// Make a marked leaf at a position eligible to be pruned. /// /// 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 - /// is recorded as of the first existing checkpoint and the associated leaves will be pruned - /// when that checkpoint is subsequently removed. + /// corresponding checkpoint would have been more than `max_checkpoints` deep, the removal is + /// recorded as of the first existing checkpoint and the associated leaves will be pruned when + /// 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( &mut self, position: Position, - as_of_checkpoint: &C, + as_of_checkpoint: Option<&C>, ) -> Result { - #[allow(clippy::blocks_in_if_conditions)] - if self.get_marked_leaf(position)?.is_some() { - if self - .store - .update_checkpoint_with(as_of_checkpoint, |checkpoint| { - checkpoint.marks_removed.insert(position); - Ok(()) - })? + match self.store.get_shard(Self::subtree_addr(position))? { + Some(shard) + if shard + .value_at_position(position) + .iter() + .any(|(_, r)| r.is_marked()) => { - return Ok(true); - } - - if let Some(cid) = self.store.min_checkpoint_id()? { - if self.store.update_checkpoint_with(&cid, |checkpoint| { - checkpoint.marks_removed.insert(position); - Ok(()) - })? { - return Ok(true); + match as_of_checkpoint { + Some(cid) if Some(cid) >= self.store.min_checkpoint_id()?.as_ref() => { + 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. + self.store.put_shard( + shard.clear_flags(BTreeMap::from([(position, RetentionFlags::MARKED)])), + )?; + Ok(true) + } } } + _ => Ok(false), } - - Ok(false) } } @@ -2786,7 +2779,7 @@ mod tests { #[test] fn shardtree_insertion() { let mut tree: ShardTree, 4, 3> = - ShardTree::empty(MemoryShardStore::empty(), 100, 0).unwrap(); + ShardTree::empty(MemoryShardStore::empty(), 100); assert_matches!( tree.batch_insert( Position::from(1), @@ -2915,11 +2908,12 @@ mod tests { } fn remove_mark(&mut self, position: Position) -> bool { - if let Ok(Some(c)) = self.store.max_checkpoint_id() { - ShardTree::remove_mark(self, position, &c).unwrap() - } else { - false - } + ShardTree::remove_mark( + self, + position, + self.store.max_checkpoint_id().unwrap().as_ref(), + ) + .unwrap() } fn checkpoint(&mut self, checkpoint_id: C) -> bool { @@ -2934,60 +2928,35 @@ mod tests { #[test] fn append() { check_append(|m| { - ShardTree::, 4, 3>::empty( - MemoryShardStore::empty(), - m, - 0, - ) - .unwrap() + ShardTree::, 4, 3>::empty(MemoryShardStore::empty(), m) }); } #[test] fn root_hashes() { check_root_hashes(|m| { - ShardTree::, 4, 3>::empty( - MemoryShardStore::empty(), - m, - 0, - ) - .unwrap() + ShardTree::, 4, 3>::empty(MemoryShardStore::empty(), m) }); } #[test] fn witnesses() { check_witnesses(|m| { - ShardTree::, 4, 3>::empty( - MemoryShardStore::empty(), - m, - 0, - ) - .unwrap() + ShardTree::, 4, 3>::empty(MemoryShardStore::empty(), m) }); } #[test] fn checkpoint_rewind() { check_checkpoint_rewind(|m| { - ShardTree::, 4, 3>::empty( - MemoryShardStore::empty(), - m, - 0, - ) - .unwrap() + ShardTree::, 4, 3>::empty(MemoryShardStore::empty(), m) }); } #[test] fn rewind_remove_mark() { check_rewind_remove_mark(|m| { - ShardTree::, 4, 3>::empty( - MemoryShardStore::empty(), - m, - 0, - ) - .unwrap() + ShardTree::, 4, 3>::empty(MemoryShardStore::empty(), m) }); } @@ -3002,8 +2971,8 @@ mod tests { ShardTree, 4, 3>, > { CombinedTree::new( - CompleteTree::new(max_checkpoints, 0), - ShardTree::empty(MemoryShardStore::empty(), max_checkpoints, 0).unwrap(), + CompleteTree::new(max_checkpoints), + ShardTree::empty(MemoryShardStore::empty(), max_checkpoints), ) }