Store to database

src/lib.rs

@@ -17,7 +17,7 @@ use sha2::{Digest, Sha256};
 pub use stored::Store;
 use stored::{
     merkle::{Snapshot, SnapshotBuilder},
-    Node, NodeHash,
+    DatabaseSet, Node, NodeHash,
 };
 
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
@@ -258,7 +258,7 @@ impl<NR> Branch<NR> {
     }
 }
 
-impl<V: Debug> Branch<NodeRef<V>> {
+impl<V> Branch<NodeRef<V>> {
     fn new_at_branch(
         word_idx: usize,
         branch_word_or_prefix: u32,
@@ -424,7 +424,42 @@ pub struct Transaction<S, V> {
     pub current_root: TrieRoot<V>,
 }
 
-impl<S: Store<V>, V: Debug + AsRef<[u8]>> Transaction<S, V> {
+impl<'a, Db: DatabaseSet<V>, V: Clone + AsRef<[u8]>> Transaction<SnapshotBuilder<'a, Db, V>, 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.
+    ///
+    /// Caching writes is the responsibility of the `DatabaseSet` implementation.
+    pub fn commit(&self) -> Result<NodeHash, String> {
+        let store_modified_branch =
+            &mut |hash: &NodeHash, branch: &Branch<NodeRef<V>>, left: NodeHash, right: NodeHash| {
+                let branch = Branch {
+                    left,
+                    right,
+                    mask: branch.mask,
+                    prior_word: branch.prior_word,
+                    prefix: branch.prefix.clone(),
+                };
+
+                self.data_store
+                    .db
+                    .set(*hash, Node::Branch(branch))
+                    .map_err(|e| e.into())
+            };
+
+        let store_modified_leaf = &mut |hash: &NodeHash, leaf: &Leaf<V>| {
+            self.data_store
+                .db
+                .set(*hash, Node::Leaf(leaf.clone()))
+                .map_err(|e| e.into())
+        };
+
+        let root_hash = self.calc_root_hash_inner(store_modified_branch, store_modified_leaf)?;
+        Ok(root_hash)
+    }
+}
+
+impl<S: Store<V>, V: AsRef<[u8]>> Transaction<S, V> {
     pub fn new(root: TrieRoot<V>, data_store: S) -> Self {
         Transaction {
             current_root: root,
@@ -433,29 +468,44 @@ impl<S: Store<V>, V: Debug + AsRef<[u8]>> Transaction<S, V> {
     }
 
     /// TODO a version of this that writes to the database.
-    pub fn calc_root_hash(&self) -> Result<NodeHash, String> {
-        let mut on_modified_leaf = |_: &_, _: &_| {};
-        let mut on_modified_branch = |_: &_, _: &_| {};
-
+    pub fn calc_root_hash_inner(
+        &self,
+        on_modified_branch: &mut impl FnMut(
+            &NodeHash,
+            &Branch<NodeRef<V>>,
+            NodeHash,
+            NodeHash,
+        ) -> Result<(), String>,
+        on_modified_leaf: &mut impl FnMut(&NodeHash, &Leaf<V>) -> Result<(), String>,
+    ) -> Result<NodeHash, String> {
         let root_hash = match &self.current_root {
             TrieRoot::Empty => return Ok([0; 32]),
             TrieRoot::Node(node_ref) => Self::calc_root_hash_node(
                 &self.data_store,
                 node_ref,
-                &mut on_modified_leaf,
-                &mut on_modified_branch,
+                on_modified_leaf,
+                on_modified_branch,
             )?,
         };
 
         Ok(root_hash)
     }
 
+    pub fn calc_root_hash(&self) -> Result<NodeHash, String> {
+        self.calc_root_hash_inner(&mut |_, _, _, _| Ok(()), &mut |_, _| Ok(()))
+    }
+
     /// TODO use this to store nodes in the data base
     fn calc_root_hash_node(
         data_store: &S,
         node_ref: &NodeRef<V>,
-        on_modified_leaf: &mut impl FnMut(&NodeHash, &Leaf<V>),
-        on_modified_branch: &mut impl FnMut(&NodeHash, &Branch<NodeRef<V>>),
+        on_modified_leaf: &mut impl FnMut(&NodeHash, &Leaf<V>) -> Result<(), String>,
+        on_modified_branch: &mut impl FnMut(
+            &NodeHash,
+            &Branch<NodeRef<V>>,
+            NodeHash,
+            NodeHash,
+        ) -> Result<(), String>,
     ) -> Result<NodeHash, String> {
         // TODO use a stack instead of recursion
         match node_ref {
@@ -474,13 +524,13 @@ impl<S: Store<V>, V: Debug + AsRef<[u8]>> Transaction<S, V> {
                 )?;
 
                 let hash = branch.hash_branch(&left, &right);
-                on_modified_branch(&hash, branch);
+                on_modified_branch(&hash, branch, left, right)?;
                 Ok(hash)
             }
             NodeRef::ModLeaf(leaf) => {
                 let hash = leaf.hash_leaf();
 
-                on_modified_leaf(&hash, leaf);
+                on_modified_leaf(&hash, leaf)?;
                 Ok(hash)
             }
             NodeRef::Stored(stored_idx) => {

src/stored.rs

@@ -1,5 +1,6 @@
 pub mod merkle;
 
+use std::cell::RefCell;
 use std::hash::Hash;
 
 use alloc::{collections::BTreeMap, fmt::Debug, string::String};
@@ -18,10 +19,20 @@ pub trait Store<V> {
     fn get_node(&self, hash_idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>, Self::Error>;
 }
 
-pub trait Database<V> {
-    type Error: Into<String> + Debug;
+pub trait DatabaseGet<V> {
+    type GetError: Into<String> + Debug;
+
+    fn get(&self, hash: &NodeHash) -> Result<Node<Branch<NodeHash>, Leaf<V>>, Self::GetError>;
+}
+
+pub trait DatabaseSet<V>: DatabaseGet<V> {
+    type SetError: Into<String> + Debug;
 
-    fn get(&self, hash: &NodeHash) -> Result<Node<Branch<NodeHash>, Leaf<V>>, Self::Error>;
+    fn set(
+        &self,
+        hash: NodeHash,
+        node: Node<Branch<NodeHash>, Leaf<V>>,
+    ) -> Result<(), Self::GetError>;
 }
 
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
@@ -47,24 +58,38 @@ pub type NodeHash = [u8; 32];
 
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
 pub struct MemoryDb<V> {
-    leaves: BTreeMap<NodeHash, Node<Branch<NodeHash>, Leaf<V>>>,
+    leaves: RefCell<BTreeMap<NodeHash, Node<Branch<NodeHash>, Leaf<V>>>>,
 }
 
 impl<V> MemoryDb<V> {
     pub fn empty() -> Self {
         Self {
-            leaves: BTreeMap::new(),
+            leaves: RefCell::default(),
         }
     }
 }
 
-impl<V: Clone> Database<V> for MemoryDb<V> {
-    type Error = Error;
+impl<V: Clone> DatabaseGet<V> for MemoryDb<V> {
+    type GetError = Error;
 
-    fn get(&self, hash_idx: &NodeHash) -> Result<Node<Branch<NodeHash>, Leaf<V>>, Self::Error> {
+    fn get(&self, hash_idx: &NodeHash) -> Result<Node<Branch<NodeHash>, Leaf<V>>, Self::GetError> {
         self.leaves
+            .borrow()
             .get(hash_idx)
             .cloned()
             .ok_or(Error::NodeNotFound)
     }
 }
+
+impl<V: Clone> DatabaseSet<V> for MemoryDb<V> {
+    type SetError = Error;
+
+    fn set(
+        &self,
+        hash_idx: NodeHash,
+        node: Node<Branch<NodeHash>, Leaf<V>>,
+    ) -> Result<(), Self::SetError> {
+        self.leaves.borrow_mut().insert(hash_idx, node);
+        Ok(())
+    }
+}

src/stored/merkle.rs

@@ -5,7 +5,7 @@ use bumpalo::Bump;
 
 use crate::{Branch, Leaf};
 
-use super::{Database, Error, Idx, Node, NodeHash, Store};
+use super::{DatabaseGet, Error, Idx, Node, NodeHash, Store};
 
 /// A snapshot of the merkle trie
 ///
@@ -104,7 +104,7 @@ impl<V: AsRef<[u8]>> Store<V> for Snapshot<V> {
 
 #[derive(Clone, Debug)]
 pub struct SnapshotBuilder<'a, Db, V> {
-    db: Db,
+    pub db: Db,
     bump: &'a Bump,
 
     /// The root of the trie is always at index 0
@@ -113,7 +113,7 @@ pub struct SnapshotBuilder<'a, Db, V> {
 
 type NodeHashMaybeNode<'a, V> = (NodeHash, Option<Node<&'a Branch<Idx>, &'a Leaf<V>>>);
 
-impl<'a, Db: Database<V>, V: Clone> Store<V> for SnapshotBuilder<'a, Db, V> {
+impl<'a, Db: DatabaseGet<V>, V: Clone> Store<V> for SnapshotBuilder<'a, Db, V> {
     type Error = Error;
 
     fn get_unvisted_hash(&self, hash_idx: Idx) -> Result<&NodeHash, Self::Error> {
@@ -203,7 +203,7 @@ impl<'a, Db, V: Clone> SnapshotBuilder<'a, Db, V> {
         Error,
     >
     where
-        Db: Database<V>,
+        Db: DatabaseGet<V>,
     {
         let Ok(node) = db.get(hash) else {
             return Err(Error::NodeNotFound);