From 1055d4723e34b8280baa428af14fdc5c189dcc79 Mon Sep 17 00:00:00 2001 From: rymnc <43716372+rymnc@users.noreply.github.com> Date: Tue, 18 Apr 2023 14:22:38 +0530 Subject: [PATCH 1/2] fix: custom error enum --- README.md | 12 ++++++------ src/database.rs | 10 +++++----- src/lib.rs | 32 ++++++++++++++++++++++++++------ src/tree.rs | 34 ++++++++++++++++------------------ tests/memory_keccak.rs | 16 ++++++++-------- tests/poseidon.rs | 32 +++++++++++++++----------------- tests/sled_keccak.rs | 20 +++++++++----------- 7 files changed, 85 insertions(+), 71 deletions(-) diff --git a/README.md b/README.md index f0a98ed..720d378 100644 --- a/README.md +++ b/README.md @@ -29,25 +29,25 @@ struct MemoryDBConfig; impl Database for MemoryDB { type Config = MemoryDBConfig; - fn new(_db_config: MemoryDBConfig) -> Result { + fn new(_db_config: MemoryDBConfig) -> PmtreeResult { Ok(MemoryDB(HashMap::new())) } - fn load(_db_config: MemoryDBConfig) -> Result { - Err(Box::new(Error("Cannot load in-memory DB".to_string()))) + fn load(_db_config: MemoryDBConfig) -> PmtreeResult { + Err(DatabaseError(DatabaseErrorKind::CannotLoadDatabase)) } - fn get(&self, key: DBKey) -> Result> { + fn get(&self, key: DBKey) -> PmtreeResult> { Ok(self.0.get(&key).cloned()) } - fn put(&mut self, key: DBKey, value: Value) -> Result<()> { + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> { self.0.insert(key, value); Ok(()) } - fn put_batch(&mut self, subtree: HashMap) -> Result<()> { + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()> { self.0.extend(subtree.into_iter()); Ok(()) diff --git a/src/database.rs b/src/database.rs index 507a489..51b3bfc 100644 --- a/src/database.rs +++ b/src/database.rs @@ -8,21 +8,21 @@ pub trait Database { type Config: Default; /// Creates new instance of db - fn new(config: Self::Config) -> Result + fn new(config: Self::Config) -> PmtreeResult where Self: Sized; /// Loades existing db (existence check required) - fn load(config: Self::Config) -> Result + fn load(config: Self::Config) -> PmtreeResult where Self: Sized; /// Returns value from db by the key - fn get(&self, key: DBKey) -> Result>; + fn get(&self, key: DBKey) -> PmtreeResult>; /// Puts the value to the db by the key - fn put(&mut self, key: DBKey, value: Value) -> Result<()>; + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()>; /// Batc - fn put_batch(&mut self, subtree: HashMap) -> Result<()>; + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()>; } diff --git a/src/lib.rs b/src/lib.rs index 3bd7984..59d158b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,17 +20,37 @@ pub type DBKey = [u8; 8]; /// Denotes values in a database pub type Value = Vec; -/// Denotes `Error` type, for handling DB interaction errors #[derive(Debug)] -pub struct Error(pub String); +pub enum TreeErrorKind { + MerkleTreeIsFull, + InvalidKey, + IndexOutOfBounds, +} -impl std::fmt::Display for Error { +#[derive(Debug)] +pub enum DatabaseErrorKind { + CannotLoadDatabase, + DatabaseExists, +} + +#[derive(Debug)] +pub enum PmtreeErrorKind { + /// Error in database + DatabaseError(DatabaseErrorKind), + /// Error in tree + TreeError(TreeErrorKind), +} + +impl std::fmt::Display for PmtreeErrorKind { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!(f, "{}", self.0) + match self { + PmtreeErrorKind::DatabaseError(e) => write!(f, "Database error: {e:?}"), + PmtreeErrorKind::TreeError(e) => write!(f, "Tree error: {e:?}"), + } } } -impl std::error::Error for Error {} +impl std::error::Error for PmtreeErrorKind {} /// Custom `Result` type with custom `Error` type -pub type Result = std::result::Result>; +pub type PmtreeResult = std::result::Result; diff --git a/src/tree.rs b/src/tree.rs index 99d14c4..0fe90c8 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -44,12 +44,12 @@ where H: Hasher, { /// Creates tree with specified depth and default "pmtree_db" dbpath. - pub fn default(depth: usize) -> Result { + pub fn default(depth: usize) -> PmtreeResult { Self::new(depth, D::Config::default()) } /// Creates new `MerkleTree` and store it to the specified path/db - pub fn new(depth: usize, db_config: D::Config) -> Result { + pub fn new(depth: usize, db_config: D::Config) -> PmtreeResult { // Create new db instance let mut db = D::new(db_config)?; @@ -85,7 +85,7 @@ where } /// Loads existing Merkle Tree from the specified path/db - pub fn load(db_config: D::Config) -> Result { + pub fn load(db_config: D::Config) -> PmtreeResult { // Load existing db instance let db = D::load(db_config)?; @@ -116,9 +116,9 @@ where } /// Sets a leaf at the specified tree index - pub fn set(&mut self, key: usize, leaf: H::Fr) -> Result<()> { + pub fn set(&mut self, key: usize, leaf: H::Fr) -> PmtreeResult<()> { if key >= self.capacity() { - return Err(Box::new(Error("Merkle Tree is full!".to_string()))); + return Err(PmtreeErrorKind::TreeError(TreeErrorKind::IndexOutOfBounds)); } self.db @@ -136,7 +136,7 @@ where } // Recalculates `Merkle Tree` from the specified key - fn recalculate_from(&mut self, key: usize) -> Result<()> { + fn recalculate_from(&mut self, key: usize) -> PmtreeResult<()> { let mut depth = self.depth; let mut i = key; @@ -156,7 +156,7 @@ where } // Hashes the correct couple for the key - fn hash_couple(&self, depth: usize, key: usize) -> Result { + fn hash_couple(&self, depth: usize, key: usize) -> PmtreeResult { let b = key & !1; Ok(H::hash(&[ self.get_elem(Key(depth, b))?, @@ -165,7 +165,7 @@ where } // Returns elem by the key - fn get_elem(&self, key: Key) -> Result { + fn get_elem(&self, key: Key) -> PmtreeResult { let res = self .db .get(key.into())? @@ -175,9 +175,9 @@ where } /// Deletes a leaf at the `key` by setting it to its default value - pub fn delete(&mut self, key: usize) -> Result<()> { + pub fn delete(&mut self, key: usize) -> PmtreeResult<()> { if key >= self.next_index { - return Err(Box::new(Error("The key doesn't exist!".to_string()))); + return Err(PmtreeErrorKind::TreeError(TreeErrorKind::InvalidKey)); } self.set(key, H::default_leaf())?; @@ -186,20 +186,18 @@ where } /// Inserts a leaf to the next available index - pub fn update_next(&mut self, leaf: H::Fr) -> Result<()> { + pub fn update_next(&mut self, leaf: H::Fr) -> PmtreeResult<()> { self.set(self.next_index, leaf)?; Ok(()) } /// Batch insertion, updates the tree in parallel. - pub fn batch_insert(&mut self, leaves: &[H::Fr]) -> Result<()> { + pub fn batch_insert(&mut self, leaves: &[H::Fr]) -> PmtreeResult<()> { let end = self.next_index + leaves.len(); if end > self.capacity() { - return Err(Box::new(Error( - "Not enough space to insert the leaves!".to_string(), - ))); + return Err(PmtreeErrorKind::TreeError(TreeErrorKind::MerkleTreeIsFull)); } let mut subtree = HashMap::::new(); @@ -253,7 +251,7 @@ where subtree: &mut HashMap, leaves: &[H::Fr], from: usize, - ) -> Result<()> { + ) -> PmtreeResult<()> { if key.0 == self.depth { if key.1 >= from { subtree.insert(key, leaves[key.1 - from]); @@ -309,9 +307,9 @@ where } /// Computes a Merkle proof for the leaf at the specified index - pub fn proof(&self, index: usize) -> Result> { + pub fn proof(&self, index: usize) -> PmtreeResult> { if index >= self.capacity() { - return Err(Box::new(Error("Index exceeds set size!".to_string()))); + return Err(PmtreeErrorKind::TreeError(TreeErrorKind::IndexOutOfBounds)); } let mut witness = Vec::with_capacity(self.depth); diff --git a/tests/memory_keccak.rs b/tests/memory_keccak.rs index 53bb623..56b7e91 100644 --- a/tests/memory_keccak.rs +++ b/tests/memory_keccak.rs @@ -12,25 +12,25 @@ struct MemoryDBConfig; impl Database for MemoryDB { type Config = MemoryDBConfig; - fn new(_db_config: MemoryDBConfig) -> Result { + fn new(_db_config: MemoryDBConfig) -> PmtreeResult { Ok(MemoryDB(HashMap::new())) } - fn load(_db_config: MemoryDBConfig) -> Result { - Err(Box::new(Error("Cannot load in-memory DB".to_string()))) + fn load(_db_config: MemoryDBConfig) -> PmtreeResult { + Err(DatabaseError(DatabaseErrorKind::CannotLoadDatabase)) } - fn get(&self, key: DBKey) -> Result> { + fn get(&self, key: DBKey) -> PmtreeResult> { Ok(self.0.get(&key).cloned()) } - fn put(&mut self, key: DBKey, value: Value) -> Result<()> { + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> { self.0.insert(key, value); Ok(()) } - fn put_batch(&mut self, subtree: HashMap) -> Result<()> { + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()> { self.0.extend(subtree.into_iter()); Ok(()) @@ -64,7 +64,7 @@ impl Hasher for MyKeccak { } #[test] -fn insert_delete() -> Result<()> { +fn insert_delete() -> PmtreeResult<()> { let mut mt = MerkleTree::::new(2, MemoryDBConfig)?; assert_eq!(mt.capacity(), 4); @@ -110,7 +110,7 @@ fn insert_delete() -> Result<()> { } #[test] -fn batch_insertions() -> Result<()> { +fn batch_insertions() -> PmtreeResult<()> { let mut mt = MerkleTree::::new(2, MemoryDBConfig)?; let leaves = [ diff --git a/tests/poseidon.rs b/tests/poseidon.rs index 8014f5b..550d0b8 100644 --- a/tests/poseidon.rs +++ b/tests/poseidon.rs @@ -49,25 +49,25 @@ struct MemoryDBConfig; impl Database for MemoryDB { type Config = MemoryDBConfig; - fn new(_db_config: MemoryDBConfig) -> Result { + fn new(_db_config: MemoryDBConfig) -> PmtreeResult { Ok(MemoryDB(HashMap::new())) } - fn load(_db_config: MemoryDBConfig) -> Result { - Err(Box::new(Error("Cannot load in-memory DB".to_string()))) + fn load(_db_config: MemoryDBConfig) -> PmtreeResult { + Err(DatabaseError(DatabaseErrorKind::CannotLoadDatabase)) } - fn get(&self, key: DBKey) -> Result> { + fn get(&self, key: DBKey) -> PmtreeResult> { Ok(self.0.get(&key).cloned()) } - fn put(&mut self, key: DBKey, value: Value) -> Result<()> { + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> { self.0.insert(key, value); Ok(()) } - fn put_batch(&mut self, subtree: HashMap) -> Result<()> { + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()> { self.0.extend(subtree.into_iter()); Ok(()) @@ -82,33 +82,31 @@ struct SledConfig { impl Database for MySled { type Config = SledConfig; - fn new(db_config: SledConfig) -> Result { + fn new(db_config: SledConfig) -> PmtreeResult { let db = sled::open(db_config.path).unwrap(); if db.was_recovered() { - return Err(Box::new(Error("Database exists, try load()!".to_string()))); + return Err(DatabaseError(DatabaseErrorKind::DatabaseExists)); } Ok(MySled(db)) } - fn load(db_config: SledConfig) -> Result { + fn load(db_config: SledConfig) -> PmtreeResult { let db = sled::open(&db_config.path).unwrap(); if !db.was_recovered() { fs::remove_dir_all(&db_config.path).expect("Error removing db"); - return Err(Box::new(Error( - "Trying to load non-existing database!".to_string(), - ))); + return Err(DatabaseError(DatabaseErrorKind::CannotLoadDatabase)); } Ok(MySled(db)) } - fn get(&self, key: DBKey) -> Result> { + fn get(&self, key: DBKey) -> PmtreeResult> { Ok(self.0.get(key).unwrap().map(|val| val.to_vec())) } - fn put(&mut self, key: DBKey, value: Value) -> Result<()> { + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> { self.0.insert(key, value).unwrap(); self.0.flush().unwrap(); @@ -116,7 +114,7 @@ impl Database for MySled { Ok(()) } - fn put_batch(&mut self, subtree: HashMap) -> Result<()> { + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()> { let mut batch = sled::Batch::default(); for (key, value) in subtree { @@ -130,7 +128,7 @@ impl Database for MySled { } #[test] -fn poseidon_memory() -> Result<()> { +fn poseidon_memory() -> PmtreeResult<()> { let mut mt = MerkleTree::::new(TEST_TREE_HEIGHT, MemoryDBConfig)?; let leaf_index = 3; @@ -254,7 +252,7 @@ fn poseidon_memory() -> Result<()> { } #[test] -fn poseidon_sled() -> Result<()> { +fn poseidon_sled() -> PmtreeResult<()> { let mut mt = MerkleTree::::new( TEST_TREE_HEIGHT, SledConfig { diff --git a/tests/sled_keccak.rs b/tests/sled_keccak.rs index 0f8a4eb..a3def2b 100644 --- a/tests/sled_keccak.rs +++ b/tests/sled_keccak.rs @@ -15,33 +15,31 @@ struct SledConfig { impl Database for MySled { type Config = SledConfig; - fn new(db_config: SledConfig) -> Result { + fn new(db_config: SledConfig) -> PmtreeResult { let db = sled::open(db_config.path).unwrap(); if db.was_recovered() { - return Err(Box::new(Error("Database exists, try load()!".to_string()))); + return Err(DatabaseError(DatabaseErrorKind::DatabaseExists)); } Ok(MySled(db)) } - fn load(db_config: SledConfig) -> Result { + fn load(db_config: SledConfig) -> PmtreeResult { let db = sled::open(&db_config.path).unwrap(); if !db.was_recovered() { fs::remove_dir_all(&db_config.path).expect("Error removing db"); - return Err(Box::new(Error( - "Trying to load non-existing database!".to_string(), - ))); + return Err(DatabaseError(DatabaseErrorKind::CannotLoadDatabase)); } Ok(MySled(db)) } - fn get(&self, key: DBKey) -> Result> { + fn get(&self, key: DBKey) -> PmtreeResult> { Ok(self.0.get(key).unwrap().map(|val| val.to_vec())) } - fn put(&mut self, key: DBKey, value: Value) -> Result<()> { + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> { self.0.insert(key, value).unwrap(); self.0.flush().unwrap(); @@ -49,7 +47,7 @@ impl Database for MySled { Ok(()) } - fn put_batch(&mut self, subtree: HashMap) -> Result<()> { + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()> { let mut batch = sled::Batch::default(); for (key, value) in subtree { @@ -89,7 +87,7 @@ impl Hasher for MyKeccak { } #[test] -fn insert_delete() -> Result<()> { +fn insert_delete() -> PmtreeResult<()> { let mut mt = MerkleTree::::new( 2, SledConfig { @@ -142,7 +140,7 @@ fn insert_delete() -> Result<()> { } #[test] -fn batch_insertions() -> Result<()> { +fn batch_insertions() -> PmtreeResult<()> { let mut mt = MerkleTree::::new( 2, SledConfig { From f777a3a45754261b895d8d3d2c698dbdefa26ac6 Mon Sep 17 00:00:00 2001 From: rymnc <43716372+rymnc@users.noreply.github.com> Date: Tue, 18 Apr 2023 14:30:03 +0530 Subject: [PATCH 2/2] fix: tests --- tests/memory_keccak.rs | 4 +++- tests/poseidon.rs | 12 +++++++++--- tests/sled_keccak.rs | 8 ++++++-- 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/tests/memory_keccak.rs b/tests/memory_keccak.rs index 56b7e91..a326e8a 100644 --- a/tests/memory_keccak.rs +++ b/tests/memory_keccak.rs @@ -17,7 +17,9 @@ impl Database for MemoryDB { } fn load(_db_config: MemoryDBConfig) -> PmtreeResult { - Err(DatabaseError(DatabaseErrorKind::CannotLoadDatabase)) + Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CannotLoadDatabase, + )) } fn get(&self, key: DBKey) -> PmtreeResult> { diff --git a/tests/poseidon.rs b/tests/poseidon.rs index 550d0b8..f02f953 100644 --- a/tests/poseidon.rs +++ b/tests/poseidon.rs @@ -54,7 +54,9 @@ impl Database for MemoryDB { } fn load(_db_config: MemoryDBConfig) -> PmtreeResult { - Err(DatabaseError(DatabaseErrorKind::CannotLoadDatabase)) + Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CannotLoadDatabase, + )) } fn get(&self, key: DBKey) -> PmtreeResult> { @@ -85,7 +87,9 @@ impl Database for MySled { fn new(db_config: SledConfig) -> PmtreeResult { let db = sled::open(db_config.path).unwrap(); if db.was_recovered() { - return Err(DatabaseError(DatabaseErrorKind::DatabaseExists)); + return Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::DatabaseExists, + )); } Ok(MySled(db)) @@ -96,7 +100,9 @@ impl Database for MySled { if !db.was_recovered() { fs::remove_dir_all(&db_config.path).expect("Error removing db"); - return Err(DatabaseError(DatabaseErrorKind::CannotLoadDatabase)); + return Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CannotLoadDatabase, + )); } Ok(MySled(db)) diff --git a/tests/sled_keccak.rs b/tests/sled_keccak.rs index a3def2b..6d7e297 100644 --- a/tests/sled_keccak.rs +++ b/tests/sled_keccak.rs @@ -18,7 +18,9 @@ impl Database for MySled { fn new(db_config: SledConfig) -> PmtreeResult { let db = sled::open(db_config.path).unwrap(); if db.was_recovered() { - return Err(DatabaseError(DatabaseErrorKind::DatabaseExists)); + return Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::DatabaseExists, + )); } Ok(MySled(db)) @@ -29,7 +31,9 @@ impl Database for MySled { if !db.was_recovered() { fs::remove_dir_all(&db_config.path).expect("Error removing db"); - return Err(DatabaseError(DatabaseErrorKind::CannotLoadDatabase)); + return Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CannotLoadDatabase, + )); } Ok(MySled(db))