From be5e8673f95204dfe944c1a6aca4997b3da4dfbb Mon Sep 17 00:00:00 2001
From: Avi Dessauer <avi.the.coder@gmail.com>
Date: Fri, 16 Aug 2024 15:56:15 -0700
Subject: [PATCH 1/3] Add VerifiedSnapshot

---
 src/stored/merkle.rs | 162 ++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 160 insertions(+), 2 deletions(-)

diff --git a/src/stored/merkle.rs b/src/stored/merkle.rs
index a052a0b..ac52713 100644
--- a/src/stored/merkle.rs
+++ b/src/stored/merkle.rs
@@ -13,6 +13,163 @@ use super::{DatabaseGet, Idx, Node, NodeHash, Store};
 
 type Result<T, E = TrieError> = core::result::Result<T, E>;
 
+pub struct VerifiedSnapshot<V> {
+    snapshot: Snapshot<V>,
+
+    /// The root hash of the snapshot is the last hash in the slice.
+    /// The indexes of each hash match the indexes of nodes in the snapshot.
+    branch_hashes: Box<[NodeHash]>,
+    leaf_hashes: Box<[NodeHash]>,
+}
+
+impl<V: PortableHash> VerifiedSnapshot<V> {
+    #[inline]
+    pub fn verify_snapshot(
+        snapshot: Snapshot<V>,
+        hasher: &mut impl PortableHasher<32>,
+    ) -> Result<Self> {
+        // Check that the snapshot is well formed.
+        let _ = snapshot.root_node_idx()?;
+
+        let mut leaf_hashes = Vec::with_capacity(snapshot.leaves.len());
+        let mut branch_hashes = Vec::with_capacity(snapshot.branches.len());
+
+        for leaf in snapshot.leaves.iter() {
+            leaf_hashes.push(leaf.hash_leaf(hasher));
+        }
+
+        let leaf_offset = snapshot.branches.len();
+        let unvisited_offset = leaf_offset + snapshot.leaves.len();
+
+        for (idx, branch) in snapshot.branches.iter().enumerate() {
+            let hash_of_child = |child| {
+                if child < idx {
+                    branch_hashes.get(child).ok_or_else(|| {
+                        format!(
+                            "Invalid snapshot: branch {} has child {},
+                            child branch index must be less than parent, branches are not in post-order traversal",
+                            idx, child
+                        )
+                    })
+                } else if child < leaf_offset {
+                    leaf_hashes.get(child).ok_or_else(|| {
+                        format!(
+                            "Invalid snapshot: branch {} has child {},
+                            child leaf does not exist",
+                            idx, child
+                        )
+                    })
+                } else {
+                    snapshot
+                        .unvisited_nodes
+                        .get(child - unvisited_offset)
+                        .ok_or_else(|| {
+                            format!(
+                                "Invalid snapshot: branch {} has child {},
+                                child unvisited node does not exist",
+                                idx, child
+                            )
+                        })
+                }
+            };
+
+            let left_hash = hash_of_child(branch.left as usize)?;
+            let right_hash = hash_of_child(branch.right as usize)?;
+
+            branch_hashes.push(branch.hash_branch(hasher, left_hash, right_hash));
+        }
+
+        Ok(VerifiedSnapshot {
+            snapshot,
+            branch_hashes: branch_hashes.into_boxed_slice(),
+            leaf_hashes: leaf_hashes.into_boxed_slice(),
+        })
+    }
+
+    #[inline]
+    pub fn trie_root(&self) -> TrieRoot<NodeRef<V>> {
+        if !self.snapshot.branches.is_empty() {
+            TrieRoot::Node(NodeRef::Stored(self.snapshot.branches.len() as Idx - 1))
+        } else {
+            // We know that the snapshot is valid, because we verified it in `verify_snapshot`.
+            // If a snapshot contains no branches a single leaf or a single unvisited node is a valid snapshot.
+            // Any other combination is invalid.
+            debug_assert_eq!(
+                self.snapshot.leaves.len() + self.snapshot.unvisited_nodes.len(),
+                1
+            );
+            TrieRoot::Node(NodeRef::Stored(0))
+        }
+    }
+}
+
+impl<V> Store<V> for VerifiedSnapshot<V> {
+    type Error = TrieError;
+
+    #[inline]
+    fn calc_subtree_hash(
+        &self,
+        _: &mut impl PortableHasher<32>,
+        node: Idx,
+    ) -> Result<NodeHash, Self::Error> {
+        let idx = node as usize;
+        let leaf_offset = self.snapshot.branches.len();
+        let unvisited_offset = leaf_offset + self.snapshot.leaves.len();
+
+        if let Some(branch) = self.branch_hashes.get(idx) {
+            Ok(*branch)
+        } else if let Some(leaf) = self.leaf_hashes.get(idx - leaf_offset) {
+            Ok(*leaf)
+        } else if let Some(hash) = self.snapshot.unvisited_nodes.get(idx - unvisited_offset) {
+            Ok(*hash)
+        } else {
+            Err(format!(
+                "Invalid arg: node {} does not exist\n\
+                Snapshot has {} nodes",
+                idx,
+                self.snapshot.branches.len()
+                    + self.snapshot.leaves.len()
+                    + self.snapshot.unvisited_nodes.len(),
+            )
+            .into())
+        }
+    }
+
+    #[inline]
+    fn get_node(&self, idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>> {
+        let idx = idx as usize;
+        let leaf_offset = self.snapshot.branches.len();
+        let unvisited_offset = leaf_offset + self.snapshot.leaves.len();
+
+        if let Some(branch) = self.snapshot.branches.get(idx) {
+            Ok(Node::Branch(branch))
+        } else if let Some(leaf) = self.snapshot.leaves.get(idx - leaf_offset) {
+            Ok(Node::Leaf(leaf))
+        } else if self
+            .snapshot
+            .unvisited_nodes
+            .get(idx - unvisited_offset)
+            .is_some()
+        {
+            Err(format!(
+                "Invalid arg: node {idx} is unvisited\n\
+                get_node can only return visited nodes"
+            )
+            .into())
+        } else {
+            Err(format!(
+                "Invalid arg: node {} does not exist\n\
+                Snapshot has {} nodes",
+                idx,
+                self.snapshot.branches.len()
+                    + self.snapshot.leaves.len()
+                    + self.snapshot.unvisited_nodes.len(),
+            )
+            .into())
+        }
+    }
+}
+
 /// A snapshot of the merkle trie
 ///
 /// Contains visited nodes and unvisited nodes
@@ -39,8 +196,9 @@ impl<V: PortableHash> Snapshot<V> {
         ) {
             // A empty tree
             ([], [], []) => Ok(TrieRoot::Empty),
-            // A tree with only one node
-            ([_], [], []) | ([], [_], []) | ([], [], [_]) => Ok(TrieRoot::Node(0)),
+            // A tree with only one node, it must be a leaf or unvisited node.
+            // It can't be a branch because branches have children.
+            ([], [_], []) | ([], [], [_]) => Ok(TrieRoot::Node(0)),
             (branches, _, _) if !branches.is_empty() => {
                 Ok(TrieRoot::Node(branches.len() as Idx - 1))
             }

From 93b4efb80c93f1eb026255d0b3b9e616ccf80183 Mon Sep 17 00:00:00 2001
From: Avi Dessauer <avi.the.coder@gmail.com>
Date: Fri, 16 Aug 2024 19:10:08 -0700
Subject: [PATCH 2/3] Change many things

---
 examples/prove-and-verify.rs |   8 +-
 src/stored.rs                |  34 ++++++---
 src/stored/merkle.rs         | 138 ++++++++++++++++++++++++-----------
 src/transaction.rs           | 106 +++++++++++++++++++--------
 tests/utils/insert_get.rs    |  15 ++--
 tests/utils/operations.rs    |  15 +++-
 6 files changed, 218 insertions(+), 98 deletions(-)

diff --git a/examples/prove-and-verify.rs b/examples/prove-and-verify.rs
index 6800c9f..aa025ff 100644
--- a/examples/prove-and-verify.rs
+++ b/examples/prove-and-verify.rs
@@ -3,7 +3,7 @@ use std::rc::Rc;
 use kairos_trie::{
     stored::{
         memory_db::MemoryDb,
-        merkle::{Snapshot, SnapshotBuilder},
+        merkle::{Snapshot, SnapshotBuilder, VerifiedSnapshot},
         Store,
     },
     DigestHasher, KeyHash, NodeHash, PortableHash, PortableHasher, Transaction, TrieRoot,
@@ -21,7 +21,7 @@ fn hash(key: &str) -> KeyHash {
     KeyHash::from_bytes(&hasher.finalize_reset())
 }
 
-fn apply_operations(txn: &mut Transaction<impl Store<u64>, u64>, operations: &[Ops]) {
+fn apply_operations(txn: &mut Transaction<impl Store<Value = u64>>, operations: &[Ops]) {
     for op in operations {
         match op {
             Ops::Add(key, value) => {
@@ -72,7 +72,9 @@ fn verifier(
 ) -> TrieRoot<NodeHash> {
     let hasher = &mut DigestHasher::<Sha256>::default();
 
-    let mut txn = Transaction::from_snapshot(snapshot).unwrap();
+    let mut txn = Transaction::from_verified_snapshot(
+        VerifiedSnapshot::verify_snapshot(snapshot, hasher).unwrap(),
+    );
 
     let pre_batch_trie_root = txn.calc_root_hash(hasher).unwrap();
     // Assert that the trie started the transaction with the correct root hash.
diff --git a/src/stored.rs b/src/stored.rs
index 0ea1bd6..490b3c2 100644
--- a/src/stored.rs
+++ b/src/stored.rs
@@ -7,13 +7,14 @@ use alloc::{rc::Rc, sync::Arc};
 
 use crate::{
     transaction::nodes::{Branch, Leaf, Node},
-    NodeHash, PortableHasher,
+    NodeHash, PortableHash, PortableHasher,
 };
 
 pub type Idx = u32;
 
-pub trait Store<V> {
+pub trait Store {
     type Error: Display;
+    type Value: Clone + PortableHash;
 
     fn calc_subtree_hash(
         &self,
@@ -21,11 +22,15 @@ pub trait Store<V> {
         hash_idx: Idx,
     ) -> Result<NodeHash, Self::Error>;
 
-    fn get_node(&self, hash_idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>, Self::Error>;
+    fn get_node(
+        &self,
+        hash_idx: Idx,
+    ) -> Result<Node<&Branch<Idx>, &Leaf<Self::Value>>, Self::Error>;
 }
 
-impl<V, S: Store<V>> Store<V> for &S {
+impl<S: Store> Store for &S {
     type Error = S::Error;
+    type Value = S::Value;
 
     #[inline(always)]
     fn calc_subtree_hash(
@@ -38,13 +43,17 @@ impl<V, S: Store<V>> Store<V> for &S {
     }
 
     #[inline(always)]
-    fn get_node(&self, hash_idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>, Self::Error> {
+    fn get_node(
+        &self,
+        hash_idx: Idx,
+    ) -> Result<Node<&Branch<Idx>, &Leaf<Self::Value>>, Self::Error> {
         (**self).get_node(hash_idx)
     }
 }
 
-impl<V, S: Store<V>> Store<V> for Rc<S> {
+impl<S: Store> Store for Rc<S> {
     type Error = S::Error;
+    type Value = S::Value;
 
     #[inline(always)]
     fn calc_subtree_hash(
@@ -56,13 +65,17 @@ impl<V, S: Store<V>> Store<V> for Rc<S> {
     }
 
     #[inline(always)]
-    fn get_node(&self, hash_idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>, Self::Error> {
+    fn get_node(
+        &self,
+        hash_idx: Idx,
+    ) -> Result<Node<&Branch<Idx>, &Leaf<Self::Value>>, Self::Error> {
         (**self).get_node(hash_idx)
     }
 }
 
-impl<V, S: Store<V>> Store<V> for Arc<S> {
+impl<S: Store> Store for Arc<S> {
     type Error = S::Error;
+    type Value = S::Value;
 
     #[inline(always)]
     fn calc_subtree_hash(
@@ -74,7 +87,10 @@ impl<V, S: Store<V>> Store<V> for Arc<S> {
     }
 
     #[inline(always)]
-    fn get_node(&self, hash_idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>, Self::Error> {
+    fn get_node(
+        &self,
+        hash_idx: Idx,
+    ) -> Result<Node<&Branch<Idx>, &Leaf<Self::Value>>, Self::Error> {
         (**self).get_node(hash_idx)
     }
 }
diff --git a/src/stored/merkle.rs b/src/stored/merkle.rs
index ac52713..91178fc 100644
--- a/src/stored/merkle.rs
+++ b/src/stored/merkle.rs
@@ -13,8 +13,13 @@ use super::{DatabaseGet, Idx, Node, NodeHash, Store};
 
 type Result<T, E = TrieError> = core::result::Result<T, E>;
 
-pub struct VerifiedSnapshot<V> {
-    snapshot: Snapshot<V>,
+/// A snapshot of the merkle trie verified
+///
+/// Contains visited nodes and unvisited nodes
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct VerifiedSnapshot<S: Store> {
+    snapshot: S,
 
     /// The root hash of the snapshot is the last hash in the slice.
     /// The indexes of each hash match the indexes of nodes in the snapshot.
@@ -22,26 +27,32 @@ pub struct VerifiedSnapshot<V> {
     leaf_hashes: Box<[NodeHash]>,
 }
 
-impl<V: PortableHash> VerifiedSnapshot<V> {
+impl<S: Store + AsRef<Snapshot<S::Value>>> VerifiedSnapshot<S> {
+    /// Verify the snapshot by checking that it is well formed and calculating the merkle hashes of all nodes.
+    /// The merkle hashes are cached such that `calc_subtree_hash` is an O(1) operation for all nodes in the snapshot.
+    /// In practice this means if a transaction gets, but does not modify a node the hash is never recalculated.
+    ///
+    /// Storing the hashes of all nodes does increase memory usage,
+    /// so for some use cases it may be better to use `Snapshot` directly.
+    /// If you choose to do that, be sure to run `calc_root_hash` before using the snapshot.
     #[inline]
-    pub fn verify_snapshot(
-        snapshot: Snapshot<V>,
-        hasher: &mut impl PortableHasher<32>,
-    ) -> Result<Self> {
+    pub fn verify_snapshot(snapshot: S, hasher: &mut impl PortableHasher<32>) -> Result<Self> {
+        let snapshot_ref = snapshot.as_ref();
+
         // Check that the snapshot is well formed.
-        let _ = snapshot.root_node_idx()?;
+        let _ = snapshot_ref.root_node_idx()?;
 
-        let mut leaf_hashes = Vec::with_capacity(snapshot.leaves.len());
-        let mut branch_hashes = Vec::with_capacity(snapshot.branches.len());
+        let mut leaf_hashes = Vec::with_capacity(snapshot_ref.leaves.len());
+        let mut branch_hashes = Vec::with_capacity(snapshot_ref.branches.len());
 
-        for leaf in snapshot.leaves.iter() {
+        for leaf in snapshot_ref.leaves.iter() {
             leaf_hashes.push(leaf.hash_leaf(hasher));
         }
 
-        let leaf_offset = snapshot.branches.len();
-        let unvisited_offset = leaf_offset + snapshot.leaves.len();
+        let leaf_offset = snapshot_ref.branches.len();
+        let unvisited_offset = leaf_offset + snapshot_ref.leaves.len();
 
-        for (idx, branch) in snapshot.branches.iter().enumerate() {
+        for (idx, branch) in snapshot_ref.branches.iter().enumerate() {
             let hash_of_child = |child| {
                 if child < idx {
                     branch_hashes.get(child).ok_or_else(|| {
@@ -60,7 +71,7 @@ impl<V: PortableHash> VerifiedSnapshot<V> {
                         )
                     })
                 } else {
-                    snapshot
+                    snapshot_ref
                         .unvisited_nodes
                         .get(child - unvisited_offset)
                         .ok_or_else(|| {
@@ -87,24 +98,40 @@ impl<V: PortableHash> VerifiedSnapshot<V> {
     }
 
     #[inline]
-    pub fn trie_root(&self) -> TrieRoot<NodeRef<V>> {
-        if !self.snapshot.branches.is_empty() {
-            TrieRoot::Node(NodeRef::Stored(self.snapshot.branches.len() as Idx - 1))
+    pub fn trie_root(&self) -> TrieRoot<NodeRef<S::Value>> {
+        let snapshot = self.snapshot.as_ref();
+
+        if !snapshot.branches.is_empty() {
+            TrieRoot::Node(NodeRef::Stored(snapshot.branches.len() as Idx - 1))
+        } else if snapshot.leaves.is_empty() && snapshot.unvisited_nodes.is_empty() {
+            TrieRoot::Empty
         } else {
             // We know that the snapshot is valid, because we verified it in `verify_snapshot`.
-            // If a snapshot contains no branches a single leaf or a single unvisited node is a valid snapshot.
+            // If a non-empty snapshot contains no branches, it must have a single leaf or unvisited node.
             // Any other combination is invalid.
-            debug_assert_eq!(
-                self.snapshot.leaves.len() + self.snapshot.unvisited_nodes.len(),
-                1
-            );
+            debug_assert_eq!(snapshot.leaves.len() + snapshot.unvisited_nodes.len(), 1);
             TrieRoot::Node(NodeRef::Stored(0))
         }
     }
+
+    /// Returns the merkle root hash of the trie in the snapshot.
+    /// The hash of all nodes has already been calculated in `VerifiedSnapshot::verify_snapshot`.
+    /// `trie_root_hash` and `calc_root_hash` are both O(1) operations on a `VerifiedSnapshot`,
+    /// unlike `Snapshot` and `SnapshotBuilder`.
+    #[inline]
+    pub fn trie_root_hash(&self) -> TrieRoot<NodeHash> {
+        self.branch_hashes
+            .last()
+            // Given a valid snapshot: if no branches exist, there can only be one leaf or one unvisited node.
+            .or_else(|| self.leaf_hashes.first())
+            .or_else(|| self.snapshot.as_ref().unvisited_nodes.first())
+            .map_or(TrieRoot::Empty, |hash| TrieRoot::Node(*hash))
+    }
 }
 
-impl<V> Store<V> for VerifiedSnapshot<V> {
+impl<S: Store + AsRef<Snapshot<S::Value>>> Store for VerifiedSnapshot<S> {
     type Error = TrieError;
+    type Value = S::Value;
 
     #[inline]
     fn calc_subtree_hash(
@@ -112,41 +139,42 @@ impl<V> Store<V> for VerifiedSnapshot<V> {
         _: &mut impl PortableHasher<32>,
         node: Idx,
     ) -> Result<NodeHash, Self::Error> {
+        let snapshot = self.snapshot.as_ref();
+
         let idx = node as usize;
-        let leaf_offset = self.snapshot.branches.len();
-        let unvisited_offset = leaf_offset + self.snapshot.leaves.len();
+        let leaf_offset = snapshot.branches.len();
+        let unvisited_offset = leaf_offset + snapshot.leaves.len();
 
         if let Some(branch) = self.branch_hashes.get(idx) {
             Ok(*branch)
         } else if let Some(leaf) = self.leaf_hashes.get(idx - leaf_offset) {
             Ok(*leaf)
-        } else if let Some(hash) = self.snapshot.unvisited_nodes.get(idx - unvisited_offset) {
+        } else if let Some(hash) = snapshot.unvisited_nodes.get(idx - unvisited_offset) {
             Ok(*hash)
         } else {
             Err(format!(
                 "Invalid arg: node {} does not exist\n\
                 Snapshot has {} nodes",
                 idx,
-                self.snapshot.branches.len()
-                    + self.snapshot.leaves.len()
-                    + self.snapshot.unvisited_nodes.len(),
+                snapshot.branches.len() + snapshot.leaves.len() + snapshot.unvisited_nodes.len(),
             )
             .into())
         }
     }
 
     #[inline]
-    fn get_node(&self, idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>> {
+    fn get_node(&self, idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<S::Value>>> {
+        let snapshot = self.snapshot.as_ref();
+
         let idx = idx as usize;
-        let leaf_offset = self.snapshot.branches.len();
-        let unvisited_offset = leaf_offset + self.snapshot.leaves.len();
+        let leaf_offset = snapshot.branches.len();
+        let unvisited_offset = leaf_offset + snapshot.leaves.len();
 
-        if let Some(branch) = self.snapshot.branches.get(idx) {
+        if let Some(branch) = snapshot.branches.get(idx) {
             Ok(Node::Branch(branch))
-        } else if let Some(leaf) = self.snapshot.leaves.get(idx - leaf_offset) {
+        } else if let Some(leaf) = snapshot.leaves.get(idx - leaf_offset) {
             Ok(Node::Leaf(leaf))
-        } else if self
-            .snapshot
+        } else if snapshot
             .unvisited_nodes
             .get(idx - unvisited_offset)
             .is_some()
@@ -161,9 +189,7 @@ impl<V> Store<V> for VerifiedSnapshot<V> {
                 "Invalid arg: node {} does not exist\n\
                 Snapshot has {} nodes",
                 idx,
-                self.snapshot.branches.len()
-                    + self.snapshot.leaves.len()
-                    + self.snapshot.unvisited_nodes.len(),
+                snapshot.branches.len() + snapshot.leaves.len() + snapshot.unvisited_nodes.len(),
             )
             .into())
         }
@@ -185,7 +211,14 @@ pub struct Snapshot<V> {
     unvisited_nodes: Box<[NodeHash]>,
 }
 
-impl<V: PortableHash> Snapshot<V> {
+impl<V> AsRef<Snapshot<V>> for Snapshot<V> {
+    #[inline]
+    fn as_ref(&self) -> &Snapshot<V> {
+        self
+    }
+}
+
+impl<V: Clone + PortableHash> Snapshot<V> {
     #[inline]
     pub fn root_node_idx(&self) -> Result<TrieRoot<Idx>> {
         // Revist this once https://github.com/rust-lang/rust/issues/37854 is stable
@@ -239,10 +272,28 @@ impl<V: PortableHash> Snapshot<V> {
             TrieRoot::Empty => Ok(TrieRoot::Empty),
         }
     }
+
+    /// Verify the snapshot by checking that it is well formed and calculating the merkle hashes of all nodes.
+    /// This is an alias for `VerifiedSnapshot::verify_snapshot`.
+    #[inline]
+    pub fn verify_ref(
+        &self,
+        hasher: &mut impl PortableHasher<32>,
+    ) -> Result<VerifiedSnapshot<&Self>> {
+        VerifiedSnapshot::verify_snapshot(self, hasher)
+    }
+
+    /// Verify the snapshot by checking that it is well formed and calculating the merkle hashes of all nodes.
+    /// This is an alias for `VerifiedSnapshot::verify_snapshot`.
+    #[inline]
+    pub fn verify(self, hasher: &mut impl PortableHasher<32>) -> Result<VerifiedSnapshot<Self>> {
+        VerifiedSnapshot::verify_snapshot(self, hasher)
+    }
 }
 
-impl<V: PortableHash> Store<V> for Snapshot<V> {
+impl<V: Clone + PortableHash> Store for Snapshot<V> {
     type Error = TrieError;
+    type Value = V;
 
     // TODO fix possible stack overflow
     // I dislike using an explicit mutable stack.
@@ -324,8 +375,9 @@ struct SnapshotBuilderInner<Db: 'static, V: 'static> {
     nodes: RefCell<Vec<NodeHashMaybeNode<'this, V>>>,
 }
 
-impl<Db: DatabaseGet<V>, V: Clone> Store<V> for SnapshotBuilder<Db, V> {
+impl<Db: DatabaseGet<V>, V: Clone + PortableHash> Store for SnapshotBuilder<Db, V> {
     type Error = TrieError;
+    type Value = V;
 
     #[inline]
     fn calc_subtree_hash(
diff --git a/src/transaction.rs b/src/transaction.rs
index 45087c8..21956a2 100644
--- a/src/transaction.rs
+++ b/src/transaction.rs
@@ -4,6 +4,7 @@ use alloc::borrow::Cow;
 use alloc::{boxed::Box, format};
 use core::mem;
 
+use crate::stored::merkle::VerifiedSnapshot;
 use crate::stored::DatabaseGet;
 use crate::{stored, KeyHash, NodeHash, PortableHash, PortableHasher};
 use crate::{
@@ -18,12 +19,12 @@ use self::nodes::{
     Branch, KeyPosition, KeyPositionAdjacent, Leaf, Node, NodeRef, StoredLeafRef, TrieRoot,
 };
 
-pub struct Transaction<S, V> {
+pub struct Transaction<S: Store> {
     pub data_store: S,
-    current_root: TrieRoot<NodeRef<V>>,
+    current_root: TrieRoot<NodeRef<S::Value>>,
 }
 
-impl<Db: DatabaseSet<V>, V: Clone + PortableHash> Transaction<SnapshotBuilder<Db, V>, V> {
+impl<Db: DatabaseSet<V>, V: Clone + PortableHash> Transaction<SnapshotBuilder<Db, V>> {
     /// Write modified nodes to the database and return the root hash.
     /// Calling this method will write all modified nodes to the database.
     /// Calling this method again will rewrite the nodes to the database.
@@ -65,7 +66,7 @@ impl<Db: DatabaseSet<V>, V: Clone + PortableHash> Transaction<SnapshotBuilder<Db
     }
 }
 
-impl<S: Store<V>, V: PortableHash> Transaction<S, V> {
+impl<S: Store> Transaction<S> {
     /// Caller must ensure that the hasher is reset before calling this method.
     #[inline]
     pub fn calc_root_hash_inner(
@@ -73,11 +74,11 @@ impl<S: Store<V>, V: PortableHash> Transaction<S, V> {
         hasher: &mut impl PortableHasher<32>,
         on_modified_branch: &mut impl FnMut(
             &NodeHash,
-            &Branch<NodeRef<V>>,
+            &Branch<NodeRef<S::Value>>,
             NodeHash,
             NodeHash,
         ) -> Result<(), TrieError>,
-        on_modified_leaf: &mut impl FnMut(&NodeHash, &Leaf<V>) -> Result<(), TrieError>,
+        on_modified_leaf: &mut impl FnMut(&NodeHash, &Leaf<S::Value>) -> Result<(), TrieError>,
     ) -> Result<TrieRoot<NodeHash>, TrieError> {
         let root_hash = match &self.current_root {
             TrieRoot::Empty => return Ok(TrieRoot::Empty),
@@ -108,11 +109,11 @@ impl<S: Store<V>, V: PortableHash> Transaction<S, V> {
     fn calc_root_hash_node(
         hasher: &mut impl PortableHasher<32>,
         data_store: &S,
-        node_ref: &NodeRef<V>,
-        on_modified_leaf: &mut impl FnMut(&NodeHash, &Leaf<V>) -> Result<(), TrieError>,
+        node_ref: &NodeRef<S::Value>,
+        on_modified_leaf: &mut impl FnMut(&NodeHash, &Leaf<S::Value>) -> Result<(), TrieError>,
         on_modified_branch: &mut impl FnMut(
             &NodeHash,
-            &Branch<NodeRef<V>>,
+            &Branch<NodeRef<S::Value>>,
             NodeHash,
             NodeHash,
         ) -> Result<(), TrieError>,
@@ -160,7 +161,7 @@ impl<S: Store<V>, V: PortableHash> Transaction<S, V> {
     }
 }
 
-impl<Db: 'static + DatabaseGet<V>, V: Clone> Transaction<SnapshotBuilder<Db, V>, V> {
+impl<Db: 'static + DatabaseGet<V>, V: Clone + PortableHash> Transaction<SnapshotBuilder<Db, V>> {
     /// This method is like standard `Transaction::get` but won't affect the Transaction or any Snapshot built from it.
     /// You should use this method to check precondition before modifying the Transaction.
     ///
@@ -244,9 +245,9 @@ impl<Db: 'static + DatabaseGet<V>, V: Clone> Transaction<SnapshotBuilder<Db, V>,
     }
 }
 
-impl<S: Store<V>, V> Transaction<S, V> {
+impl<S: Store> Transaction<S> {
     #[inline]
-    pub fn get(&self, key_hash: &KeyHash) -> Result<Option<&V>, TrieError> {
+    pub fn get(&self, key_hash: &KeyHash) -> Result<Option<&S::Value>, TrieError> {
         match &self.current_root {
             TrieRoot::Empty => Ok(None),
             TrieRoot::Node(node_ref) => Self::get_node(&self.data_store, node_ref, key_hash),
@@ -256,9 +257,9 @@ impl<S: Store<V>, V> Transaction<S, V> {
     #[inline]
     fn get_node<'root, 's: 'root>(
         data_store: &'s S,
-        mut node_ref: &'root NodeRef<V>,
+        mut node_ref: &'root NodeRef<S::Value>,
         key_hash: &KeyHash,
-    ) -> Result<Option<&'root V>, TrieError> {
+    ) -> Result<Option<&'root S::Value>, TrieError> {
         loop {
             match node_ref {
                 NodeRef::ModBranch(branch) => match branch.key_position(key_hash) {
@@ -285,7 +286,7 @@ impl<S: Store<V>, V> Transaction<S, V> {
         data_store: &'s S,
         mut stored_idx: stored::Idx,
         key_hash: &KeyHash,
-    ) -> Result<Option<&'s V>, TrieError> {
+    ) -> Result<Option<&'s S::Value>, TrieError> {
         loop {
             let node = data_store
                 .get_node(stored_idx)
@@ -317,7 +318,7 @@ impl<S: Store<V>, V> Transaction<S, V> {
     }
 
     #[inline]
-    pub fn insert(&mut self, key_hash: &KeyHash, value: V) -> Result<(), TrieError> {
+    pub fn insert(&mut self, key_hash: &KeyHash, value: S::Value) -> Result<(), TrieError> {
         match &mut self.current_root {
             TrieRoot::Empty => {
                 self.current_root = TrieRoot::Node(NodeRef::ModLeaf(Box::new(Leaf {
@@ -335,9 +336,9 @@ impl<S: Store<V>, V> Transaction<S, V> {
     #[inline(always)]
     fn insert_node<'root, 's: 'root>(
         data_store: &'s mut S,
-        mut node_ref: &'root mut NodeRef<V>,
+        mut node_ref: &'root mut NodeRef<S::Value>,
         key_hash: &KeyHash,
-        value: V,
+        value: S::Value,
     ) -> Result<(), TrieError> {
         loop {
             match node_ref {
@@ -430,7 +431,7 @@ impl<S: Store<V>, V> Transaction<S, V> {
     }
 }
 
-impl<S: Store<V>, V: PortableHash + Clone> Transaction<S, V> {
+impl<S: Store> Transaction<S> {
     /// This method allows for getting, inserting, and updating a entry in the trie with a single lookup.
     /// We match the standard library's `Entry` API for the most part.
     ///
@@ -438,7 +439,10 @@ impl<S: Store<V>, V: PortableHash + Clone> Transaction<S, V> {
     /// This incurs allocations, now and unnecessary rehashing later when calculating the root hash.
     /// For this reason you should prefer `get` if you have a high probability of not modifying the entry.
     #[inline]
-    pub fn entry<'txn>(&'txn mut self, key_hash: &KeyHash) -> Result<Entry<'txn, V>, TrieError> {
+    pub fn entry<'txn>(
+        &'txn mut self,
+        key_hash: &KeyHash,
+    ) -> Result<Entry<'txn, S::Value>, TrieError> {
         let mut key_position = KeyPositionAdjacent::PrefixOfWord(usize::MAX);
 
         match self.current_root {
@@ -528,7 +532,7 @@ impl<S: Store<V>, V: PortableHash + Clone> Transaction<S, V> {
     }
 }
 
-impl<Db, V: PortableHash + Clone> Transaction<SnapshotBuilder<Db, V>, V> {
+impl<Db: DatabaseGet<V>, V: PortableHash + Clone> Transaction<SnapshotBuilder<Db, V>> {
     /// An alias for `SnapshotBuilder::new_with_db`.
     ///
     /// Builds a snapshot of the trie before the transaction.
@@ -551,8 +555,8 @@ impl<Db, V: PortableHash + Clone> Transaction<SnapshotBuilder<Db, V>, V> {
     }
 }
 
-impl<Db, V: PortableHash + Clone> TryFrom<SnapshotBuilder<Db, V>>
-    for Transaction<SnapshotBuilder<Db, V>, V>
+impl<Db: DatabaseGet<V>, V: PortableHash + Clone> TryFrom<SnapshotBuilder<Db, V>>
+    for Transaction<SnapshotBuilder<Db, V>>
 {
     type Error = TrieError;
 
@@ -562,10 +566,10 @@ impl<Db, V: PortableHash + Clone> TryFrom<SnapshotBuilder<Db, V>>
     }
 }
 
-impl<'s, V: PortableHash + Clone> Transaction<&'s Snapshot<V>, V> {
+impl<'s, V: PortableHash + Clone> Transaction<&'s Snapshot<V>> {
     /// Create a `Transaction` from a borrowed `Snapshot`.
     #[inline]
-    pub fn from_snapshot(snapshot: &'s Snapshot<V>) -> Result<Self, TrieError> {
+    pub fn from_unverified_snapshot_ref(snapshot: &'s Snapshot<V>) -> Result<Self, TrieError> {
         Ok(Transaction {
             current_root: snapshot.trie_root()?,
             data_store: snapshot,
@@ -573,10 +577,10 @@ impl<'s, V: PortableHash + Clone> Transaction<&'s Snapshot<V>, V> {
     }
 }
 
-impl<V: PortableHash + Clone> Transaction<Snapshot<V>, V> {
+impl<V: PortableHash + Clone> Transaction<Snapshot<V>> {
     /// Create a `Transaction` from a owned `Snapshot`.
     #[inline]
-    pub fn from_snapshot_owned(snapshot: Snapshot<V>) -> Result<Self, TrieError> {
+    pub fn from_unverified_snapshot(snapshot: Snapshot<V>) -> Result<Self, TrieError> {
         Ok(Transaction {
             current_root: snapshot.trie_root()?,
             data_store: snapshot,
@@ -584,21 +588,61 @@ impl<V: PortableHash + Clone> Transaction<Snapshot<V>, V> {
     }
 }
 
-impl<'s, V: PortableHash + Clone> TryFrom<&'s Snapshot<V>> for Transaction<&'s Snapshot<V>, V> {
+impl<'s, V: PortableHash + Clone> TryFrom<&'s Snapshot<V>> for Transaction<&'s Snapshot<V>> {
     type Error = TrieError;
 
     #[inline]
     fn try_from(value: &'s Snapshot<V>) -> Result<Self, Self::Error> {
-        Self::from_snapshot(value)
+        Self::from_unverified_snapshot_ref(value)
     }
 }
 
-impl<V: PortableHash + Clone> TryFrom<Snapshot<V>> for Transaction<Snapshot<V>, V> {
+impl<V: PortableHash + Clone> TryFrom<Snapshot<V>> for Transaction<Snapshot<V>> {
     type Error = TrieError;
 
     #[inline]
     fn try_from(value: Snapshot<V>) -> Result<Self, Self::Error> {
-        Self::from_snapshot_owned(value)
+        Self::from_unverified_snapshot(value)
+    }
+}
+
+impl<'s, S: Store + AsRef<Snapshot<S::Value>>> Transaction<&'s VerifiedSnapshot<S>> {
+    /// Create a `Transaction` from a borrowed `VerifiedSnapshot`.
+    #[inline]
+    pub fn from_verified_snapshot_ref(snapshot: &'s VerifiedSnapshot<S>) -> Self {
+        Transaction {
+            current_root: snapshot.trie_root(),
+            data_store: snapshot,
+        }
+    }
+}
+
+impl<S: Store + AsRef<Snapshot<S::Value>>> Transaction<VerifiedSnapshot<S>> {
+    /// Create a `Transaction` from a owned `VerifiedSnapshot`.
+    #[inline]
+    pub fn from_verified_snapshot(snapshot: VerifiedSnapshot<S>) -> Self {
+        Transaction {
+            current_root: snapshot.trie_root(),
+            data_store: snapshot,
+        }
+    }
+}
+
+impl<'s, S: Store + AsRef<Snapshot<S::Value>>> From<&'s VerifiedSnapshot<S>>
+    for Transaction<&'s VerifiedSnapshot<S>>
+{
+    #[inline]
+    fn from(value: &'s VerifiedSnapshot<S>) -> Self {
+        Transaction::from_verified_snapshot_ref(value)
+    }
+}
+
+impl<S: Store + AsRef<Snapshot<S::Value>>> From<VerifiedSnapshot<S>>
+    for Transaction<VerifiedSnapshot<S>>
+{
+    #[inline]
+    fn from(value: VerifiedSnapshot<S>) -> Self {
+        Transaction::from_verified_snapshot(value)
     }
 }
 
diff --git a/tests/utils/insert_get.rs b/tests/utils/insert_get.rs
index 9a5b494..05ef548 100644
--- a/tests/utils/insert_get.rs
+++ b/tests/utils/insert_get.rs
@@ -4,7 +4,7 @@ use std::collections::HashMap;
 
 use kairos_trie::{
     stored::{
-        merkle::{Snapshot, SnapshotBuilder},
+        merkle::{Snapshot, SnapshotBuilder, VerifiedSnapshot},
         DatabaseSet,
     },
     DigestHasher, KeyHash, NodeHash, Transaction, TrieRoot,
@@ -37,14 +37,13 @@ pub fn run_against_snapshot(
     new_root_hash: TrieRoot<NodeHash>,
     old_root_hash: TrieRoot<NodeHash>,
 ) {
-    assert_eq!(
-        old_root_hash,
-        snapshot
-            .calc_root_hash(&mut DigestHasher::<Sha256>::default())
-            .unwrap()
-    );
+    let mut hasher = &mut DigestHasher::<Sha256>::default();
+
+    assert_eq!(old_root_hash, snapshot.calc_root_hash(hasher).unwrap());
 
-    let mut txn = Transaction::from_snapshot(&snapshot).unwrap();
+    let mut txn = Transaction::from_verified_snapshot(
+        VerifiedSnapshot::verify_snapshot(snapshot, hasher).unwrap(),
+    );
 
     for (key, value) in new.iter() {
         txn.insert(key, value.to_le_bytes()).unwrap();
diff --git a/tests/utils/operations.rs b/tests/utils/operations.rs
index 3d294af..4e2ea62 100644
--- a/tests/utils/operations.rs
+++ b/tests/utils/operations.rs
@@ -10,7 +10,7 @@ use proptest::{prelude::*, sample::SizeRange};
 use kairos_trie::{
     stored::{
         memory_db::MemoryDb,
-        merkle::{Snapshot, SnapshotBuilder},
+        merkle::{Snapshot, SnapshotBuilder, VerifiedSnapshot},
         Store,
     },
     DigestHasher, KeyHash, NodeHash, Transaction, TrieRoot,
@@ -132,8 +132,15 @@ pub fn run_against_snapshot(
             .unwrap()
     );
 
+    let verified_snapshot =
+        VerifiedSnapshot::verify_snapshot(&snapshot, &mut DigestHasher::<Sha256>::default())
+            .unwrap();
+
+    assert_eq!(old_root_hash, verified_snapshot.trie_root_hash());
+
     // Create a transaction against the snapshot at the old root hash
-    let mut txn = Transaction::from_snapshot(&snapshot).unwrap();
+    // let mut txn = Transaction::from_verified_snapshot(verified_snapshot);
+    let mut txn = Transaction::from_unverified_snapshot(snapshot).unwrap();
 
     // Apply the operations to the transaction
     for op in batch {
@@ -150,9 +157,9 @@ pub fn run_against_snapshot(
     assert_eq!(root_hash, new_root_hash);
 }
 
-fn trie_op<S: Store<Value>>(
+fn trie_op<S: Store<Value = [u8; 8]>>(
     op: &Operation,
-    txn: &mut Transaction<S, Value>,
+    txn: &mut Transaction<S>,
 ) -> (Option<Value>, Option<Value>) {
     match op {
         Operation::Insert(key, value) => {

From 2a6c733eebdb77820940697c908ebd4191b8c481 Mon Sep 17 00:00:00 2001
From: Avi Dessauer <avi.the.coder@gmail.com>
Date: Fri, 16 Aug 2024 20:00:39 -0700
Subject: [PATCH 3/3] Fix tests

---
 src/stored/merkle.rs | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/stored/merkle.rs b/src/stored/merkle.rs
index 91178fc..b0dceb6 100644
--- a/src/stored/merkle.rs
+++ b/src/stored/merkle.rs
@@ -54,18 +54,18 @@ impl<S: Store + AsRef<Snapshot<S::Value>>> VerifiedSnapshot<S> {
 
         for (idx, branch) in snapshot_ref.branches.iter().enumerate() {
             let hash_of_child = |child| {
-                if child < idx {
+                if child < leaf_offset {
                     branch_hashes.get(child).ok_or_else(|| {
                         format!(
-                            "Invalid snapshot: branch {} has child {},
+                            "Invalid snapshot: branch {} has child {},\
                             child branch index must be less than parent, branches are not in post-order traversal",
                             idx, child
                         )
                     })
-                } else if child < leaf_offset {
-                    leaf_hashes.get(child).ok_or_else(|| {
+                } else if child < unvisited_offset {
+                    leaf_hashes.get(child - leaf_offset).ok_or_else(|| {
                         format!(
-                            "Invalid snapshot: branch {} has child {},
+                            "Invalid snapshot: branch {} has child {},\
                             child leaf does not exist",
                             idx, child
                         )
@@ -76,7 +76,7 @@ impl<S: Store + AsRef<Snapshot<S::Value>>> VerifiedSnapshot<S> {
                         .get(child - unvisited_offset)
                         .ok_or_else(|| {
                             format!(
-                                "Invalid snapshot: branch {} has child {},
+                                "Invalid snapshot: branch {} has child {},\
                                 child unvisited node does not exist",
                                 idx, child
                             )