Merge pull request #98 from nuttycom/shardtree/insert_frontier

shardtree: Add `ShardTree::insert_frontier`

shardtree/CHANGELOG.md

@@ -8,19 +8,20 @@ and this project adheres to Rust's notion of
 ## Unreleased
 
 ## Added
-* `Shardtree::store`
+- `ShardTree::{store, store_mut}`
+- `ShardTree::insert_frontier`
 
 ## [0.2.0] - 2023-11-07
 
 ## Added
-* `Shardtree::{root_at_checkpoint_id, root_at_checkpoint_id_caching}`
-* `Shardtree::{witness_at_checkpoint_id, witness_at_checkpoint_id_caching}`
+- `ShardTree::{root_at_checkpoint_id, root_at_checkpoint_id_caching}`
+- `ShardTree::{witness_at_checkpoint_id, witness_at_checkpoint_id_caching}`
 
 ## Changed
-* `Shardtree::root_at_checkpoint` and `Shardtree::root_at_checkpoint_caching` have
+- `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.
-* `Shardtree::witness` and `Shardtree::witness_caching` have
+- `ShardTree::witness` and `ShardTree::witness_caching` have
   been renamed to `witness_at_checkpoint_depth` and `witness_at_checkpoint_depth_caching`,
   respectively.
 

shardtree/src/error.rs

@@ -78,6 +78,10 @@ pub enum InsertionError {
     /// An input data structure had malformed data when attempting to insert a value
     /// at the given address
     InputMalformed(Address),
+    // The caller attempted to mark the empty tree state as corresponding to the state
+    // for a spendable note.
+    // TODO: Add this proper error type for `shardtree-0.3.0`
+    //MarkedRetentionInvalid,
 }
 
 impl fmt::Display for InsertionError {
@@ -104,7 +108,9 @@ impl fmt::Display for InsertionError {
             InsertionError::TreeFull => write!(f, "Note commitment tree is full."),
             InsertionError::InputMalformed(addr) => {
                 write!(f, "Input malformed for insertion at address {:?}", addr)
-            }
+            } //InsertionError::MarkedRetentionInvalid => {
+              //    write!(f, "Cannot use `Marked` retention for the empty tree.")
+              //}
         }
     }
 }

shardtree/src/lib.rs

@@ -24,6 +24,7 @@
 
 use core::fmt::Debug;
 use either::Either;
+use incrementalmerkletree::frontier::Frontier;
 use std::collections::{BTreeMap, BTreeSet};
 use std::sync::Arc;
 use tracing::trace;
@@ -96,6 +97,11 @@ impl<
         &self.store
     }
 
+    /// Returns a mutable reference to the underlying [`ShardStore`].
+    pub fn store_mut(&mut self) -> &mut S {
+        &mut self.store
+    }
+
     /// Returns the root address of the tree.
     pub fn root_addr() -> Address {
         Address::from_parts(Level::from(DEPTH), 0)
@@ -257,6 +263,42 @@ impl<
         Ok(())
     }
 
+    /// Add the leaf and ommers of the provided frontier to the tree.
+    ///
+    /// The leaf and ommers will be added as nodes within the subtree corresponding to the
+    /// frontier's position, and the tree state corresponding to that frontier will be marked as
+    /// specified by the leaf retention.
+    ///
+    /// This method may be used to add a checkpoint for the empty tree; note that
+    /// [`Retention::Marked`] is invalid for the empty tree.
+    pub fn insert_frontier(
+        &mut self,
+        frontier: Frontier<H, DEPTH>,
+        leaf_retention: Retention<C>,
+    ) -> Result<(), ShardTreeError<S::Error>> {
+        if let Some(nonempty_frontier) = frontier.take() {
+            self.insert_frontier_nodes(nonempty_frontier, leaf_retention)
+        } else {
+            match leaf_retention {
+                Retention::Ephemeral => Ok(()),
+                Retention::Checkpoint {
+                    id,
+                    is_marked: false,
+                } => self
+                    .store
+                    .add_checkpoint(id, Checkpoint::tree_empty())
+                    .map_err(ShardTreeError::Storage),
+                Retention::Checkpoint {
+                    is_marked: true, ..
+                }
+                | Retention::Marked => Err(ShardTreeError::Insert(
+                    //TODO: use InsertionError::MarkedRetentionInvalid for `shardtree-0.3.0`
+                    InsertionError::CheckpointOutOfOrder,
+                )),
+            }
+        }
+    }
+
     /// Add the leaf and ommers of the provided frontier as nodes within the subtree corresponding
     /// to the frontier's position, and update the cap to include the ommer nodes at levels greater
     /// than or equal to the shard height.