🧺 use typed errors in public interface

Cargo.toml

@@ -24,27 +24,27 @@ migration = []
 
 [dependencies]
 anyhow = "1.0.38"
+blake3 = { version = "1.4.0", optional = true, features = ["traits-preview"] } 
 borsh = { version = "1.3.0" , features = ["derive", "de_strict_order"]}
 digest = "0.10" 
 hashbrown = "0.13.2"
+hex = "0.4"
+ics23 = { version = "0.11.0", optional = true}
 itertools = { version = "0.10.0", default-features = false }
 mirai-annotations = "1.10.1"
 num-derive = "0.3.3"
 num-traits = "0.2.14"
 parking_lot = { version = "0.12.1", optional = true } 
 serde = { version = "1.0.124", features = ["derive"] }
-thiserror = { version = "1.0.24", optional = true } 
 sha2 = { version = "0.10", optional = true } 
-blake3 = { version = "1.4.0", optional = true, features = ["traits-preview"] } 
-hex = "0.4"
+thiserror = { version = "1.0.24", optional = true } 
 tracing = "0.1"
-ics23 = { version = "0.11.0", optional = true}
 
 [dev-dependencies]
 hex = { version = "0.4", features = ["serde"] }
-rand = { version = "0.8.3" }
 parking_lot = { version = "0.12.1" } 
-serde_json = { version  = "1.0.95" }
 proptest = { version = "1.0.0" }
 proptest-derive = { version = "0.3.0" }
+rand = { version = "0.8.3" }
+serde_json = { version  = "1.0.95" }
 sha2 = "0.10"

src/iterator.rs

@@ -8,11 +8,11 @@
 
 use alloc::{sync::Arc, vec::Vec};
 
-use anyhow::{bail, ensure, format_err, Result};
+use anyhow::{bail, ensure};
 
 use crate::{
     node_type::{Child, InternalNode, Node, NodeKey},
-    storage::TreeReader,
+    storage::{TreeReader, TreeReaderExt},
     types::{
         nibble::{nibble_path::NibblePath, Nibble, ROOT_NIBBLE_HEIGHT},
         Version,
@@ -114,14 +114,35 @@ pub struct JellyfishMerkleIterator<R> {
     done: bool,
 }
 
+/// Errors that can occur when iterating across a [`JellyfishMerkleIterator<R>`].
+#[derive(Debug, thiserror::Error)]
+#[non_exhaustive]
+pub enum IteratorError<E> {
+    #[error(transparent)]
+    Read(E),
+    #[error("node with key `{key:?}` was missing")]
+    NodeMissing { key: NodeKey },
+    #[error(
+        "could not find a value at version `{version}` associated with key hash `{key_hash:?}`"
+    )]
+    ValueNotFound { version: u64, key_hash: KeyHash },
+    #[error("Should not reach a null node.")]
+    NullNodeReached,
+}
+
 impl<R> JellyfishMerkleIterator<R>
 where
     R: TreeReader,
