shardtree: Rework rewind & checkpoint depth handling.

shardtree/CHANGELOG.md

@@ -7,11 +7,90 @@ and this project adheres to Rust's notion of
 
 ## Unreleased
 
-### Changed
+This release includes a significant refactoring and rework of several methods
-- MSRV is now 1.64
+of the `shardtree::ShardTree` type and the `shardtree::store::ShardStore`
+trait. Please read the notes for this release carefully as the semantics of
+important methods have changed. These changes may require changes to clients of
+this crate; in particular, the existence of a checkpoint is required for all
+witnessing and rewinding operations.
 
 ### Added
 - `shardtree::store::ShardStore::for_each_checkpoint`
+- `shardtree::ShardTree::truncate_to_checkpoint_depth`. This replaces
+  the `ShardTree::truncate_to_depth` method, with modified semantics such that
+  the provided `checkpoint_depth` argument is now treated strictly as a
+  zero-based index into the checkpoints applied to the tree, in reverse order
+  of checkpoint identifier. It is no longer possible to truncate the tree if no
+  checkpoints have been created.
+- `shardtree::ShardTree::truncate_to_checkpoint` replaces
+  `ShardTree::truncate_removing_checkpoint`. Instead of removing
+  the checkpoint, this replacement method removes all state from the tree
+  related to leaves having positions greater than the checkpointed position,
+  but unlike `truncate_removing_checkpoint` it leaves the checkpoint itself
+  in place.
+- `shardtree::store::ShardStore::truncate_shards` replaces
+  `ShardStore::truncate`. Instead of taking an `Address` and requiring that
+  implementations impose additional runtime restriction on the level of that
+  address, the replacement method directly takes a shard index.
+- `ShardStore::truncate_truncate_checkpoints_retaining` replaces
+  `ShardStore::truncate_checkpoints`. Instead of removing the checkpoint
+  having the specified identifier, the associated checkpoint should be retained
+  but any additional metadata stored with the checkpoint, such as information
+  about marks removed after the creation of the checkpoint, should be deleted.
+
+### Changed
+- MSRV is now 1.64
+- `shardtree::ShardTree`:
+  - `ShardTree::max_leaf_position` now takes its `checkpoint_depth` argument
+    as `Option<usize>` instead of `usize`. The semantics of this method are now
+    changed such that if a checkpoint depth is provided, it is now treated
+    strictly as a zero-based index into the checkpoints of the tree in reverse
+    order of checkpoint identifier, and an error is returned if no checkpoint
+    exists at the given index. A `None` value passed for this argument causes
+    it to return the maximum position among all leaves added to the tree.
+  - `ShardTree::root_at_checkpoint_id` and `root_at_checkpoint_id_caching` now
+    each return the computed root as an optional value, returning `Ok(None)` if
+    the checkpoint corresponding to the requested identifier does not exist.
+  - `ShardTree::root_at_checkpoint_depth` and `root_at_checkpoint_depth_caching`
+    now takes the `checkpoint_depth` argument as `Option<usize>` instead of
+    `usize`. The semantics of this method are now changed such that if a
+    checkpoint depth is provided, it is now treated strictly as a zero-based
+    index into the checkpoints of the tree in reverse order of checkpoint
+    identifier, and an error is returned if no checkpoint exists at the given
+    index. A `None` value passed for this argument causes it to return the root
+    computed over all of the leaves in the tree. These methods now each return
+    the computed root as an optional value, returning `Ok(None)` if no checkpoint
+    exist at the requested checkpoint depth.
+  - `ShardTree::witness_at_checkpoint_id` and `witness_at_checkpoint_id_caching`
+    now each return the computed witness as an optional value, returning
+    `Ok(None)` if the checkpoint corresponding to the requested identifier does
+    not exist.
+  - `ShardTree::witness_at_checkpoint_depth` and `witness_at_checkpoint_depth_caching`
+    now each return the computed witness as an optional value, returning
+    `Ok(None)` if no checkpoint was available at the given checkpoint depth.The
+    semantics of this method are now changed such that if a checkpoint depth is
+    provided, it is now treated strictly as a zero-based index into the
+    checkpoints of the tree in reverse order of checkpoint identifier, and an
+    error is returned if no checkpoint exists at the given index. IT IS NO LONGER
+    POSSIBLE TO COMPUTE A WITNESS WITHOUT A CHECKPOINT BEING PRESENT IN THE TREE.
+- `shardtree::store::ShardStore`:
+  - The semantics of `ShardStore::get_checkpoint_at_depth` HAVE CHANGED WITHOUT
+    CHANGES TO THE METHOD SIGNATURE. The `checkpoint_depth` argument to this
+    method is now treated strictly as a zero-based index into the checkpoints
+    of the tree in reverse order of checkpoint identifier. Previously, this
+    method always returned `None` for `checkpoint_depth == 0`, and
+    `checkpoint_depth` was otherwise treated as a 1-based index instead of a
+    0-based index.
+
+### Removed
+- `shardtree::ShardTree::truncate_to_depth` has been replaced by
+  `ShardTree::truncate_to_checkpoint_depth`
+- `shardtree::ShardTree::truncate_removing_checkpoint` has been replaced by
+  `ShardTree::truncate_to_checkpoint`
+- `shardtree::store::ShardStore::truncate` has been replaced by
+  `ShardStore::truncate_shards`
+- `shardtree::store::ShardStore::truncate_checkpoints` has been replaced by
+  `ShardStore::truncate_checkpoints_retaining`
 
 ## [0.4.0] - 2024-08-12
 

shardtree/Cargo.toml

@@ -38,7 +38,12 @@ proptest.workspace = true
 legacy-api = ["incrementalmerkletree/legacy-api"]
 # The test-depenencies feature can be enabled to expose types and functions
 # that are useful for testing `shardtree` functionality.
-test-dependencies = ["proptest", "assert_matches", "incrementalmerkletree/test-dependencies"]
+test-dependencies = [
+  "dep:proptest", 
+  "dep:assert_matches", 
+  "dep:incrementalmerkletree-testing", 
+  "incrementalmerkletree/test-dependencies"
+]
 
 [target.'cfg(unix)'.dev-dependencies]
 tempfile = ">=3, <3.7.0" # 3.7 has MSRV 1.63

shardtree/proptest-regressions/lib.txt

@@ -20,3 +20,6 @@ cc d53a73021238de143764ee1d48b944abb93bd4bc54f35d16e514261220d3eb78 # shrinks to
 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)]
