feat: add Error handling

src/database.rs

@@ -3,14 +3,18 @@ use crate::*;
 /// Trait that must be implemented for a Database
 pub trait Database {
     /// Creates new instance of db
-    fn new(dbpath: &str) -> Self;
+    fn new(dbpath: &str) -> Result<Self>
+    where
+        Self: Sized;
 
     /// Loades existing db (existence check required)
-    fn load(dbpath: &str) -> Self;
+    fn load(dbpath: &str) -> Result<Self>
+    where
+        Self: Sized;
 
     /// Returns value from db by the key
-    fn get(&self, key: DBKey) -> Option<Value>;
+    fn get(&self, key: DBKey) -> Result<Option<Value>>;
 
     /// Puts the value to the db by the key
-    fn put(&mut self, key: DBKey, value: Value);
+    fn put(&mut self, key: DBKey, value: Value) -> Result<()>;
 }

src/lib.rs

@@ -20,7 +20,7 @@ pub type DBKey = [u8; 8];
 /// Denotes values in a database
 pub type Value = Vec<u8>;
 
-/// Denotes Error type, for handling DB interaction errors
+/// Denotes `Error` type, for handling DB interaction errors
 #[derive(Debug)]
 pub struct Error(String);
 
@@ -29,3 +29,6 @@ impl std::fmt::Display for Error {
         write!(f, "{}", self.0)
     }
 }
+
+/// Custom `Result` type with custom `Error` type
+pub type Result<T> = std::result::Result<T, Error>;

src/tree.rs

@@ -1,4 +1,5 @@
 use crate::*;
+use std::cmp::max;
 use std::marker::PhantomData;
 
 // db[DEPTH_KEY] = depth
