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

Merge pull request #72 from nuttycom/no_initial_checkpoint 

Remove initial checkpoints.
This commit is contained in:
Kris Nuttycombe 2023-05-24 12:57:02 -06:00 committed by GitHub
parent a26844a451 2b8e2d62fa
commit beb6ead57f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 140 additions and 194 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

63
bridgetree/src/lib.rs
View File

@ -470,13 +470,13 @@ impl<H, C, const DEPTH: u8> BridgeTree<H, C, DEPTH> {
/// ///
/// Panics if `max_checkpoints < 1` because mark/rewind logic depends upon the presence /// Panics if `max_checkpoints < 1` because mark/rewind logic depends upon the presence
/// of checkpoints to function. /// 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); assert!(max_checkpoints >= 1);
Self { Self {
prior_bridges: vec![], prior_bridges: vec![],
current_bridge: None, current_bridge: None,
saved: BTreeMap::new(), saved: BTreeMap::new(),
checkpoints: VecDeque::from(vec![Checkpoint::at_length(0, initial_checkpoint_id)]), checkpoints: VecDeque::new(),
max_checkpoints, max_checkpoints,
} }
} }
@ -535,12 +535,8 @@ impl<H, C, const DEPTH: u8> BridgeTree<H, C, DEPTH> {
impl<H: Hashable + Clone + Ord, C: Clone + Ord, const DEPTH: u8> BridgeTree<H, C, DEPTH> { impl<H: Hashable + Clone + Ord, C: Clone + Ord, const DEPTH: u8> BridgeTree<H, C, DEPTH> {
/// Construct a new BridgeTree that will start recording changes from the state of /// Construct a new BridgeTree that will start recording changes from the state of
/// the specified frontier. /// the specified frontier.
pub fn from_frontier( pub fn from_frontier(max_checkpoints: usize, frontier: NonEmptyFrontier<H>) -> Self {
max_checkpoints: usize, Self {
frontier: NonEmptyFrontier<H>,
checkpoint_id: C,
) -> Self {
let mut bridge = Self {
prior_bridges: vec![], prior_bridges: vec![],
current_bridge: Some(MerkleBridge::from_parts( current_bridge: Some(MerkleBridge::from_parts(
None, None,
@ -551,9 +547,7 @@ impl<H: Hashable + Clone + Ord, C: Clone + Ord, const DEPTH: u8> BridgeTree<H, C
saved: BTreeMap::new(), saved: BTreeMap::new(),
checkpoints: VecDeque::new(), checkpoints: VecDeque::new(),
max_checkpoints, max_checkpoints,
}; }
bridge.checkpoint(checkpoint_id);
bridge
} }
/// Construct a new BridgeTree from its constituent parts, checking for internal /// Construct a new BridgeTree from its constituent parts, checking for internal
@ -717,11 +711,9 @@ impl<H: Hashable + Clone + Ord, C: Clone + Ord, const DEPTH: u8> BridgeTree<H, C
// mark the position as having been marked in the current checkpoint // mark the position as having been marked in the current checkpoint
if let std::collections::btree_map::Entry::Vacant(e) = self.saved.entry(pos) { if let std::collections::btree_map::Entry::Vacant(e) = self.saved.entry(pos) {
let c = self if let Some(c) = self.checkpoints.back_mut() {
.checkpoints c.marked.insert(pos);
.back_mut() }
.expect("Checkpoints deque must never be empty");
c.marked.insert(pos);
e.insert(self.prior_bridges.len() - 1); e.insert(self.prior_bridges.len() - 1);
} }
@ -749,11 +741,11 @@ impl<H: Hashable + Clone + Ord, C: Clone + Ord, const DEPTH: u8> BridgeTree<H, C
/// false if we were already not maintaining a mark at this position. /// false if we were already not maintaining a mark at this position.
pub fn remove_mark(&mut self, position: Position) -> bool { pub fn remove_mark(&mut self, position: Position) -> bool {
if self.saved.contains_key(&position) { if self.saved.contains_key(&position) {
let c = self if let Some(c) = self.checkpoints.back_mut() {
.checkpoints c.forgotten.insert(position);
.back_mut() } else {
.expect("Checkpoints deque must never be empty."); self.saved.remove(&position);
c.forgotten.insert(position); }
true true
} else { } else {
false false
@ -805,12 +797,7 @@ impl<H: Hashable + Clone + Ord, C: Clone + Ord, const DEPTH: u8> BridgeTree<H, C
/// at that tree state have been removed using `rewind`. This function /// at that tree state have been removed using `rewind`. This function
/// return false and leave the tree unmodified if no checkpoints exist. /// return false and leave the tree unmodified if no checkpoints exist.
pub fn rewind(&mut self) -> bool { pub fn rewind(&mut self) -> bool {
if self.checkpoints.len() > 1 { if let Some(c) = self.checkpoints.pop_back() {
let c = self
.checkpoints
.pop_back()
.expect("Checkpoints deque is known to be non-empty.");
// Remove marks for positions that were marked during the lifetime of this checkpoint. // Remove marks for positions that were marked during the lifetime of this checkpoint.
for pos in c.marked { for pos in c.marked {
self.saved.remove(&pos); self.saved.remove(&pos);
@ -1063,7 +1050,7 @@ mod tests {
#[test] #[test]
fn tree_depth() { fn tree_depth() {
let mut tree = BridgeTree::<String, usize, 3>::new(100, 0); let mut tree = BridgeTree::<String, usize, 3>::new(100);
for c in 'a'..'i' { for c in 'a'..'i' {
assert!(tree.append(c.to_string())) 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()); 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| { proptest::collection::vec(arb_operation(item_gen, pos_gen), 0..max_count).prop_map(|ops| {
let mut tree: BridgeTree<G::Value, usize, 8> = BridgeTree::new(10, 0); let mut tree: BridgeTree<G::Value, usize, 8> = BridgeTree::new(10);
for (i, op) in ops.into_iter().enumerate() { for (i, op) in ops.into_iter().enumerate() {
apply_operation(&mut tree, op.map_checkpoint_id(|_| i)); apply_operation(&mut tree, op.map_checkpoint_id(|_| i));
} }
@ -1130,33 +1117,31 @@ mod tests {
#[test] #[test]
fn root_hashes() { fn root_hashes() {
check_root_hashes(|max_checkpoints| { check_root_hashes(BridgeTree::<String, usize, 4>::new);
BridgeTree::<String, usize, 4>::new(max_checkpoints, 0)
});
} }
#[test] #[test]
fn witness() { fn witness() {
check_witnesses(|max_checkpoints| BridgeTree::<String, usize, 4>::new(max_checkpoints, 0)); check_witnesses(BridgeTree::<String, usize, 4>::new);
} }
#[test] #[test]
fn checkpoint_rewind() { fn checkpoint_rewind() {
check_checkpoint_rewind(|max_checkpoints| { check_checkpoint_rewind(|max_checkpoints| {
BridgeTree::<String, usize, 4>::new(max_checkpoints, 0) BridgeTree::<String, usize, 4>::new(max_checkpoints)
}); });
} }
#[test] #[test]
fn rewind_remove_mark() { fn rewind_remove_mark() {
check_rewind_remove_mark(|max_checkpoints| { check_rewind_remove_mark(|max_checkpoints| {
BridgeTree::<String, usize, 4>::new(max_checkpoints, 0) BridgeTree::<String, usize, 4>::new(max_checkpoints)
}); });
} }
#[test] #[test]
fn garbage_collect() { fn garbage_collect() {
let mut tree: BridgeTree<String, usize, 7> = BridgeTree::new(1000, 0); let mut tree: BridgeTree<String, usize, 7> = BridgeTree::new(1000);
let empty_root = tree.root(0); let empty_root = tree.root(0);
tree.append("a".to_string()); tree.append("a".to_string());
for i in 0..100 { for i in 0..100 {
@ -1167,7 +1152,7 @@ mod tests {
tree.rewind(); tree.rewind();
assert!(tree.root(0) != empty_root); assert!(tree.root(0) != empty_root);
let mut t = BridgeTree::<String, usize, 7>::new(10, 0); let mut t = BridgeTree::<String, usize, 7>::new(10);
let mut to_unmark = vec![]; let mut to_unmark = vec![];
let mut has_witness = vec![]; let mut has_witness = vec![];
for i in 0u64..100 { for i in 0u64..100 {
@ -1213,8 +1198,8 @@ mod tests {
max_checkpoints: usize, max_checkpoints: usize,
) -> CombinedTree<H, usize, CompleteTree<H, usize, 4>, BridgeTree<H, usize, 4>> { ) -> CombinedTree<H, usize, CompleteTree<H, usize, 4>, BridgeTree<H, usize, 4>> {
CombinedTree::new( CombinedTree::new(
CompleteTree::<H, usize, 4>::new(max_checkpoints, 0), CompleteTree::<H, usize, 4>::new(max_checkpoints),
BridgeTree::<H, usize, 4>::new(max_checkpoints, 0), BridgeTree::<H, usize, 4>::new(max_checkpoints),
) )
} }

66
incrementalmerkletree/src/testing/complete_tree.rs
View File

@ -71,11 +71,11 @@ pub struct CompleteTree<H, C: Ord, const DEPTH: u8> {
impl<H: Hashable, C: Clone + Ord + core::fmt::Debug, const DEPTH: u8> CompleteTree<H, C, DEPTH> { impl<H: Hashable, C: Clone + Ord + core::fmt::Debug, const DEPTH: u8> CompleteTree<H, C, DEPTH> {
/// Creates a new, empty binary tree /// 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 { Self {
leaves: vec![], leaves: vec![],
marks: BTreeSet::new(), marks: BTreeSet::new(),
checkpoints: BTreeMap::from([(initial_checkpoint_id, Checkpoint::at_length(0))]), checkpoints: BTreeMap::new(),
max_checkpoints, max_checkpoints,
} }
} }
@ -140,22 +140,18 @@ impl<H: Hashable, C: Clone + Ord + core::fmt::Debug, const DEPTH: u8> CompleteTr
/// Marks the current tree state leaf as a value that we're interested in /// 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. /// marking. Returns the current position if the tree is non-empty.
fn mark(&mut self) -> Option<Position> { fn mark(&mut self) -> Option<Position> {
match self.current_position() { if let Some(pos) = self.current_position() {
Some(pos) => { if !self.marks.contains(&pos) {
if !self.marks.contains(&pos) { self.marks.insert(pos);
self.marks.insert(pos);
self.checkpoints if let Some(checkpoint) = self.checkpoints.values_mut().rev().next() {
.iter_mut() checkpoint.marked.insert(pos);
.rev()
.next()
.unwrap()
.1
.marked
.insert(pos);
} }
Some(pos)
} }
None => None,
Some(pos)
} else {
None
} }
} }
@ -282,14 +278,11 @@ impl<H: Hashable + PartialEq + Clone, C: Ord + Clone + core::fmt::Debug, const D
fn remove_mark(&mut self, position: Position) -> bool { fn remove_mark(&mut self, position: Position) -> bool {
if self.marks.contains(&position) { if self.marks.contains(&position) {
self.checkpoints if let Some(c) = self.checkpoints.values_mut().rev().next() {
.iter_mut() c.forgotten.insert(position);
.rev() } else {
.next() self.marks.remove(&position);
.unwrap() }
.1
.forgotten
.insert(position);
true true
} else { } else {
false false
@ -297,7 +290,7 @@ impl<H: Hashable + PartialEq + Clone, C: Ord + Clone + core::fmt::Debug, const D
} }
fn checkpoint(&mut self, id: C) -> bool { fn checkpoint(&mut self, id: C) -> 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()); Self::checkpoint(self, id, self.current_position());
true true
} else { } else {
@ -306,8 +299,7 @@ impl<H: Hashable + PartialEq + Clone, C: Ord + Clone + core::fmt::Debug, const D
} }
fn rewind(&mut self) -> bool { fn rewind(&mut self) -> bool {
if self.checkpoints.len() > 1 { if let Some((id, c)) = self.checkpoints.iter().rev().next() {
let (id, c) = self.checkpoints.iter().rev().next().unwrap();
self.leaves.truncate(c.leaves_len); self.leaves.truncate(c.leaves_len);
for pos in c.marked.iter() { for pos in c.marked.iter() {
self.marks.remove(pos); self.marks.remove(pos);
@ -342,7 +334,7 @@ mod tests {
expected = SipHashable::combine(lvl.into(), &expected, &expected); expected = SipHashable::combine(lvl.into(), &expected, &expected);
} }
let tree = CompleteTree::<SipHashable, (), DEPTH>::new(100, ()); let tree = CompleteTree::<SipHashable, (), DEPTH>::new(100);
assert_eq!(tree.root(0).unwrap(), expected); assert_eq!(tree.root(0).unwrap(), expected);
} }
@ -351,7 +343,7 @@ mod tests {
const DEPTH: u8 = 3; const DEPTH: u8 = 3;
let values = (0..(1 << DEPTH)).map(SipHashable); let values = (0..(1 << DEPTH)).map(SipHashable);
let mut tree = CompleteTree::<SipHashable, (), DEPTH>::new(100, ()); let mut tree = CompleteTree::<SipHashable, (), DEPTH>::new(100);
for value in values { for value in values {
assert!(tree.append(value, Retention::Ephemeral).is_ok()); assert!(tree.append(value, Retention::Ephemeral).is_ok());
} }
@ -376,21 +368,17 @@ mod tests {
#[test] #[test]
fn append() { fn append() {
check_append(|max_checkpoints| CompleteTree::<String, usize, 4>::new(max_checkpoints, 0)); check_append(CompleteTree::<String, usize, 4>::new);
} }
#[test] #[test]
fn root_hashes() { fn root_hashes() {
check_root_hashes(|max_checkpoints| { check_root_hashes(CompleteTree::<String, usize, 4>::new);
CompleteTree::<String, usize, 4>::new(max_checkpoints, 0)
});
} }
#[test] #[test]
fn witnesses() { fn witnesses() {
check_witnesses(|max_checkpoints| { check_witnesses(CompleteTree::<String, usize, 4>::new);
CompleteTree::<String, usize, 4>::new(max_checkpoints, 0)
});
} }
#[test] #[test]
@ -400,7 +388,7 @@ mod tests {
const DEPTH: u8 = 3; const DEPTH: u8 = 3;
let values = (0..(1 << DEPTH)).map(SipHashable); let values = (0..(1 << DEPTH)).map(SipHashable);
let mut tree = CompleteTree::<SipHashable, (), DEPTH>::new(100, ()); let mut tree = CompleteTree::<SipHashable, (), DEPTH>::new(100);
for value in values { for value in values {
assert!(Tree::append(&mut tree, value, Retention::Marked)); assert!(Tree::append(&mut tree, value, Retention::Marked));
} }
@ -435,14 +423,14 @@ mod tests {
#[test] #[test]
fn checkpoint_rewind() { fn checkpoint_rewind() {
check_checkpoint_rewind(|max_checkpoints| { check_checkpoint_rewind(|max_checkpoints| {
CompleteTree::<String, usize, 4>::new(max_checkpoints, 0) CompleteTree::<String, usize, 4>::new(max_checkpoints)
}); });
} }
#[test] #[test]
fn rewind_remove_mark() { fn rewind_remove_mark() {
check_rewind_remove_mark(|max_checkpoints| { check_rewind_remove_mark(|max_checkpoints| {
CompleteTree::<String, usize, 4>::new(max_checkpoints, 0) CompleteTree::<String, usize, 4>::new(max_checkpoints)
}); });
} }
} }

4
shardtree/proptest-regressions/lib.txt
View File

@ -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 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 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 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)]

201
shardtree/src/lib.rs
View File

@ -1761,21 +1761,12 @@ impl<
where where
S::Error: From<InsertionError> + From<QueryError>, S::Error: From<InsertionError> + From<QueryError>,
{ {
/// Creates a new empty tree and establishes a checkpoint for the empty tree at the given /// Creates a new empty tree.
/// checkpoint identifier. pub fn empty(store: S, max_checkpoints: usize) -> Self {
pub fn empty( Self {
store: S,
max_checkpoints: usize,
initial_checkpoint_id: C,
) -> Result<Self, S::Error> {
let mut result = Self {
store, store,
max_checkpoints, max_checkpoints,
}; }
result
.store
.add_checkpoint(initial_checkpoint_id, Checkpoint::tree_empty())?;
Ok(result)
} }
/// Constructs a wrapper around the provided shard store without initialization. /// Constructs a wrapper around the provided shard store without initialization.
@ -1797,11 +1788,16 @@ where
Level::from(SHARD_HEIGHT - 1) 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. /// Returns the leaf value at the specified position, if it is a marked leaf.
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 }))
} }
@ -1923,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 {
@ -2031,19 +2027,13 @@ where
// Update the rightmost subtree to add the `CHECKPOINT` flag to the right-most leaf (which // 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). // 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(subtree) = self.store.last_shard()? {
if let Some((replacement, checkpoint_position)) = go(subtree.root_addr, &subtree.root) { if let Some((replacement, pos)) = go(subtree.root_addr, &subtree.root) {
if self self.store.put_shard(LocatedTree {
.store root_addr: subtree.root_addr,
.put_shard(LocatedTree { root: replacement,
root_addr: subtree.root_addr, })?;
root: replacement,
})
.is_err()
{
return Ok(false);
}
self.store 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 // early return once we've updated the tree state
self.prune_excess_checkpoints()?; self.prune_excess_checkpoints()?;
@ -2075,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| {
@ -2132,36 +2122,32 @@ where
) -> Result<bool, S::Error> { ) -> Result<bool, S::Error> {
Ok(if checkpoint_depth == 0 { Ok(if checkpoint_depth == 0 {
true true
} else if self.store.checkpoint_count()? > 1 { } else if let Some((checkpoint_id, c)) =
if let Some((checkpoint_id, c)) = self.store.get_checkpoint_at_depth(checkpoint_depth)?
self.store.get_checkpoint_at_depth(checkpoint_depth)? {
{ match c.tree_state {
match c.tree_state { TreeState::Empty => {
TreeState::Empty => { self.store
self.store .truncate(Address::from_parts(Self::subtree_level(), 0))?;
.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)?; self.store.truncate_checkpoints(&checkpoint_id)?;
true true
} } else {
TreeState::AtPosition(position) => { false
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 { } else {
false false
@ -2332,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(
@ -2355,37 +2341,44 @@ 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: &C, as_of_checkpoint: Option<&C>,
) -> Result<bool, S::Error> { ) -> Result<bool, S::Error> {
#[allow(clippy::blocks_in_if_conditions)] match self.store.get_shard(Self::subtree_addr(position))? {
if self.get_marked_leaf(position)?.is_some() { Some(shard)
if self if shard
.store .value_at_position(position)
.update_checkpoint_with(as_of_checkpoint, |checkpoint| { .iter()
checkpoint.marks_removed.insert(position); .any(|(_, r)| r.is_marked()) =>
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| {
if let Some(cid) = self.store.min_checkpoint_id()? { checkpoint.marks_removed.insert(position);
if self.store.update_checkpoint_with(&cid, |checkpoint| { Ok(())
checkpoint.marks_removed.insert(position); })
Ok(()) }
})? { _ => {
return Ok(true); // 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] #[test]
fn shardtree_insertion() { fn shardtree_insertion() {
let mut tree: ShardTree<MemoryShardStore<String, usize>, 4, 3> = let mut tree: ShardTree<MemoryShardStore<String, usize>, 4, 3> =
ShardTree::empty(MemoryShardStore::empty(), 100, 0).unwrap(); ShardTree::empty(MemoryShardStore::empty(), 100);
assert_matches!( assert_matches!(
tree.batch_insert( tree.batch_insert(
Position::from(1), Position::from(1),
@ -2915,11 +2908,12 @@ mod tests {
} }
fn remove_mark(&mut self, position: Position) -> bool { fn remove_mark(&mut self, position: Position) -> bool {
if let Ok(Some(c)) = self.store.max_checkpoint_id() { ShardTree::remove_mark(
ShardTree::remove_mark(self, position, &c).unwrap() self,
} else { position,
false self.store.max_checkpoint_id().unwrap().as_ref(),
} )
.unwrap()
} }
fn checkpoint(&mut self, checkpoint_id: C) -> bool { fn checkpoint(&mut self, checkpoint_id: C) -> bool {
@ -2934,60 +2928,35 @@ mod tests {
#[test] #[test]
fn append() { fn append() {
check_append(|m| { check_append(|m| {
ShardTree::<MemoryShardStore<String, usize>, 4, 3>::empty( ShardTree::<MemoryShardStore<String, usize>, 4, 3>::empty(MemoryShardStore::empty(), m)
MemoryShardStore::empty(),
m,
0,
)
.unwrap()
}); });
} }
#[test] #[test]
fn root_hashes() { fn root_hashes() {
check_root_hashes(|m| { check_root_hashes(|m| {
ShardTree::<MemoryShardStore<String, usize>, 4, 3>::empty( ShardTree::<MemoryShardStore<String, usize>, 4, 3>::empty(MemoryShardStore::empty(), m)
MemoryShardStore::empty(),
m,
0,
)
.unwrap()
}); });
} }
#[test] #[test]
fn witnesses() { fn witnesses() {
check_witnesses(|m| { check_witnesses(|m| {
ShardTree::<MemoryShardStore<String, usize>, 4, 3>::empty( ShardTree::<MemoryShardStore<String, usize>, 4, 3>::empty(MemoryShardStore::empty(), m)
MemoryShardStore::empty(),
m,
0,
)
.unwrap()
}); });
} }
#[test] #[test]
fn checkpoint_rewind() { fn checkpoint_rewind() {
check_checkpoint_rewind(|m| { check_checkpoint_rewind(|m| {
ShardTree::<MemoryShardStore<String, usize>, 4, 3>::empty( ShardTree::<MemoryShardStore<String, usize>, 4, 3>::empty(MemoryShardStore::empty(), m)
MemoryShardStore::empty(),
m,
0,
)
.unwrap()
}); });
} }
#[test] #[test]
fn rewind_remove_mark() { fn rewind_remove_mark() {
check_rewind_remove_mark(|m| { check_rewind_remove_mark(|m| {
ShardTree::<MemoryShardStore<String, usize>, 4, 3>::empty( ShardTree::<MemoryShardStore<String, usize>, 4, 3>::empty(MemoryShardStore::empty(), m)
MemoryShardStore::empty(),
m,
0,
)
.unwrap()
}); });
} }
@ -3002,8 +2971,8 @@ mod tests {
ShardTree<MemoryShardStore<H, usize>, 4, 3>, ShardTree<MemoryShardStore<H, usize>, 4, 3>,
> { > {
CombinedTree::new( CombinedTree::new(
CompleteTree::new(max_checkpoints, 0), CompleteTree::new(max_checkpoints),
ShardTree::empty(MemoryShardStore::empty(), max_checkpoints, 0).unwrap(), ShardTree::empty(MemoryShardStore::empty(), max_checkpoints),
) )
} }