Finish end to end test

src/lib.rs

@@ -207,6 +207,7 @@ impl<NR> Branch<NR> {
         let word_idx = self.mask.bit_idx as usize / 32;
         debug_assert!(word_idx < 8);
 
+        debug_assert!(self.prefix.len() <= word_idx);
         let prefix_offset = word_idx - self.prefix.len();
 
         let prefix_diff = iter::zip(self.prefix.iter(), key_hash.0.iter().skip(prefix_offset))
@@ -221,6 +222,7 @@ impl<NR> Branch<NR> {
             };
         }
 
+        // If sub wraps around to the last word, the prior word is 0.
         let prior_word_idx = word_idx.wrapping_sub(1);
         let prior_word = key_hash.0.get(prior_word_idx).unwrap_or(&0);
 
@@ -548,13 +550,14 @@ impl<S: Store<V>, V: AsRef<[u8]>> Transaction<S, V> {
                 on_modified_leaf(&hash, leaf)?;
                 Ok(hash)
             }
-            NodeRef::Stored(stored_idx) => {
-                let hash = data_store
-                    .get_unvisited_hash(*stored_idx)
-                    .copied()
-                    .map_err(|e| format!("Error in `calc_root_hash_node`: {e}"))?;
-                Ok(hash)
-            }
+            NodeRef::Stored(stored_idx) => data_store.calc_subtree_hash(*stored_idx).map_err(|e| {
+                format!(
+                    "Error in `calc_root_hash_node`: {e} at {file}:{line}:{column}",
+                    file = file!(),
+                    line = line!(),
+                    column = column!()
+                )
+            }),
         }
     }
 
