diff --git a/shardtree/src/lib.rs b/shardtree/src/lib.rs index f2f78893..24937904 100644 --- a/shardtree/src/lib.rs +++ b/shardtree/src/lib.rs @@ -68,6 +68,8 @@ pub struct ShardTree { store: S, /// The maximum number of checkpoints to retain before pruning. max_checkpoints: usize, + /// The set of checkpoints to be explicitly retained in pruning operations. + to_retain: BTreeSet, } impl< @@ -83,6 +85,7 @@ impl< Self { store, max_checkpoints, + to_retain: BTreeSet::new(), } } @@ -111,6 +114,12 @@ impl< (0x1 << (DEPTH - SHARD_HEIGHT)) - 1 } + /// Adds the provided checkpoint to the set of checkpoints to be retained + /// across pruning operations. + pub fn ensure_retained(&mut self, checkpoint_id: C) { + self.to_retain.insert(checkpoint_id); + } + /// Returns the leaf value at the specified position, if it is a marked leaf. pub fn get_marked_leaf( &self, @@ -442,10 +451,11 @@ impl< checkpoint_count, self.max_checkpoints, ); - if checkpoint_count > self.max_checkpoints { + let retain_count = self.max_checkpoints + self.to_retain.len(); + if checkpoint_count > retain_count { // Batch removals by subtree & create a list of the checkpoint identifiers that // will be removed from the checkpoints map. - let remove_count = checkpoint_count - self.max_checkpoints; + let remove_count = checkpoint_count - retain_count; let mut checkpoints_to_delete = vec![]; let mut clear_positions: BTreeMap> = BTreeMap::new(); @@ -454,8 +464,10 @@ impl< // When removing is true, we are iterating through the range of // checkpoints being removed. When remove is false, we are // iterating through the range of checkpoints that are being - // retained. - let removing = checkpoints_to_delete.len() < remove_count; + // retained, or skipping over a particular checkpoint that we + // have been explicitly asked to retain. + let removing = + checkpoints_to_delete.len() < remove_count && !self.to_retain.contains(cid); if removing { checkpoints_to_delete.push(cid.clone()); @@ -1355,6 +1367,42 @@ mod tests { ), Ok(()), ); + + // Append a leaf we want to retain + assert_eq!(tree.append('e'.to_string(), Retention::Marked), Ok(()),); + + // Now a few more leaves and then checkpoint + for c in 'f'..='i' { + tree.append(c.to_string(), Retention::Ephemeral).unwrap(); + } + + // Checkpoint the tree. We'll want to retain this checkpoint. + assert_eq!(tree.checkpoint(12), Ok(true)); + tree.ensure_retained(12); + + // Simulate adding yet another block + for c in 'j'..='m' { + tree.append(c.to_string(), Retention::Ephemeral).unwrap(); + } + + assert_eq!(tree.checkpoint(13), Ok(true)); + + // Witness `e` as of checkpoint 12 + let e_witness_12 = tree + .witness_at_checkpoint_id(Position::from(4), &12) + .unwrap(); + + // Now add some more checkpoints, which would ordinarily cause checkpoint 12 + // to be pruned (but will not, because we explicitly retained it.) + for i in 14..24 { + assert_eq!(tree.checkpoint(i), Ok(true)); + } + + // Verify that we can still compute the same root + assert_matches!( + tree.witness_at_checkpoint_id(Position::from(4), &12), + Ok(w) if w == e_witness_12 + ); } // Combined tree tests