+cc 7cc1bd3b98a0ddbc7409077c010220ed8611936693eb955748df9d86167b1488 # shrinks to ops = [Append(SipHashable(0), Ephemeral), Append(SipHashable(0), Marked), Checkpoint(()), Witness(Position(1), 1)]
+cc 7e6c5a3b74e14acaf101f9d411e3d7af878c335abacc6cdbbe92615426a6c277 # shrinks to ops = [Append("a", Ephemeral), Append("a", Ephemeral), Append("a", Ephemeral), CurrentPosition, CurrentPosition, Append("a", Ephemeral), Append("a", Ephemeral), Append("a", Ephemeral), Append("a", Marked), Checkpoint(()), Checkpoint(()), Checkpoint(()), Checkpoint(()), Checkpoint(()), Witness(Position(6), 5)]
+cc 04866a1bb2691ed38b853576489ad858a2d0e09e1579e120fd825e9f809d544b # shrinks to ops = [Append("a", Checkpoint { id: (), marking: Marked }), Append("a", Checkpoint { id: (), marking: Marked }), Checkpoint(()), Checkpoint(()), Append("a", Checkpoint { id: (), marking: Marked }), Append("a", Ephemeral), Append("a", Ephemeral), Append("a", Ephemeral), Append("a", Ephemeral), Append("a", Checkpoint { id: (), marking: Marked }), Witness(Position(0), 1)]

shardtree/src/batch.rs

@@ -36,6 +36,8 @@ impl<
     ///
     /// This method operates on a single thread. If you have parallelism available, consider using
     /// [`LocatedPrunableTree::from_iter`] and [`Self::insert_tree`] instead.
+    ///
+    /// [`Node::Nil`]: crate::tree::Node::Nil
     #[allow(clippy::type_complexity)]
     pub fn batch_insert<I: Iterator<Item = (H, Retention<C>)>>(
         &mut self,
@@ -96,6 +98,8 @@ pub struct BatchInsertionResult<H, C: Ord, I: Iterator<Item = (H, Retention<C>)>
     /// The vector of addresses of [`Node::Nil`] nodes that were inserted into the tree as part of
     /// the insertion operation, for nodes that are required in order to construct a witness for
     /// each inserted leaf with [`Retention::Marked`] retention.
+    ///
+    /// [`Node::Nil`]: crate::tree::Node::Nil
     pub incomplete: Vec<IncompleteAt>,
     /// The maximum position at which a leaf was inserted.
     pub max_insert_position: Option<Position>,
@@ -507,8 +511,8 @@ mod tests {
     fn batch_insert_matches_insert_tree() {
         {
             let (lhs, rhs) = build_insert_tree_and_batch_insert(vec![]);
-            assert_eq!(lhs.max_leaf_position(0), Ok(None));
+            assert_eq!(lhs.max_leaf_position(None), Ok(None));
-            assert_eq!(rhs.max_leaf_position(0), Ok(None));
+            assert_eq!(rhs.max_leaf_position(None), Ok(None));
         }
 
         for i in 0..64 {
@@ -524,11 +528,23 @@ mod tests {
                 });
 
             let (lhs, rhs) = build_insert_tree_and_batch_insert(leaves);
-            assert_eq!(lhs.max_leaf_position(0), Ok(Some(Position::from(i as u64))));
+            assert_eq!(
-            assert_eq!(rhs.max_leaf_position(0), Ok(Some(Position::from(i as u64))));
+                lhs.max_leaf_position(None),
+                Ok(Some(Position::from(i as u64)))
+            );
+            assert_eq!(
+                rhs.max_leaf_position(None),
+                Ok(Some(Position::from(i as u64)))
+            );
 
-            assert_eq!(lhs.root_at_checkpoint_depth(0).unwrap(), expected_root);
+            assert_eq!(
-            assert_eq!(rhs.root_at_checkpoint_depth(0).unwrap(), expected_root);
+                lhs.root_at_checkpoint_depth(None).unwrap().as_ref(),
+                Some(&expected_root)
+            );
+            assert_eq!(
+                rhs.root_at_checkpoint_depth(None).unwrap().as_ref(),
+                Some(&expected_root)
+            );
         }
     }
 }
@@ -548,7 +564,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_depth(0), right.root_at_checkpoint_depth(0));
+            assert_eq!(left.root_at_checkpoint_depth(None), right.root_at_checkpoint_depth(None));
         }
     }
 }

shardtree/src/lib.rs

@@ -596,23 +596,25 @@ impl<
         Ok(())
     }
 
-    /// Truncates the tree, discarding all information after the checkpoint at the specified depth.
+    /// Truncates the tree, discarding all information after the checkpoint at the specified
+    /// checkpoint depth.
     ///
-    /// This will also discard all checkpoints with depth less than or equal to the specified depth.
     /// Returns `true` if the truncation succeeds or has no effect, or `false` if no checkpoint
