shardtree: Add `root_at_checkpoint_id` methods.

It can be useful to refer to a specific checkpoint, rather than just a
checkpoint depth, for which one wishes to compute the root.

shardtree/CHANGELOG.md

@@ -7,6 +7,14 @@ and this project adheres to Rust's notion of
 
 ## Unreleased
 
+## Added
+* `Shardtree::{root_at_checkpoint_id, root_at_checkpoint_id_caching}`
+
+## Changed
+* `Shardtree::root_at_checkpoint` and `Shardtree::root_at_checkpoint_caching` have
+  been renamed to `root_at_checkpoint_depth` and `root_at_checkpoint_depth_caching`,
+  respectively.
+
 ## [0.1.0] - 2023-09-08
 
 Initial release!

shardtree/src/batch.rs

@@ -533,8 +533,8 @@ mod tests {
             assert_eq!(lhs.max_leaf_position(0), Ok(Some(Position::from(i as u64))));
             assert_eq!(rhs.max_leaf_position(0), Ok(Some(Position::from(i as u64))));
 
-            assert_eq!(lhs.root_at_checkpoint(0).unwrap(), expected_root);
-            assert_eq!(rhs.root_at_checkpoint(0).unwrap(), expected_root);
+            assert_eq!(lhs.root_at_checkpoint_depth(0).unwrap(), expected_root);
+            assert_eq!(rhs.root_at_checkpoint_depth(0).unwrap(), expected_root);
         }
     }
 }
@@ -554,7 +554,7 @@ mod proptests {
             let (left, right) = build_insert_tree_and_batch_insert(leaves);
 
             // Check that the resulting trees are equal.
-            assert_eq!(left.root_at_checkpoint(0), right.root_at_checkpoint(0));
+            assert_eq!(left.root_at_checkpoint_depth(0), right.root_at_checkpoint_depth(0));
         }
     }
 }

shardtree/src/lib.rs

@@ -641,7 +641,7 @@ impl<
     /// given position, and parents of such nodes, will be replaced by the empty root for the
     /// associated level.
     ///
-    /// Use [`Self::root_at_checkpoint`] to obtain the root of the overall tree.
+    /// Use [`Self::root_at_checkpoint_id`] to obtain the root of the overall tree.
     pub fn root(
         &self,
         address: Address,
@@ -970,12 +970,47 @@ impl<
         })
     }
 
+    /// Computes the root of the tree as of the checkpointed position having the specified
+    /// checkpoint id.
+    pub fn root_at_checkpoint_id(&self, checkpoint: &C) -> Result<H, ShardTreeError<S::Error>> {
+        let position = self
+            .store
+            .get_checkpoint(checkpoint)
+            .map_err(ShardTreeError::Storage)?
+            .map(|c| c.position())
+            .ok_or(QueryError::CheckpointPruned)?;
+
+        position.map_or_else(
+            || Ok(H::empty_root(Self::root_addr().level())),
+            |pos| self.root(Self::root_addr(), pos + 1),
+        )
+    }
+
+    /// Computes the root of the tree as of the checkpointed position having the specified
+    /// checkpoint id, caching intermediate values produced while computing the root.
+    pub fn root_at_checkpoint_id_caching(
+        &mut self,
+        checkpoint: &C,
+    ) -> Result<H, ShardTreeError<S::Error>> {
+        let position = self
+            .store
+            .get_checkpoint(checkpoint)
+            .map_err(ShardTreeError::Storage)?
+            .map(|c| c.position())
+            .ok_or(QueryError::CheckpointPruned)?;
+
+        position.map_or_else(
+            || Ok(H::empty_root(Self::root_addr().level())),
+            |pos| self.root_caching(Self::root_addr(), pos + 1),
+        )
+    }
+
     /// Computes the root of the tree as of the checkpointed position at the specified depth.
     ///
     /// Returns the root as of the most recently appended leaf if `checkpoint_depth == 0`. Note
     /// that if the most recently appended leaf is also a checkpoint, this will return the same
     /// result as `checkpoint_depth == 1`.
-    pub fn root_at_checkpoint(
+    pub fn root_at_checkpoint_depth(
         &self,
         checkpoint_depth: usize,
     ) -> Result<H, ShardTreeError<S::Error>> {
@@ -985,7 +1020,9 @@ impl<
         )
     }
 
-    pub fn root_at_checkpoint_caching(
+    /// Computes the root of the tree as of the checkpointed position at the specified depth,
+    /// caching intermediate values produced while computing the root.
+    pub fn root_at_checkpoint_depth_caching(
         &mut self,
         checkpoint_depth: usize,
     ) -> Result<H, ShardTreeError<S::Error>> {

shardtree/src/testing.rs

@@ -185,7 +185,7 @@ impl<
     }
 
     fn root(&self, checkpoint_depth: usize) -> Option<H> {
-        match ShardTree::root_at_checkpoint(self, checkpoint_depth) {
+        match ShardTree::root_at_checkpoint_depth(self, checkpoint_depth) {
             Ok(v) => Some(v),
             Err(err) => panic!("root computation failed: {:?}", err),
         }
@@ -254,7 +254,7 @@ pub fn check_shardtree_insertion<
     );
 
     assert_matches!(
-        tree.root_at_checkpoint(1),
+        tree.root_at_checkpoint_depth(1),
         Err(ShardTreeError::Query(QueryError::TreeIncomplete(v))) if v == vec![Address::from_parts(Level::from(0), 0)]
     );
 
@@ -271,12 +271,12 @@ pub fn check_shardtree_insertion<
     );
 
     assert_matches!(
-        tree.root_at_checkpoint(0),
+        tree.root_at_checkpoint_depth(0),
         Ok(h) if h == *"abcd____________"
     );
 
     assert_matches!(
-        tree.root_at_checkpoint(1),
+        tree.root_at_checkpoint_depth(1),
         Ok(h) if h == *"ab______________"
     );
 
@@ -308,7 +308,7 @@ pub fn check_shardtree_insertion<
     );
 
     assert_matches!(
-        tree.root_at_checkpoint(0),
+        tree.root_at_checkpoint_depth(0),
         // The (0, 13) and (1, 7) incomplete subtrees are
         // not considered incomplete here because they appear
         // at the tip of the tree.
@@ -331,12 +331,12 @@ pub fn check_shardtree_insertion<
     );
 
     assert_matches!(
-        tree.root_at_checkpoint(0),
+        tree.root_at_checkpoint_depth(0),
         Ok(h) if h == *"abcdefghijkl____"
     );
 
     assert_matches!(
-        tree.root_at_checkpoint(1),
+        tree.root_at_checkpoint_depth(1),
         Ok(h) if h == *"ab______________"
     );
 }