+    <R as TreeReader>::Error: std::error::Error + Send + Sync + 'static,
 {
     /// Constructs a new iterator. This puts the internal state in the correct position, so the
     /// following `next` call will yield the smallest key that is greater or equal to
     /// `starting_key`.
-    pub fn new(reader: Arc<R>, version: Version, starting_key: KeyHash) -> Result<Self> {
+    pub fn new(
+        reader: Arc<R>,
+        version: Version,
+        starting_key: KeyHash,
+    ) -> Result<Self, anyhow::Error> {
         let mut parent_stack = Vec::new();
         let mut done = false;
 
@@ -201,7 +222,11 @@ where
 
     /// Constructs a new iterator. This puts the internal state in the correct position, so the
     /// following `next` call will yield the leaf at `start_idx`.
-    pub fn new_by_index(reader: Arc<R>, version: Version, start_idx: usize) -> Result<Self> {
+    pub fn new_by_index(
+        reader: Arc<R>,
+        version: Version,
+        start_idx: usize,
+    ) -> Result<Self, anyhow::Error> {
         let mut parent_stack = Vec::new();
 
         let mut current_node_key = NodeKey::new_empty_path(version);
@@ -256,7 +281,7 @@ where
         internal_node: &'a InternalNode,
         leaves_skipped: &mut usize,
         target_leaf_idx: usize,
-    ) -> Result<(Nibble, &'a Child)> {
+    ) -> Result<(Nibble, &'a Child), anyhow::Error> {
         for (nibble, child) in internal_node.children_sorted() {
             let child_leaf_count = child.leaf_count();
             // n.b. The index is 0-based, so to reach leaf at N, N previous ones need to be skipped.
@@ -274,8 +299,9 @@ where
 impl<R> Iterator for JellyfishMerkleIterator<R>
 where
     R: TreeReader,
+    <R as TreeReader>::Error: std::error::Error + Send + Sync + 'static,
 {
-    type Item = Result<(KeyHash, OwnedValue)>;
+    type Item = Result<(KeyHash, OwnedValue), IteratorError<R::Error>>;
 
     fn next(&mut self) -> Option<Self::Item> {
         if self.done {
@@ -284,27 +310,37 @@ where
 
         if self.parent_stack.is_empty() {
             let root_node_key = NodeKey::new_empty_path(self.version);
-            match self.reader.get_node(&root_node_key) {
-                Ok(Node::Leaf(leaf_node)) => {
+            match self
+                .reader
+                .get_node_option(&root_node_key)
+                .map_err(IteratorError::Read)
+            {
+                Ok(Some(Node::Leaf(leaf_node))) => {
                     // This means the entire tree has a single leaf node. The key of this leaf node
                     // is greater or equal to `starting_key` (otherwise we would have set `done` to
                     // true in `new`). Return the node and mark `self.done` so next time we return
                     // None.
                     self.done = true;
                     return match self
                         .reader
-                        .get_value(root_node_key.version(), leaf_node.key_hash())
+                        .get_value_option(root_node_key.version(), leaf_node.key_hash())
+                        .map_err(IteratorError::Read)
                     {
-                        Ok(value) => Some(Ok((leaf_node.key_hash(), value))),
+                        Ok(Some(value)) => Some(Ok((leaf_node.key_hash(), value))),
+                        Ok(None) => Some(Err(IteratorError::ValueNotFound {
+                            version: root_node_key.version(),
+                            key_hash: leaf_node.key_hash(),
+                        })),
                         Err(e) => Some(Err(e)),
                     };
                 }
-                Ok(Node::Internal(_)) => {
+                Ok(Some(Node::Internal(_))) => {
                     // This means `starting_key` is bigger than every key in this tree, or we have
                     // iterated past the last key.
                     return None;
                 }
-                Ok(Node::Null) => unreachable!("We would have set done to true in new."),
+                Ok(Some(Node::Null)) => unreachable!("We would have set done to true in new."),
+                Ok(None) => return Some(Err(IteratorError::NodeMissing { key: root_node_key })),
                 Err(err) => return Some(Err(err)),
             }
         }
@@ -324,25 +360,35 @@ where
                     .version,
                 child_index,
             );
-            match self.reader.get_node(&node_key) {
-                Ok(Node::Internal(internal_node)) => {
+            match self
+                .reader
+                .get_node_option(&node_key)
+                .map_err(IteratorError::Read)
+            {
+                Ok(Some(Node::Internal(internal_node))) => {
                     let visit_info = NodeVisitInfo::new(node_key, internal_node);
                     self.parent_stack.push(visit_info);
                 }
-                Ok(Node::Leaf(leaf_node)) => {
+                Ok(Some(Node::Leaf(leaf_node))) => {
                     return match self
                         .reader
-                        .get_value(node_key.version(), leaf_node.key_hash())
+                        .get_value_option(node_key.version(), leaf_node.key_hash())
+                        .map_err(IteratorError::Read)
                     {
-                        Ok(value) => {
+                        Ok(Some(value)) => {
                             let ret = (leaf_node.key_hash(), value);
                             Self::cleanup_stack(&mut self.parent_stack);
                             Some(Ok(ret))
                         }
+                        Ok(None) => Some(Err(IteratorError::ValueNotFound {
+                            version: node_key.version(),
+                            key_hash: leaf_node.key_hash(),
+                        })),
                         Err(e) => Some(Err(e)),
                     }
                 }
-                Ok(Node::Null) => return Some(Err(format_err!("Should not reach a null node."))),
+                Ok(Some(Node::Null)) => return Some(Err(IteratorError::NullNodeReached)),
+                Ok(None) => return Some(Err(IteratorError::NodeMissing { key: node_key })),
                 Err(err) => return Some(Err(err)),
             }
         }

src/lib.rs

@@ -110,7 +110,7 @@ pub use types::Version;
 pub mod storage {
     pub use node_type::{LeafNode, Node, NodeKey};
     pub use reader::HasPreimage;
-    pub use reader::TreeReader;
+    pub use reader::{TreeReader, TreeReaderExt};
     pub use types::nibble::nibble_path::NibblePath;
     pub use writer::{
         NodeBatch, NodeStats, StaleNodeIndex, StaleNodeIndexBatch, TreeUpdateBatch, TreeWriter,

src/node_type.rs

@@ -8,31 +8,29 @@
 //! [`JellyfishMerkleTree`](crate::JellyfishMerkleTree). [`InternalNode`] represents a 4-level
 //! binary tree to optimize for IOPS: it compresses a tree with 31 nodes into one node with 16
 //! chidren at the lowest level. [`LeafNode`] stores the full key and the value associated.
-use crate::storage::TreeReader;
 
-use crate::SimpleHasher;
-use alloc::format;
-use alloc::vec::Vec;
-use alloc::{boxed::Box, vec};
+use alloc::{boxed::Box, format, vec::Vec};
+
 use anyhow::Context;
 use borsh::{BorshDeserialize, BorshSerialize};
 use num_derive::{FromPrimitive, ToPrimitive};
-#[cfg(any(test))]
-use proptest::prelude::*;
-#[cfg(any(test))]
-use proptest_derive::Arbitrary;
 use serde::{Deserialize, Serialize};
 
-use crate::proof::SparseMerkleNode;
 use crate::{
+    storage::{TreeReader, TreeReaderExt},
     types::{
         nibble::{nibble_path::NibblePath, Nibble},
-        proof::{SparseMerkleInternalNode, SparseMerkleLeafNode},
+        proof::{SparseMerkleInternalNode, SparseMerkleLeafNode, SparseMerkleNode},
         Version,
     },
-    KeyHash, ValueHash, SPARSE_MERKLE_PLACEHOLDER_HASH,
+    KeyHash, SimpleHasher, ValueHash, SPARSE_MERKLE_PLACEHOLDER_HASH,
 };
 
+#[cfg(any(test))]
+use proptest::prelude::*;
+#[cfg(any(test))]
+use proptest_derive::Arbitrary;
+
 /// The unique key of each node.
 #[derive(
     Clone,
@@ -499,16 +497,21 @@ impl InternalNode {
     }
 
     /// [`build_sibling`] builds the sibling contained in the merkle tree between
-    /// [start; start+width) under the internal node (`self`) using the `TreeReader` as
-    /// a node reader to get the leaves/internal nodes at the bottom level of this internal node
-    fn build_sibling<H: SimpleHasher>(
+    /// [start; start+width) under the internal node (`self`) using the [`TreeReader`] as
+    /// a node reader to get the leaves/internal nodes at the bottom level of this internal node.
+    fn build_sibling<H, R>(
         &self,
-        tree_reader: &impl TreeReader,
+        tree_reader: &R,
         node_key: &NodeKey,
         start: u8,
         width: u8,
         (existence_bitmap, leaf_bitmap): (u16, u16),
-    ) -> SparseMerkleNode {
+    ) -> SparseMerkleNode
+    where
+        H: SimpleHasher,
+        R: TreeReader,
+        <R as TreeReader>::Error: std::error::Error + Send + Sync + 'static,
+    {
         // Given a bit [start, 1 << nibble_height], return the value of that range.
         let (range_existence_bitmap, range_leaf_bitmap) =
             Self::range_bitmaps(start, width, (existence_bitmap, leaf_bitmap));
@@ -673,13 +676,18 @@ impl InternalNode {
     ///     |   MSB|<---------------------- uint 16 ---------------------------->|LSB
     ///  height    chs: `child_half_start`         shs: `sibling_half_start`
     /// ```
-    fn get_child_with_siblings_helper<H: SimpleHasher>(
+    fn get_child_with_siblings_helper<H, R>(
         &self,
-        tree_reader: &impl TreeReader,
+        tree_reader: &R,
         node_key: &NodeKey,
         n: Nibble,
         get_only_child: bool,
-    ) -> (Option<NodeKey>, Vec<SparseMerkleNode>) {
+    ) -> (Option<NodeKey>, Vec<SparseMerkleNode>)
+    where
+        H: SimpleHasher,
+        R: TreeReader,
+        <R as TreeReader>::Error: std::error::Error + Send + Sync + 'static,
+    {
         let mut siblings: Vec<SparseMerkleNode> = vec![];
         let (existence_bitmap, leaf_bitmap) = self.generate_bitmaps();
 
@@ -692,7 +700,7 @@ impl InternalNode {
             let width = 1 << h;
             let (child_half_start, sibling_half_start) = get_child_and_sibling_half_start(n, h);
             // Compute the root hash of the subtree rooted at the sibling of `r`.
-            siblings.push(self.build_sibling::<H>(
+            siblings.push(self.build_sibling::<H, R>(
                 tree_reader,
                 node_key,
                 sibling_half_start,
@@ -762,13 +770,18 @@ impl InternalNode {
 
     /// [`get_child_with_siblings`] will return the child from this subtree that matches the nibble n in addition
     /// to building the list of its sibblings. This function has the same behavior as [`child`].
-    pub(crate) fn get_child_with_siblings<H: SimpleHasher>(
+    pub(crate) fn get_child_with_siblings<H, R>(
         &self,
-        tree_cache: &impl TreeReader,
+        tree_cache: &R,
         node_key: &NodeKey,
         n: Nibble,
-    ) -> (Option<NodeKey>, Vec<SparseMerkleNode>) {
-        self.get_child_with_siblings_helper::<H>(tree_cache, node_key, n, false)
+    ) -> (Option<NodeKey>, Vec<SparseMerkleNode>)
+    where
+        H: SimpleHasher,
+        R: TreeReader,
+        <R as TreeReader>::Error: std::error::Error + Send + Sync + 'static,
+    {
+        self.get_child_with_siblings_helper::<H, R>(tree_cache, node_key, n, false)
     }
 
     /// [`get_only_child_with_siblings`] will **either** return the child that matches the nibble n or the only
@@ -777,13 +790,18 @@ impl InternalNode {
     /// Even this leaf child is not the n-th child, it should be returned instead of
     /// `None` because it's existence indirectly proves the n-th child doesn't exist.
     /// Please read proof format for details.
-    pub(crate) fn get_only_child_with_siblings<H: SimpleHasher>(
+    pub(crate) fn get_only_child_with_siblings<H: SimpleHasher, R>(
         &self,
-        tree_reader: &impl TreeReader,
+        tree_reader: &R,
         node_key: &NodeKey,
         n: Nibble,
-    ) -> (Option<NodeKey>, Vec<SparseMerkleNode>) {
-        self.get_child_with_siblings_helper::<H>(tree_reader, node_key, n, true)
+    ) -> (Option<NodeKey>, Vec<SparseMerkleNode>)
+    where
+        H: SimpleHasher,
+        R: TreeReader,
+        <R as TreeReader>::Error: std::error::Error + Send + Sync + 'static,
+    {
+        self.get_child_with_siblings_helper::<H, R>(tree_reader, node_key, n, true)
     }
 
     #[cfg(test)]