@@ -821,7 +824,7 @@ impl<S: Store<V>, V: AsRef<[u8]>> Transaction<S, V> {
                         }
                         stored::Node::Leaf(leaf) => {
                             *next = NodeRef::ModBranch(Branch::new_from_leafs(
-                                branch.mask.word_idx() - 1,
+                                branch.mask.word_idx().saturating_sub(1),
                                 StoredLeafRef::new(leaf, *stored_idx),
                                 Box::new(Leaf {
                                     key_hash: *key_hash,

src/stored.rs

@@ -11,20 +11,20 @@ pub type Idx = u32;
 pub trait Store<V> {
     type Error: Display;
 
-    /// Must return a hash of a node that has not been visited.
-    /// May return a hash of a node that has already been visited.
-    fn get_unvisited_hash(&self, hash_idx: Idx) -> Result<&NodeHash, Self::Error>;
+    fn calc_subtree_hash(&self, hash_idx: Idx) -> Result<NodeHash, Self::Error>;
 
     fn get_node(&self, hash_idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>, Self::Error>;
 }
 
 impl<V, S: Store<V>> Store<V> for &S {
     type Error = S::Error;
 
-    fn get_unvisited_hash(&self, hash_idx: Idx) -> Result<&NodeHash, Self::Error> {
-        (**self).get_unvisited_hash(hash_idx)
+    #[inline(always)]
+    fn calc_subtree_hash(&self, hash_idx: Idx) -> Result<NodeHash, Self::Error> {
+        (**self).calc_subtree_hash(hash_idx)
     }
 
+    #[inline(always)]
     fn get_node(&self, hash_idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>, Self::Error> {
         (**self).get_node(hash_idx)
     }

src/stored/memory_db.rs

@@ -0,0 +1,45 @@
+use core::cell::RefCell;
+
+use alloc::collections::BTreeMap;
+
+use crate::{Branch, Leaf};
+
+use super::{DatabaseGet, DatabaseSet, Node, NodeHash};
+
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
+pub struct MemoryDb<V> {
+    leaves: RefCell<BTreeMap<NodeHash, Node<Branch<NodeHash>, Leaf<V>>>>,
+}
+
+impl<V> MemoryDb<V> {
+    pub fn empty() -> Self {
+        Self {
+            leaves: RefCell::default(),
+        }
+    }
+}
+
+impl<V: Clone> DatabaseGet<V> for MemoryDb<V> {
+    type GetError = String;
+
+    fn get(&self, hash: &NodeHash) -> Result<Node<Branch<NodeHash>, Leaf<V>>, Self::GetError> {
+        self.leaves
+            .borrow()
+            .get(hash)
+            .cloned()
+            .ok_or_else(|| format!("Hash: `{}` not found", hash))
+    }
+}
+
+impl<V: Clone> DatabaseSet<V> for MemoryDb<V> {
+    type SetError = String;
+
+    fn set(
+        &self,
+        hash: NodeHash,
+        node: Node<Branch<NodeHash>, Leaf<V>>,
+    ) -> Result<(), Self::SetError> {
+        self.leaves.borrow_mut().insert(hash, node);
+        Ok(())
+    }
+}

src/stored/merkle.rs

@@ -62,22 +62,26 @@ impl<V: AsRef<[u8]>> Snapshot<V> {
     /// Always check that the snapshot is of the merkle tree you expect.
     pub fn calc_root_hash(&self) -> Result<TrieRoot<NodeHash>> {
         match self.root_node_idx()? {
-            TrieRoot::Node(idx) => Ok(TrieRoot::Node(self.calc_root_hash_inner(idx)?)),
+            TrieRoot::Node(idx) => Ok(TrieRoot::Node(self.calc_subtree_hash(idx)?)),
             TrieRoot::Empty => Ok(TrieRoot::Empty),
         }
     }
+}
+
+impl<V: AsRef<[u8]>> Store<V> for Snapshot<V> {
+    type Error = Error;
 
     // TODO fix possible stack overflow
     // I dislike using an explicit mutable stack.
     // I have an idea for abusing async for high performance segmented stacks
-    fn calc_root_hash_inner(&self, node: Idx) -> Result<NodeHash> {
+    fn calc_subtree_hash(&self, node: Idx) -> Result<NodeHash> {
         let idx = node as usize;
         let leaf_offset = self.branches.len();
         let unvisited_offset = leaf_offset + self.leaves.len();
 
         if let Some(branch) = self.branches.get(idx) {
-            let left = self.calc_root_hash_inner(branch.left)?;
-            let right = self.calc_root_hash_inner(branch.right)?;
+            let left = self.calc_subtree_hash(branch.left)?;
+            let right = self.calc_subtree_hash(branch.right)?;
 
             Ok(branch.hash_branch(&left, &right))
         } else if let Some(leaf) = self.leaves.get(idx - leaf_offset) {
@@ -95,32 +99,8 @@ impl<V: AsRef<[u8]>> Snapshot<V> {
             ))
         }
     }
-}
-
-impl<V: AsRef<[u8]>> Store<V> for Snapshot<V> {
-    type Error = Error;
-
-    fn get_unvisited_hash(&self, idx: Idx) -> Result<&NodeHash> {
-        let error = || {
-            format!(
-                "Invalid snapshot: no unvisited node at index {}\n\
-                Snapshot has {} branches, {} leaves, and {} unvisited nodes",
-                idx,
-                self.branches.len(),
-                self.leaves.len(),
-                self.unvisited_nodes.len(),
-            )
-        };
-
-        let idx = idx as usize;
-        if idx < self.branches.len() + self.leaves.len() {
-            return Err(error());
-        }
-        let idx = idx - self.branches.len() - self.leaves.len();
-
-        self.unvisited_nodes.get(idx).ok_or_else(error)
-    }
 
+    #[inline]
     fn get_node(&self, idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>> {
         let idx = idx as usize;
         let leaf_offset = self.branches.len();
@@ -157,13 +137,13 @@ type NodeHashMaybeNode<'a, V> = (&'a NodeHash, Option<Node<&'a Branch<Idx>, &'a
 impl<'a, Db: DatabaseGet<V>, V: Clone> Store<V> for SnapshotBuilder<'a, Db, V> {
     type Error = Error;
 
-    fn get_unvisited_hash(&self, hash_idx: Idx) -> Result<&NodeHash, Self::Error> {
+    fn calc_subtree_hash(&self, hash_idx: Idx) -> Result<NodeHash, Self::Error> {
         let hash_idx = hash_idx as usize;
 
         self.nodes
             .borrow()
             .get(hash_idx)
-            .map(|(hash, _)| *hash)
+            .map(|(hash, _)| **hash)
             .ok_or_else(|| {
                 format!(
                     "Invalid snapshot: no unvisited node at index {}\n\

tests/build_store_modify.rs

@@ -16,7 +16,7 @@ prop_compose! {
 
 prop_compose! {
     fn arb_hashmap()(
-        map in prop::collection::hash_map(arb_key_hash(), 0u64.., 0..2)
+        map in prop::collection::hash_map(arb_key_hash(), 0u64.., 0..1_000)
     ) -> HashMap<KeyHash, u64> {
         map
     }
@@ -25,7 +25,7 @@ prop_compose! {
 proptest! {
     #[test]
     fn prop_end_to_end_example(
-        maps in prop::collection::vec(arb_hashmap(), 1..3)
+        maps in prop::collection::vec(arb_hashmap(), 1..100)
     ) {
         end_to_end_example(maps)
     }

tests/utils/mod.rs

@@ -35,7 +35,6 @@ pub fn run_against_snapshot(
     new_root_hash: TrieRoot<NodeHash>,
     old_root_hash: TrieRoot<NodeHash>,
 ) {
-    dbg!(&snapshot);
     assert_eq!(old_root_hash, snapshot.calc_root_hash().unwrap());
 
     let mut txn = Transaction::from_snapshot(&snapshot).unwrap();