@@ -37,49 +38,49 @@ where
     H: Hasher,
 {
     /// Creates new `MerkleTree` and store it to the specified path/db
-    pub fn new(depth: usize, dbpath: &str) -> Self {
+    pub fn new(depth: usize, dbpath: &str) -> Result<Self> {
         // Create new db instance
-        let mut db = D::new(dbpath);
+        let mut db = D::new(dbpath)?;
 
         // Insert depth val into db
         let depth_val = depth.to_be_bytes().to_vec();
-        db.put(DEPTH_KEY, depth_val);
+        db.put(DEPTH_KEY, depth_val)?;
 
         // Insert next_index val into db
         let next_index = 0usize;
         let next_index_val = next_index.to_be_bytes().to_vec();
-        db.put(NEXT_INDEX_KEY, next_index_val);
+        db.put(NEXT_INDEX_KEY, next_index_val)?;
 
         // Cache nodes
         let mut cache = vec![H::default_leaf(); depth + 1];
 
         // Initialize one branch of the `Merkle Tree` from bottom to top
         cache[depth] = H::default_leaf();
-        db.put(Key(depth, 0).into(), cache[depth].into());
+        db.put(Key(depth, 0).into(), cache[depth].into())?;
         for i in (0..depth).rev() {
             cache[i] = H::hash(&[cache[i + 1], cache[i + 1]]);
-            db.put(Key(i, 0).into(), cache[i].into());
+            db.put(Key(i, 0).into(), cache[i].into())?;
         }
 
-        Self {
+        Ok(Self {
             db,
             h: PhantomData,
             depth,
             next_index,
             cache,
-        }
+        })
     }
 
     /// Loads existing Merkle Tree from the specified path/db
-    pub fn load(dbpath: &str) -> Self {
+    pub fn load(dbpath: &str) -> Result<Self> {
         // Load existing db instance
-        let db = D::load(dbpath);
+        let db = D::load(dbpath)?;
 
         // Load depth & next_index values from db
-        let depth = db.get(DEPTH_KEY).unwrap().try_into().unwrap();
+        let depth = db.get(DEPTH_KEY)?.unwrap().try_into().unwrap();
         let depth = usize::from_be_bytes(depth);
 
-        let next_index = db.get(NEXT_INDEX_KEY).unwrap().try_into().unwrap();
+        let next_index = db.get(NEXT_INDEX_KEY)?.unwrap().try_into().unwrap();
         let next_index = usize::from_be_bytes(next_index);
 
         // Load cache vec
@@ -89,79 +90,88 @@ where
             cache[i] = H::hash(&[cache[i + 1], cache[i + 1]]);
         }
 
-        Self {
+        Ok(Self {
             db,
             h: PhantomData,
             depth,
             next_index,
             cache,
-        }
+        })
     }
 
-    /// Inserts a leaf to the next available index
+    /// Sets a leaf at the specified tree index
-    pub fn insert(&mut self, leaf: H::Fr) {
+    pub fn set(&mut self, key: usize, leaf: H::Fr) -> Result<()> {
-        // Check if the Merkle Tree is not full
+        if key < self.capacity() {
-        assert!(self.next_index < 1 << self.depth, "Merkle Tree is full!");
+            return Err(Error("Merkle Tree is full!".to_string()));
+        }
 
-        // Update the tree
+        self.db.put(Key(self.depth, key).into(), leaf.into())?;
-        self.set(self.next_index, leaf);
+        self.recalculate_from(key)?;
 
         // Update next_index in memory
-        self.next_index += 1;
+        self.next_index = max(self.next_index, key + 1);
 
         // Update next_index in db
         let next_index_val = self.next_index.to_be_bytes().to_vec();
-        self.db.put(NEXT_INDEX_KEY, next_index_val);
+        self.db.put(NEXT_INDEX_KEY, next_index_val)?;
-    }
 
-    /// Deletes a leaf at the `key` by setting it to its default value
+        Ok(())
-    pub fn delete(&mut self, key: usize) {
-        // Check if the key exists
-        assert!(key < self.next_index, "The key doesn't exist!");
-
-        // Update the tree
-        self.set(key, H::default_leaf());
-    }
-
-    // Sets a leaf at the specified tree index
-    fn set(&mut self, key: usize, leaf: H::Fr) {
-        self.db.put(Key(self.depth, key).into(), leaf.into());
-        self.recalculate_from(key);
     }
 
     // Recalculates `Merkle Tree` from the specified key
-    fn recalculate_from(&mut self, key: usize) {
+    fn recalculate_from(&mut self, key: usize) -> Result<()> {
         let mut depth = self.depth;
         let mut i = key;
 
         while depth != 0 {
-            let value = self.hash_couple(depth, i);
+            let value = self.hash_couple(depth, i)?;
             i >>= 1;
             depth -= 1;
-            self.db.put(Key(depth, i).into(), value.into());
+            self.db.put(Key(depth, i).into(), value.into())?;
         }
+
+        Ok(())
     }
 
     // Hashes the correct couple for the key
-    fn hash_couple(&self, depth: usize, key: usize) -> H::Fr {
+    fn hash_couple(&self, depth: usize, key: usize) -> Result<H::Fr> {
         let b = key & !1;
-        H::hash(&[
+        Ok(H::hash(&[
-            self.get_elem(Key(depth, b)),
+            self.get_elem(Key(depth, b))?,
-            self.get_elem(Key(depth, b + 1)),
+            self.get_elem(Key(depth, b + 1))?,
-        ])
+        ]))
     }
 
     // Returns elem by the key
-    fn get_elem(&self, key: Key) -> H::Fr {
+    fn get_elem(&self, key: Key) -> Result<H::Fr> {
-        self.db
+        Ok(self
-            .get(key.into())
+            .db
+            .get(key.into())?
             .unwrap_or_else(|| self.cache[key.0].into())
-            .into()
+            .into())
+    }
+
+    /// Deletes a leaf at the `key` by setting it to its default value
+    pub fn delete(&mut self, key: usize) -> Result<()> {
+        if key < self.next_index {
+            return Err(Error("The key doesn't exist!".to_string()));
+        }
+
+        self.set(key, H::default_leaf())?;
+
+        Ok(())
+    }
+
+    /// Inserts a leaf to the next available index
+    pub fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
+        self.set(self.next_index, leaf)?;
+
+        Ok(())
     }
 
     /// Returns the root of the tree
-    pub fn root(&self) -> H::Fr {
+    pub fn root(&self) -> Result<H::Fr> {
-        self.db.get(Key(0, 0).into()).unwrap().into()
+        Ok(self.db.get(Key(0, 0).into())?.unwrap().into())
     }
 
     /// Returns the total number of leaves set

tests/memory_keccak.rs

@@ -85,7 +85,7 @@ fn insert_delete() {
     ];
 
     for i in 0..leaves.len() {
-        mt.insert(MyFr(leaves[i]));
+        mt.update_next(MyFr(leaves[i]));
         assert_eq!(mt.root(), MyFr(roots[i]));
     }
 

tests/sled_keccak.rs

@@ -97,7 +97,7 @@ fn insert_delete() {
     ];
 
     for i in 0..leaves.len() {
-        mt.insert(MyFr(leaves[i]));
+        mt.update_next(MyFr(leaves[i]));
         assert_eq!(mt.root(), MyFr(roots[i]));
     }