-    /// exists at the specified depth.
+    /// exists at the specified depth. Depth 0 refers to the most recent checkpoint in the tree;
-    pub fn truncate_to_depth(
+    ///
+    /// ## Parameters
+    /// - `checkpoint_depth`: A zero-based index over the checkpoints that have been added to the
+    ///   tree, in reverse checkpoint identifier order.
+    pub fn truncate_to_checkpoint_depth(
         &mut self,
         checkpoint_depth: usize,
     ) -> Result<bool, ShardTreeError<S::Error>> {
-        if checkpoint_depth == 0 {
+        if let Some((checkpoint_id, c)) = self
-            Ok(true)
-        } else if let Some((checkpoint_id, c)) = self
             .store
             .get_checkpoint_at_depth(checkpoint_depth)
             .map_err(ShardTreeError::Storage)?
         {
-            self.truncate_removing_checkpoint_internal(&checkpoint_id, &c)?;
+            self.truncate_to_checkpoint_internal(&checkpoint_id, &c)?;
             Ok(true)
         } else {
             Ok(false)
@@ -621,10 +623,9 @@ impl<
 
     /// Truncates the tree, discarding all information after the specified checkpoint.
     ///
-    /// This will also discard all checkpoints with depth less than or equal to the specified depth.
     /// Returns `true` if the truncation succeeds or has no effect, or `false` if no checkpoint
     /// exists for the specified checkpoint identifier.
-    pub fn truncate_removing_checkpoint(
+    pub fn truncate_to_checkpoint(
         &mut self,
         checkpoint_id: &C,
     ) -> Result<bool, ShardTreeError<S::Error>> {
@@ -633,14 +634,14 @@ impl<
             .get_checkpoint(checkpoint_id)
             .map_err(ShardTreeError::Storage)?
         {
-            self.truncate_removing_checkpoint_internal(checkpoint_id, &c)?;
+            self.truncate_to_checkpoint_internal(checkpoint_id, &c)?;
             Ok(true)
         } else {
             Ok(false)
         }
     }
 
-    fn truncate_removing_checkpoint_internal(
+    fn truncate_to_checkpoint_internal(
         &mut self,
         checkpoint_id: &C,
         checkpoint: &Checkpoint,
@@ -648,10 +649,10 @@ impl<
         match checkpoint.tree_state() {
             TreeState::Empty => {
                 self.store
-                    .truncate(Address::from_parts(Self::subtree_level(), 0))
+                    .truncate_shards(0)
                     .map_err(ShardTreeError::Storage)?;
                 self.store
-                    .truncate_checkpoints(checkpoint_id)
+                    .truncate_checkpoints_retaining(checkpoint_id)
                     .map_err(ShardTreeError::Storage)?;
                 self.store
                     .put_cap(Tree::empty())
@@ -670,21 +671,21 @@ impl<
                     root: self.store.get_cap().map_err(ShardTreeError::Storage)?,
                 };
 
-                if let Some(truncated) = cap_tree.truncate_to_position(position) {
+                if let Some(truncated_cap) = cap_tree.truncate_to_position(position) {
                     self.store
-                        .put_cap(truncated.root)
+                        .put_cap(truncated_cap.root)
                         .map_err(ShardTreeError::Storage)?;
                 };
 
                 if let Some(truncated) = replacement {
                     self.store
-                        .truncate(subtree_addr)
+                        .truncate_shards(subtree_addr.index())
                         .map_err(ShardTreeError::Storage)?;
                     self.store
                         .put_shard(truncated)
                         .map_err(ShardTreeError::Storage)?;
                     self.store
-                        .truncate_checkpoints(checkpoint_id)
+                        .truncate_checkpoints_retaining(checkpoint_id)
                         .map_err(ShardTreeError::Storage)?;
                 }
             }
@@ -1000,100 +1001,140 @@ impl<
         }
     }
 
-    /// Returns the position of the rightmost leaf inserted as of the given checkpoint.
+    /// Returns the position of the rightmost leaf as of the checkpoint at the given depth.
     ///
-    /// Returns the maximum leaf position if `checkpoint_depth == 0` (or `Ok(None)` in this
+    /// Returns either the position of the checkpoint at the requested depth, the maximum leaf
-    /// case if the tree is empty) or an error if the checkpointed position cannot be restored
+    /// position in the tree if no checkpoint depth is provided, or `Ok(None)` if the tree is
-    /// because it has been pruned. Note that no actual level-0 leaf may exist at this position.
+    /// empty.
+    ///
+    /// Returns `ShardTreeError::Query(QueryError::CheckpointPruned)` if
+    /// `checkpoint_depth.is_some()` and no checkpoint exists at the requested depth.
+    ///
+    /// ## Parameters
+    /// - `checkpoint_depth`: A zero-based index over the checkpoints that have been added to the
+    ///   tree, in reverse checkpoint identifier order, or `None` to request the overall maximum
+    ///   leaf position in the tree.
     pub fn max_leaf_position(
         &self,
-        checkpoint_depth: usize,
+        checkpoint_depth: Option<usize>,
     ) -> Result<Option<Position>, ShardTreeError<S::Error>> {
-        Ok(if checkpoint_depth == 0 {
+        match checkpoint_depth {
+            None => {
                 // TODO: This relies on the invariant that the last shard in the subtrees vector is
                 // never created without a leaf then being added to it. However, this may be a
                 // difficult invariant to maintain when adding empty roots, so perhaps we need a
                 // better way of tracking the actual max position of the tree; we might want to
                 // just store it directly.
-            self.store
+                Ok(self
+                    .store
                     .last_shard()
                     .map_err(ShardTreeError::Storage)?
-                .and_then(|t| t.max_position())
+                    .and_then(|t| t.max_position()))
-        } else {
+            }
-            match self
+            Some(depth) => self
-                .store
+                .store
-                .get_checkpoint_at_depth(checkpoint_depth)
+                .get_checkpoint_at_depth(depth)
-                .map_err(ShardTreeError::Storage)?
+                .map_err(ShardTreeError::Storage)?
-            {
+                .map(|(_, c)| c.position())
-                Some((_, c)) => Ok(c.position()),
+                .ok_or(ShardTreeError::Query(QueryError::CheckpointPruned)),
-                None => {
-                    // There is no checkpoint at the specified depth, so we report it as pruned.
-                    Err(QueryError::CheckpointPruned)
         }
-            }?
-        })
     }
 
     /// 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
+    /// Returns `Ok(None)` if no checkpoint matches the specified ID.
-            .store
+    pub fn root_at_checkpoint_id(
+        &self,
+        checkpoint: &C,
+    ) -> Result<Option<H>, ShardTreeError<S::Error>> {
+        self.store
             .get_checkpoint(checkpoint)
             .map_err(ShardTreeError::Storage)?
-            .map(|c| c.position())
+            .map(|c| {
-            .ok_or(QueryError::CheckpointPruned)?;
+                c.position().map_or_else(
-
-        position.map_or_else(
                     || Ok(H::empty_root(Self::root_addr().level())),
                     |pos| self.root(Self::root_addr(), pos + 1),
                 )
+            })
+            .transpose()
     }
 
     /// Computes the root of the tree as of the checkpointed position having the specified
     /// checkpoint id, caching intermediate values produced while computing the root.
+    ///
+    /// Returns `Ok(None)` if no checkpoint matches the specified ID.
     pub fn root_at_checkpoint_id_caching(
         &mut self,
         checkpoint: &C,
-    ) -> Result<H, ShardTreeError<S::Error>> {
+    ) -> Result<Option<H>, ShardTreeError<S::Error>> {
-        let position = self
+        self.store
-            .store
             .get_checkpoint(checkpoint)
             .map_err(ShardTreeError::Storage)?
-            .map(|c| c.position())
+            .map(|c| {
-            .ok_or(QueryError::CheckpointPruned)?;
+                c.position().map_or_else(
-
-        position.map_or_else(
                     || Ok(H::empty_root(Self::root_addr().level())),
                     |pos| self.root_caching(Self::root_addr(), pos + 1),
                 )
+            })
+            .transpose()
     }
 
     /// 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
+    /// Returns `Ok(None)` if no checkpoint exists at the requested depth.
-    /// that if the most recently appended leaf is also a checkpoint, this will return the same
+    ///
-    /// result as `checkpoint_depth == 1`.
+    /// ## Parameters
+    /// - `checkpoint_depth`: A zero-based index over the checkpoints that have been added to the
+    ///   tree, in reverse checkpoint identifier order, or `None` to request the root computed over
+    ///   all of the leaves in the tree.
     pub fn root_at_checkpoint_depth(
         &self,
-        checkpoint_depth: usize,
+        checkpoint_depth: Option<usize>,
-    ) -> Result<H, ShardTreeError<S::Error>> {
+    ) -> Result<Option<H>, ShardTreeError<S::Error>> {
-        self.max_leaf_position(checkpoint_depth)?.map_or_else(
+        self.max_leaf_position(checkpoint_depth).map_or_else(
+            |err| match err {
+                ShardTreeError::Query(QueryError::CheckpointPruned) => Ok(None),
+                err => Err(err),
+            },
+            |position| {
+                position
+                    .map_or_else(
                         || Ok(H::empty_root(Self::root_addr().level())),
                         |pos| self.root(Self::root_addr(), pos + 1),
                     )
+                    .map(Some)
+            },
+        )
     }
 
     /// Computes the root of the tree as of the checkpointed position at the specified depth,
     /// caching intermediate values produced while computing the root.
+    ///
+    /// Returns `Ok(None)` if no checkpoint exists at the requested depth.
+    ///
+    /// ## Parameters
+    /// - `checkpoint_depth`: A zero-based index over the checkpoints that have been added to the
+    ///   tree, in reverse checkpoint identifier order, or `None` to request the root computed over
+    ///   all of the leaves in the tree.
     pub fn root_at_checkpoint_depth_caching(
         &mut self,
-        checkpoint_depth: usize,
+        checkpoint_depth: Option<usize>,
-    ) -> Result<H, ShardTreeError<S::Error>> {
+    ) -> Result<Option<H>, ShardTreeError<S::Error>> {
-        self.max_leaf_position(checkpoint_depth)?.map_or_else(
+        self.max_leaf_position(checkpoint_depth).map_or_else(
+            |err| match err {
+                ShardTreeError::Query(QueryError::CheckpointPruned) => Ok(None),
+                err => Err(err),
+            },
+            |position| {
+                position
+                    .map_or_else(
                         || Ok(H::empty_root(Self::root_addr().level())),
                         |pos| self.root_caching(Self::root_addr(), pos + 1),
                     )
+                    .map(Some)
+            },
+        )
     }
 
     fn witness_helper<Ctx>(
@@ -1155,25 +1196,27 @@ impl<
     /// Computes the witness for the leaf at the specified position, as of the given checkpoint
     /// depth.
     ///
-    /// Returns the witness as of the most recently appended leaf if `checkpoint_depth == 0`. Note
+    /// Returns `ShardTreeError::Query(QueryError::CheckpointPruned)` if no checkpoint exists at
-    /// that if the most recently appended leaf is also a checkpoint, this will return the same
+    /// the requested depth.
-    /// result as `checkpoint_depth == 1`.
     pub fn witness_at_checkpoint_depth(
         &self,
         position: Position,
         checkpoint_depth: usize,
-    ) -> Result<MerklePath<H, DEPTH>, ShardTreeError<S::Error>> {
+    ) -> Result<Option<MerklePath<H, DEPTH>>, ShardTreeError<S::Error>> {
-        let as_of = self.max_leaf_position(checkpoint_depth).and_then(|v| {
+        let checkpoint = self
-            v.ok_or_else(|| QueryError::TreeIncomplete(vec![Self::root_addr()]).into())
+            .store
+            .get_checkpoint_at_depth(checkpoint_depth)
+            .map_err(ShardTreeError::Storage)?;
+
+        match checkpoint {
+            None => Ok(None),
+            Some((_, c)) => {
+                let as_of = c.position().filter(|p| position <= *p).ok_or_else(|| {
+                    QueryError::NotContained(Address::from_parts(Level::from(0), position.into()))
                 })?;
 
-        if position > as_of {
+                self.witness_internal(position, as_of).map(Some)
-            Err(
+            }
-                QueryError::NotContained(Address::from_parts(Level::from(0), position.into()))
-                    .into(),
-            )
-        } else {
-            self.witness_internal(position, as_of)
         }
     }
 
@@ -1183,39 +1226,56 @@ impl<
     /// This implementation will mutate the tree to cache intermediate root (ommer) values that are
     /// computed in the process of constructing the witness, so as to avoid the need to recompute
     /// those values from potentially large numbers of subtree roots in the future.
+    ///
+    /// Returns `ShardTreeError::Query(QueryError::CheckpointPruned)` if no checkpoint exists at
+    /// the requested depth. It is not possible to
     pub fn witness_at_checkpoint_depth_caching(
         &mut self,
         position: Position,
         checkpoint_depth: usize,
-    ) -> Result<MerklePath<H, DEPTH>, ShardTreeError<S::Error>> {
+    ) -> Result<Option<MerklePath<H, DEPTH>>, ShardTreeError<S::Error>> {
-        let as_of = self.max_leaf_position(checkpoint_depth).and_then(|v| {
+        let checkpoint = self
-            v.ok_or_else(|| QueryError::TreeIncomplete(vec![Self::root_addr()]).into())
+            .store
+            .get_checkpoint_at_depth(checkpoint_depth)
+            .map_err(ShardTreeError::Storage)?;
+
+        match checkpoint {
+            None => Ok(None),
+            Some((_, c)) => {
+                let as_of = c.position().filter(|p| position <= *p).ok_or_else(|| {
+                    QueryError::NotContained(Address::from_parts(Level::from(0), position.into()))
                 })?;
 
-        if position > as_of {
+                self.witness_internal_caching(position, as_of).map(Some)
-            Err(
+            }
-                QueryError::NotContained(Address::from_parts(Level::from(0), position.into()))
-                    .into(),
-            )
-        } else {
-            self.witness_internal_caching(position, as_of)
         }
     }
 
     /// Computes the witness for the leaf at the specified position, as of the given checkpoint.
+    ///
+    /// Returns Ok(None) if no such checkpoint exists.
     pub fn witness_at_checkpoint_id(
         &self,
         position: Position,
         checkpoint_id: &C,
-    ) -> Result<MerklePath<H, DEPTH>, ShardTreeError<S::Error>> {
+    ) -> Result<Option<MerklePath<H, DEPTH>>, ShardTreeError<S::Error>> {
-        let as_of = self
+        self.store
-            .store
             .get_checkpoint(checkpoint_id)
             .map_err(ShardTreeError::Storage)?
-            .and_then(|c| c.position())
+            .map(|checkpoint| {
-            .ok_or(QueryError::CheckpointPruned)?;
+                let as_of = checkpoint
+                    .position()
+                    .filter(|p| position <= *p)
+                    .ok_or_else(|| {
+                        QueryError::NotContained(Address::from_parts(
+                            Level::from(0),
+                            position.into(),
+                        ))
+                    })?;
 
                 self.witness_internal(position, as_of)
+            })
+            .transpose()
     }
 
     /// Computes the witness for the leaf at the specified position, as of the given checkpoint.
@@ -1223,19 +1283,30 @@ impl<
     /// This implementation will mutate the tree to cache intermediate root (ommer) values that are
     /// computed in the process of constructing the witness, so as to avoid the need to recompute
     /// those values from potentially large numbers of subtree roots in the future.
+    ///
+    /// Returns Ok(None) if no such checkpoint exists.
     pub fn witness_at_checkpoint_id_caching(
         &mut self,
         position: Position,
         checkpoint_id: &C,
-    ) -> Result<MerklePath<H, DEPTH>, ShardTreeError<S::Error>> {
+    ) -> Result<Option<MerklePath<H, DEPTH>>, ShardTreeError<S::Error>> {
-        let as_of = self
+        self.store
-            .store
             .get_checkpoint(checkpoint_id)
             .map_err(ShardTreeError::Storage)?
-            .and_then(|c| c.position())
+            .map(|checkpoint| {
-            .ok_or(QueryError::CheckpointPruned)?;
+                let as_of = checkpoint
+                    .position()
+                    .filter(|p| position <= *p)
+                    .ok_or_else(|| {
+                        QueryError::NotContained(Address::from_parts(
+                            Level::from(0),
+                            position.into(),
+                        ))
+                    })?;
 
                 self.witness_internal_caching(position, as_of)
+            })
+            .transpose()
     }
 
     /// Make a marked leaf at a position eligible to be pruned.
@@ -1427,7 +1498,7 @@ mod tests {
     fn avoid_pruning_reference() {
         fn test_with_marking(
             frontier_marking: Marking,
-        ) -> Result<MerklePath<String, 6>, ShardTreeError<Infallible>> {
+        ) -> Result<Option<MerklePath<String, 6>>, ShardTreeError<Infallible>> {
             let mut tree = ShardTree::<MemoryShardStore<String, usize>, 6, 3>::new(
                 MemoryShardStore::empty(),
                 5,
@@ -1508,7 +1579,7 @@ mod tests {
         .unwrap();
 
         let witness = test_with_marking(Marking::Reference).unwrap();
-        assert_eq!(witness, expected_witness);
+        assert_eq!(witness, Some(expected_witness));
     }
 
     // Combined tree tests
@@ -1578,17 +1649,21 @@ mod tests {
         fn check_shardtree_caching(
             (mut tree, _, marked_positions) in arb_shardtree(arb_char_str())
         ) {
-            if let Some(max_leaf_pos) = tree.max_leaf_position(0).unwrap() {
+            if let Some(max_leaf_pos) = tree.max_leaf_position(None).unwrap() {
                 let max_complete_addr = Address::above_position(max_leaf_pos.root_level(), max_leaf_pos);
                 let root = tree.root(max_complete_addr, max_leaf_pos + 1);
                 let caching_root = tree.root_caching(max_complete_addr, max_leaf_pos + 1);
                 assert_matches!(root, Ok(_));
                 assert_eq!(root, caching_root);
+            }
 
+            if let Ok(Some(checkpoint_position)) = tree.max_leaf_position(Some(0)) {
                 for pos in marked_positions {
                     let witness = tree.witness_at_checkpoint_depth(pos, 0);
                     let caching_witness = tree.witness_at_checkpoint_depth_caching(pos, 0);
+                    if pos <= checkpoint_position {
                         assert_matches!(witness, Ok(_));
+                    }
                     assert_eq!(witness, caching_witness);
                 }
             }

shardtree/src/store.rs

@@ -65,11 +65,8 @@ pub trait ShardStore {
     fn get_shard_roots(&self) -> Result<Vec<Address>, Self::Error>;
 
     /// Removes subtrees from the underlying store having root addresses at indices greater
-    /// than or equal to that of the specified address.
+    /// than or equal to that of the specified index.
-    ///
+    fn truncate_shards(&mut self, shard_index: u64) -> Result<(), Self::Error>;
-    /// Implementations of this method MUST enforce the constraint that the root address
-    /// provided has level `SHARD_HEIGHT`.
-    fn truncate(&mut self, from: Address) -> Result<(), Self::Error>;
 
     /// A tree that is used to cache the known roots of subtrees in the "cap" - the top part of the
     /// tree, which contains parent nodes produced by hashing the roots of the individual shards.
@@ -96,9 +93,11 @@ pub trait ShardStore {
     /// Returns the number of checkpoints maintained by the data store
     fn checkpoint_count(&self) -> Result<usize, Self::Error>;
 
-    /// Returns the id and position of the checkpoint, if any. Returns `None` if
+    /// Returns the id and position of the checkpoint at the specified depth, if it exists.
-    /// `checkpoint_depth == 0` or if insufficient checkpoints exist to seek back
+    ///
-    /// to the requested depth.
+    /// Returns `None` if if insufficient checkpoints exist to seek back to the requested depth.
+    /// Depth 0 refers to the most recent checkpoint in the tree; depth 1 refers to the previous
+    /// checkpoint, and so forth.
     fn get_checkpoint_at_depth(
         &self,
         checkpoint_depth: usize,
@@ -140,8 +139,9 @@ pub trait ShardStore {
     /// If no checkpoint exists with the given ID, this does nothing.
     fn remove_checkpoint(&mut self, checkpoint_id: &Self::CheckpointId) -> Result<(), Self::Error>;
 
-    /// Removes checkpoints with identifiers greater than or equal to the given identifier.
+    /// Removes checkpoints with identifiers greater than to the given identifier, and removes mark
-    fn truncate_checkpoints(
+    /// removal metadata from the specified checkpoint.
+    fn truncate_checkpoints_retaining(
         &mut self,
         checkpoint_id: &Self::CheckpointId,
     ) -> Result<(), Self::Error>;
@@ -179,8 +179,8 @@ impl<S: ShardStore> ShardStore for &mut S {
         S::put_cap(*self, cap)
     }
 
-    fn truncate(&mut self, from: Address) -> Result<(), Self::Error> {
+    fn truncate_shards(&mut self, shard_index: u64) -> Result<(), Self::Error> {
-        S::truncate(*self, from)
+        S::truncate_shards(*self, shard_index)
     }
 
     fn min_checkpoint_id(&self) -> Result<Option<Self::CheckpointId>, Self::Error> {
@@ -246,11 +246,11 @@ impl<S: ShardStore> ShardStore for &mut S {
         S::remove_checkpoint(self, checkpoint_id)
     }
 
-    fn truncate_checkpoints(
+    fn truncate_checkpoints_retaining(
         &mut self,
         checkpoint_id: &Self::CheckpointId,
     ) -> Result<(), Self::Error> {
-        S::truncate_checkpoints(self, checkpoint_id)
+        S::truncate_checkpoints_retaining(self, checkpoint_id)
     }
 }
 

shardtree/src/store/caching.rs

@@ -9,9 +9,9 @@ use crate::{LocatedPrunableTree, PrunableTree};
 
 #[derive(Debug)]
 enum Action<C> {
-    Truncate(Address),
+    TruncateShards(u64),
     RemoveCheckpoint(C),
-    TruncateCheckpoints(C),
+    TruncateCheckpointsRetaining(C),
 }
 
 /// An implementation of [`ShardStore`] that caches all state in memory.
@@ -63,12 +63,12 @@ where
     pub fn flush(mut self) -> Result<S, S::Error> {
         for action in &self.deferred_actions {
             match action {
-                Action::Truncate(from) => self.backend.truncate(*from),
+                Action::TruncateShards(index) => self.backend.truncate_shards(*index),
                 Action::RemoveCheckpoint(checkpoint_id) => {
                     self.backend.remove_checkpoint(checkpoint_id)
                 }
-                Action::TruncateCheckpoints(checkpoint_id) => {
+                Action::TruncateCheckpointsRetaining(checkpoint_id) => {
-                    self.backend.truncate_checkpoints(checkpoint_id)
+                    self.backend.truncate_checkpoints_retaining(checkpoint_id)
                 }
             }?;
         }
@@ -131,9 +131,10 @@ where
         self.cache.get_shard_roots()
     }
 
-    fn truncate(&mut self, from: Address) -> Result<(), Self::Error> {
+    fn truncate_shards(&mut self, shard_index: u64) -> Result<(), Self::Error> {
-        self.deferred_actions.push(Action::Truncate(from));
+        self.deferred_actions
-        self.cache.truncate(from)
+            .push(Action::TruncateShards(shard_index));
+        self.cache.truncate_shards(shard_index)
     }
 
     fn get_cap(&self) -> Result<PrunableTree<Self::H>, Self::Error> {
@@ -208,13 +209,13 @@ where
         self.cache.remove_checkpoint(checkpoint_id)
     }
 
-    fn truncate_checkpoints(
+    fn truncate_checkpoints_retaining(
         &mut self,
         checkpoint_id: &Self::CheckpointId,
     ) -> Result<(), Self::Error> {
         self.deferred_actions
-            .push(Action::TruncateCheckpoints(checkpoint_id.clone()));
+            .push(Action::TruncateCheckpointsRetaining(checkpoint_id.clone()));
-        self.cache.truncate_checkpoints(checkpoint_id)
+        self.cache.truncate_checkpoints_retaining(checkpoint_id)
     }
 }
 
@@ -269,22 +270,22 @@ mod tests {
             );
 
             assert_eq!(
-                tree.root(0).unwrap(),
+                tree.root(None).unwrap(),
                 String::combine_all(tree.depth(), &[]),
             );
             assert!(tree.append(String::from_u64(0), Ephemeral));
             assert_eq!(
-                tree.root(0).unwrap(),
+                tree.root(None).unwrap(),
                 String::combine_all(tree.depth(), &[0]),
             );
             assert!(tree.append(String::from_u64(1), Ephemeral));
             assert_eq!(
-                tree.root(0).unwrap(),
+                tree.root(None).unwrap(),
                 String::combine_all(tree.depth(), &[0, 1]),
             );
             assert!(tree.append(String::from_u64(2), Ephemeral));
             assert_eq!(
-                tree.root(0).unwrap(),
+                tree.root(None).unwrap(),
                 String::combine_all(tree.depth(), &[0, 1, 2]),
             );
 
@@ -310,7 +311,7 @@ mod tests {
                 assert!(t.append(String::from_u64(0), Ephemeral));
             }
             assert_eq!(
-                t.root(0).unwrap(),
+                t.root(None).unwrap(),
                 String::combine_all(t.depth(), &[0, 0, 0, 0])
             );
 
@@ -391,7 +392,13 @@ mod tests {
                 ShardTree::<_, 4, 3>::new(&mut rhs, 100),
             );
 
-            assert!(tree.append(String::from_u64(0), Marked));
+            assert!(tree.append(
+                String::from_u64(0),
+                Checkpoint {
+                    id: 0,
+                    marking: Marking::Marked
+                }
+            ));
             assert_eq!(
                 tree.witness(Position::from(0), 0),
                 Some(vec![
@@ -402,7 +409,13 @@ mod tests {
                 ])
             );
 
-            assert!(tree.append(String::from_u64(1), Ephemeral));
+            assert!(tree.append(
+                String::from_u64(1),
+                Checkpoint {
+                    id: 1,
+                    marking: Marking::None
+                }
+            ));
             assert_eq!(
                 tree.witness(0.into(), 0),
                 Some(vec![
@@ -413,7 +426,13 @@ mod tests {
                 ])
             );
 
-            assert!(tree.append(String::from_u64(2), Marked));
+            assert!(tree.append(
+                String::from_u64(2),
+                Checkpoint {
+                    id: 2,
+                    marking: Marking::Marked
+                }
+            ));
             assert_eq!(
                 tree.witness(Position::from(2), 0),
                 Some(vec![
@@ -424,7 +443,13 @@ mod tests {
                 ])
             );
 
-            assert!(tree.append(String::from_u64(3), Ephemeral));
+            assert!(tree.append(
+                String::from_u64(3),
+                Checkpoint {
+                    id: 3,
+                    marking: Marking::None
+                }
+            ));
             assert_eq!(
                 tree.witness(Position::from(2), 0),
                 Some(vec![
@@ -435,7 +460,13 @@ mod tests {
                 ])
             );
 
-            assert!(tree.append(String::from_u64(4), Ephemeral));
+            assert!(tree.append(
+                String::from_u64(4),
+                Checkpoint {
+                    id: 4,
+                    marking: Marking::None
+                }
+            ));
             assert_eq!(
                 tree.witness(Position::from(2), 0),
                 Some(vec![
@@ -462,7 +493,13 @@ mod tests {
                 assert!(tree.append(String::from_u64(i), Ephemeral));
             }
             assert!(tree.append(String::from_u64(6), Marked));
-            assert!(tree.append(String::from_u64(7), Ephemeral));
+            assert!(tree.append(
+                String::from_u64(7),
+                Checkpoint {
+                    id: 0,
+                    marking: Marking::None
+                }
+            ));
 
             assert_eq!(
                 tree.witness(0.into(), 0),
@@ -491,7 +528,13 @@ mod tests {
             assert!(tree.append(String::from_u64(3), Marked));
             assert!(tree.append(String::from_u64(4), Marked));
             assert!(tree.append(String::from_u64(5), Marked));
-            assert!(tree.append(String::from_u64(6), Ephemeral));
+            assert!(tree.append(
+                String::from_u64(6),
+                Checkpoint {
+                    id: 0,
+                    marking: Marking::None
+                }
+            ));
 
             assert_eq!(
                 tree.witness(Position::from(5), 0),
@@ -518,7 +561,13 @@ mod tests {
                 assert!(tree.append(String::from_u64(i), Ephemeral));
             }
             assert!(tree.append(String::from_u64(10), Marked));
-            assert!(tree.append(String::from_u64(11), Ephemeral));
+            assert!(tree.append(
+                String::from_u64(11),
+                Checkpoint {
+                    id: 0,
+                    marking: Marking::None
+                }
+            ));
 
             assert_eq!(
                 tree.witness(Position::from(10), 0),
@@ -548,7 +597,7 @@ mod tests {
                     marking: Marking::Marked,
                 },
             ));
-            assert!(tree.rewind());
+            assert!(tree.rewind(0));
             for i in 1..4 {
                 assert!(tree.append(String::from_u64(i), Ephemeral));
             }
@@ -556,6 +605,7 @@ mod tests {
             for i in 5..8 {
                 assert!(tree.append(String::from_u64(i), Ephemeral));
             }
+            tree.checkpoint(2);
             assert_eq!(
                 tree.witness(0.into(), 0),
                 Some(vec![
@@ -591,7 +641,7 @@ mod tests {
                 },
             ));
             assert!(tree.append(String::from_u64(7), Ephemeral));
-            assert!(tree.rewind());
+            assert!(tree.rewind(0));
             assert_eq!(
                 tree.witness(Position::from(2), 0),
                 Some(vec![
@@ -619,7 +669,13 @@ mod tests {
             assert!(tree.append(String::from_u64(12), Marked));
             assert!(tree.append(String::from_u64(13), Marked));
             assert!(tree.append(String::from_u64(14), Ephemeral));
-            assert!(tree.append(String::from_u64(15), Ephemeral));
+            assert!(tree.append(
+                String::from_u64(15),
+                Checkpoint {
+                    id: 0,
+                    marking: Marking::None
+                }
+            ));
 
             assert_eq!(
                 tree.witness(Position::from(12), 0),
@@ -638,7 +694,13 @@ mod tests {
             let ops = (0..=11)
                 .map(|i| Append(String::from_u64(i), Marked))
                 .chain(Some(Append(String::from_u64(12), Ephemeral)))
-                .chain(Some(Append(String::from_u64(13), Ephemeral)))
+                .chain(Some(Append(
+                    String::from_u64(13),
+                    Checkpoint {
+                        id: 0,
+                        marking: Marking::None,
+                    },
+                )))
                 .chain(Some(Witness(11u64.into(), 0)))
                 .collect::<Vec<_>>();
 
@@ -700,7 +762,7 @@ mod tests {
                         marking: Marking::None,
                     },
                 ),
-                Witness(3u64.into(), 5),
+                Witness(3u64.into(), 4),
             ];
 
             let mut lhs = MemoryShardStore::empty();
@@ -750,7 +812,7 @@ mod tests {
                 ),
                 Append(String::from_u64(0), Ephemeral),
                 Append(String::from_u64(0), Ephemeral),
-                Witness(Position::from(3), 1),
+                Witness(Position::from(3), 0),
             ];
 
             let mut lhs = MemoryShardStore::empty();
@@ -796,8 +858,8 @@ mod tests {
                         marking: Marking::None,
                     },
                 ),
-                Rewind,
+                Rewind(0),
-                Rewind,
+                Rewind(1),
                 Witness(Position::from(7), 2),
             ];
 
@@ -808,6 +870,7 @@ mod tests {
                 ShardTree::<_, 4, 3>::new(&mut rhs, 100),
             );
 
+            // There is no checkpoint at depth 2, so we cannot compute a witness
             assert_eq!(Operation::apply_all(&ops, &mut tree), None);
 
             check_equal(lhs, rhs);
@@ -831,7 +894,7 @@ mod tests {
                         marking: Marking::None,
                     },
                 ),
-                Witness(Position::from(2), 2),
+                Witness(Position::from(2), 1),
             ];
 
             let mut lhs = MemoryShardStore::empty();
@@ -868,7 +931,7 @@ mod tests {
                 ShardTree::<_, 4, 3>::new(&mut rhs, 100),
             );
 
-            assert!(!t.rewind());
+            assert!(!t.rewind(0));
 
             check_equal(lhs, rhs);
         }
@@ -881,7 +944,7 @@ mod tests {
             );
 
             assert!(t.checkpoint(1));
-            assert!(t.rewind());
+            assert!(t.rewind(0));
 
             check_equal(lhs, rhs);
         }
@@ -897,7 +960,7 @@ mod tests {
             t.append("a".to_string(), Retention::Ephemeral);
             assert!(t.checkpoint(1));
             t.append("b".to_string(), Retention::Marked);
-            assert!(t.rewind());
+            assert!(t.rewind(0));
             assert_eq!(Some(Position::from(0)), t.current_position());
 
             check_equal(lhs, rhs);
@@ -913,7 +976,7 @@ mod tests {
 
             t.append("a".to_string(), Retention::Marked);
             assert!(t.checkpoint(1));
-            assert!(t.rewind());
+            assert!(t.rewind(0));
 
             check_equal(lhs, rhs);
         }
@@ -929,7 +992,7 @@ mod tests {
             t.append("a".to_string(), Retention::Marked);
             assert!(t.checkpoint(1));
             t.append("a".to_string(), Retention::Ephemeral);
-            assert!(t.rewind());
+            assert!(t.rewind(0));
             assert_eq!(Some(Position::from(0)), t.current_position());
 
             check_equal(lhs, rhs);
@@ -946,11 +1009,11 @@ mod tests {
             t.append("a".to_string(), Retention::Ephemeral);
             assert!(t.checkpoint(1));
             assert!(t.checkpoint(2));
-            assert!(t.rewind());
+            assert!(t.rewind(1));
             t.append("b".to_string(), Retention::Ephemeral);
-            assert!(t.rewind());
+            assert!(t.rewind(0));
             t.append("b".to_string(), Retention::Ephemeral);
-            assert_eq!(t.root(0).unwrap(), "ab______________");
+            assert_eq!(t.root(None).unwrap(), "ab______________");
 
             check_equal(lhs, rhs);
         }
@@ -1016,7 +1079,7 @@ mod tests {
 
             tree.append("e".to_string(), Retention::Marked);
             assert!(tree.checkpoint(1));
-            assert!(tree.rewind());
+            assert!(tree.rewind(0));
             assert!(tree.remove_mark(0u64.into()));
 
             check_equal(lhs, rhs);
@@ -1056,14 +1119,14 @@ mod tests {
             vec![
                 append_str("x", Retention::Marked),
                 Checkpoint(1),
-                Rewind,
+                Rewind(0),
                 unmark(0),
             ],
             vec![
                 append_str("d", Retention::Marked),
                 Checkpoint(1),
                 unmark(0),
-                Rewind,
+                Rewind(0),
                 unmark(0),
             ],
             vec![
@@ -1071,22 +1134,22 @@ mod tests {
                 Checkpoint(1),
                 Checkpoint(2),
                 unmark(0),
-                Rewind,
+                Rewind(0),
-                Rewind,
+                Rewind(1),
             ],
             vec![
                 append_str("s", Retention::Marked),
                 append_str("m", Retention::Ephemeral),
                 Checkpoint(1),
                 unmark(0),
-                Rewind,
+                Rewind(0),
                 unmark(0),
                 unmark(0),
             ],
             vec![
                 append_str("a", Retention::Marked),
                 Checkpoint(1),
-                Rewind,
+                Rewind(0),
                 append_str("a", Retention::Marked),
             ],
         ];
@@ -1189,7 +1252,7 @@ mod tests {
                 Checkpoint(1),
                 unmark(0),
                 Checkpoint(2),
-                Rewind,
+                Rewind(1),
                 append_str("b", Retention::Ephemeral),
                 witness(0, 0),
             ],
@@ -1197,7 +1260,7 @@ mod tests {
                 append_str("a", Retention::Marked),
                 Checkpoint(1),
                 Checkpoint(2),
-                Rewind,
+                Rewind(1),
                 append_str("a", Retention::Ephemeral),
                 unmark(0),
                 witness(0, 1),

shardtree/src/store/memory.rs

@@ -68,8 +68,8 @@ impl<H: Clone, C: Clone + Ord> ShardStore for MemoryShardStore<H, C> {
         Ok(self.shards.iter().map(|s| s.root_addr).collect())
     }
 
-    fn truncate(&mut self, from: Address) -> Result<(), Self::Error> {
+    fn truncate_shards(&mut self, shard_index: u64) -> Result<(), Self::Error> {
-        let shard_idx = usize::try_from(from.index()).expect("SHARD_HEIGHT > 64 is unsupported");
+        let shard_idx = usize::try_from(shard_index).expect("SHARD_HEIGHT > 64 is unsupported");
         self.shards.truncate(shard_idx);
         Ok(())
     }
@@ -107,15 +107,12 @@ impl<H: Clone, C: Clone + Ord> ShardStore for MemoryShardStore<H, C> {
         &self,
         checkpoint_depth: usize,
     ) -> Result<Option<(C, Checkpoint)>, Self::Error> {
-        Ok(if checkpoint_depth == 0 {
+        Ok(self
-            None
+            .checkpoints
-        } else {
-            self.checkpoints
             .iter()
             .rev()
-                .nth(checkpoint_depth - 1)
+            .nth(checkpoint_depth)
-                .map(|(id, c)| (id.clone(), c.clone()))
+            .map(|(id, c)| (id.clone(), c.clone())))
-        })
     }
 
     fn min_checkpoint_id(&self) -> Result<Option<C>, Self::Error> {
@@ -169,8 +166,15 @@ impl<H: Clone, C: Clone + Ord> ShardStore for MemoryShardStore<H, C> {
         Ok(())
     }
 
-    fn truncate_checkpoints(&mut self, checkpoint_id: &C) -> Result<(), Self::Error> {
+    fn truncate_checkpoints_retaining(
-        self.checkpoints.split_off(checkpoint_id);
+        &mut self,
+        checkpoint_id: &Self::CheckpointId,
+    ) -> Result<(), Self::Error> {
+        let mut rest = self.checkpoints.split_off(checkpoint_id);
+        if let Some(mut c) = rest.remove(checkpoint_id) {
+            c.marks_removed.clear();
+            self.checkpoints.insert(checkpoint_id.clone(), c);
+        }
         Ok(())
     }
 }

shardtree/src/testing.rs

@@ -165,7 +165,7 @@ impl<
     }
 
     fn current_position(&self) -> Option<Position> {
-        match ShardTree::max_leaf_position(self, 0) {
+        match ShardTree::max_leaf_position(self, None) {
             Ok(v) => v,
             Err(err) => panic!("current position query failed: {:?}", err),
         }
@@ -185,16 +185,16 @@ impl<
         }
     }
 
-    fn root(&self, checkpoint_depth: usize) -> Option<H> {
+    fn root(&self, checkpoint_depth: Option<usize>) -> Option<H> {
         match ShardTree::root_at_checkpoint_depth(self, checkpoint_depth) {
-            Ok(v) => Some(v),
+            Ok(v) => v,
             Err(err) => panic!("root computation failed: {:?}", err),
         }
     }
 
     fn witness(&self, position: Position, checkpoint_depth: usize) -> Option<Vec<H>> {
         match ShardTree::witness_at_checkpoint_depth(self, position, checkpoint_depth) {
-            Ok(p) => Some(p.path_elems().to_vec()),
+            Ok(p) => p.map(|p| p.path_elems().to_vec()),
             Err(ShardTreeError::Query(
                 QueryError::NotContained(_)
                 | QueryError::TreeIncomplete(_)
@@ -220,8 +220,12 @@ impl<
         ShardTree::checkpoint(self, checkpoint_id).unwrap()
     }
 
-    fn rewind(&mut self) -> bool {
+    fn checkpoint_count(&self) -> usize {
-        ShardTree::truncate_to_depth(self, 1).unwrap()
+        ShardStore::checkpoint_count(self.store()).unwrap()
+    }
+
+    fn rewind(&mut self, checkpoint_depth: usize) -> bool {
+        ShardTree::truncate_to_checkpoint_depth(self, checkpoint_depth).unwrap()
     }
 }
 
@@ -255,7 +259,7 @@ pub fn check_shardtree_insertion<
     );
 
     assert_matches!(
-        tree.root_at_checkpoint_depth(1),
+        tree.root_at_checkpoint_depth(Some(0)),
         Err(ShardTreeError::Query(QueryError::TreeIncomplete(v))) if v == vec![Address::from_parts(Level::from(0), 0)]
     );
 
@@ -272,13 +276,13 @@ pub fn check_shardtree_insertion<
     );
 
     assert_matches!(
-        tree.root_at_checkpoint_depth(0),
+        tree.root_at_checkpoint_depth(None),
-        Ok(h) if h == *"abcd____________"
+        Ok(Some(h)) if h == *"abcd____________"
     );
 
     assert_matches!(
-        tree.root_at_checkpoint_depth(1),
+        tree.root_at_checkpoint_depth(Some(0)),
-        Ok(h) if h == *"ab______________"
+        Ok(Some(h)) if h == *"ab______________"
     );
 
     assert_matches!(
@@ -309,7 +313,7 @@ pub fn check_shardtree_insertion<
     );
 
     assert_matches!(
-        tree.root_at_checkpoint_depth(0),
+        tree.root_at_checkpoint_depth(None),
         // The (0, 13) and (1, 7) incomplete subtrees are
         // not considered incomplete here because they appear
         // at the tip of the tree.
@@ -319,7 +323,7 @@ pub fn check_shardtree_insertion<
         ]
     );
 
-    assert_matches!(tree.truncate_to_depth(1), Ok(true));
+    assert_matches!(tree.truncate_to_checkpoint_depth(0), Ok(true));
 
     assert_matches!(
         tree.batch_insert(
@@ -330,13 +334,13 @@ pub fn check_shardtree_insertion<
     );
 
     assert_matches!(
-        tree.root_at_checkpoint_depth(0),
+        tree.root_at_checkpoint_depth(None),
-        Ok(h) if h == *"abcdefghijkl____"
+        Ok(Some(h)) if h == *"abcdefghijkl____"
     );
 
     assert_matches!(
-        tree.root_at_checkpoint_depth(1),
+        tree.root_at_checkpoint_depth(Some(1)),
-        Ok(h) if h == *"ab______________"
+        Ok(Some(h)) if h == *"ab______________"
     );
 }
 
@@ -399,7 +403,7 @@ pub fn check_witness_with_pruned_subtrees<
         .witness_at_checkpoint_depth(Position::from(26), 0)
         .unwrap();
     assert_eq!(
-        witness.path_elems(),
+        witness.expect("can produce a witness").path_elems(),
         &[
             "d",
             "ab",