From bd9711d954f9ae9610190a74f8288b4753aff179 Mon Sep 17 00:00:00 2001 From: Magamedrasul Ibragimov Date: Sun, 6 Nov 2022 15:00:09 +0300 Subject: [PATCH] feat: add Error handling --- src/database.rs | 12 +++-- src/lib.rs | 5 +- src/tree.rs | 110 ++++++++++++++++++++++------------------- tests/memory_keccak.rs | 2 +- tests/sled_keccak.rs | 2 +- 5 files changed, 74 insertions(+), 57 deletions(-) diff --git a/src/database.rs b/src/database.rs index cdcba9a..329940b 100644 --- a/src/database.rs +++ b/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 + where + Self: Sized; /// Loades existing db (existence check required) - fn load(dbpath: &str) -> Self; + fn load(dbpath: &str) -> Result + where + Self: Sized; /// Returns value from db by the key - fn get(&self, key: DBKey) -> Option; + fn get(&self, key: DBKey) -> Result>; /// 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<()>; } diff --git a/src/lib.rs b/src/lib.rs index 11ddab3..e699836 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,7 +20,7 @@ pub type DBKey = [u8; 8]; /// Denotes values in a database pub type Value = Vec; -/// 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 = std::result::Result; diff --git a/src/tree.rs b/src/tree.rs index 56210dd..4dfb1e0 100644 --- a/src/tree.rs +++ b/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 { // 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 { // 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 - pub fn insert(&mut self, leaf: H::Fr) { - // Check if the Merkle Tree is not full - assert!(self.next_index < 1 << self.depth, "Merkle Tree is full!"); + /// Sets a leaf at the specified tree index + pub fn set(&mut self, key: usize, leaf: H::Fr) -> Result<()> { + if key < self.capacity() { + return Err(Error("Merkle Tree is full!".to_string())); + } - // Update the tree - self.set(self.next_index, leaf); + self.db.put(Key(self.depth, key).into(), leaf.into())?; + 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 - 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); + Ok(()) } // 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 { let b = key & !1; - H::hash(&[ - self.get_elem(Key(depth, b)), - self.get_elem(Key(depth, b + 1)), - ]) + Ok(H::hash(&[ + self.get_elem(Key(depth, b))?, + self.get_elem(Key(depth, b + 1))?, + ])) } // Returns elem by the key - fn get_elem(&self, key: Key) -> H::Fr { - self.db - .get(key.into()) + fn get_elem(&self, key: Key) -> Result { + Ok(self + .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 { - self.db.get(Key(0, 0).into()).unwrap().into() + pub fn root(&self) -> Result { + Ok(self.db.get(Key(0, 0).into())?.unwrap().into()) } /// Returns the total number of leaves set diff --git a/tests/memory_keccak.rs b/tests/memory_keccak.rs index 7ec83d4..73e1772 100644 --- a/tests/memory_keccak.rs +++ b/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])); } diff --git a/tests/sled_keccak.rs b/tests/sled_keccak.rs index eb296b6..0dee330 100644 --- a/tests/sled_keccak.rs +++ b/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])); }