From f4d43b5d48011a46b83522bdc6f9426c2f33a9a8 Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@nutty.land>
Date: Thu, 18 Jan 2024 17:28:06 -0700
Subject: [PATCH] shardtree: Add the ability to avoid pruning specific
 checkpoints.

---
 shardtree/src/lib.rs | 56 ++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 52 insertions(+), 4 deletions(-)

diff --git a/shardtree/src/lib.rs b/shardtree/src/lib.rs
index f2f78893..02d5faee 100644
--- a/shardtree/src/lib.rs
+++ b/shardtree/src/lib.rs
@@ -68,6 +68,8 @@ pub struct ShardTree<S: ShardStore, const DEPTH: u8, const SHARD_HEIGHT: u8> {
     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<S::CheckpointId>
 }
 
 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<Address, BTreeMap<Position, RetentionFlags>> =
                 BTreeMap::new();
@@ -454,8 +464,9 @@ 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 +1366,43 @@ 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