Merge pull request #13 from vacp2p/pmtree-upstream

fix: include db.close, error handling and public db field

Cargo.toml

@@ -7,11 +7,10 @@ repository = "https://github.com/Rate-Limiting-Nullifier/pmtree"
 license = "MIT OR Apache-2.0"
 
 [dev-dependencies]
-hex-literal = "0.3.4"
-tiny-keccak = { version = "2.0.2", features = ["keccak"] }
-sled = "0.34.7"
-rln = { git = "https://github.com/vacp2p/zerokit", rev = "490206a" }
-ark-serialize = "0.3.0"
+hex-literal = "=0.3.4"
+tiny-keccak = { version = "=2.0.2", features = ["keccak"] }
+sled = "=0.34.7"
+ark-serialize = "=0.3.0"
 
 [dependencies]
-rayon = { version = "1.6.1", optional =  false }
+rayon = { version = "=1.7.0", optional =  false }

README.md

@@ -52,6 +52,10 @@ impl Database for MemoryDB {
 
         Ok(())
     }
+
+    fn close(&mut self) -> PmtreeResult<()> {
+        Ok(())
+    }
 }
 
 impl Hasher for MyKeccak {
@@ -90,5 +94,7 @@ fn main() {
         "c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95"
     ))
     .unwrap();
+    // closes the connection to the database
+    mt.close().unwrap();
 }
 ```

src/database.rs

@@ -25,4 +25,7 @@ pub trait Database {
 
     /// Puts the leaves batch to the db
     fn put_batch(&mut self, subtree: HashMap<DBKey, Value>) -> PmtreeResult<()>;
+
+    /// Closes the db connection
+    fn close(&mut self) -> PmtreeResult<()>;
 }

src/tree.rs

@@ -10,6 +10,9 @@ const DEPTH_KEY: DBKey = (u64::MAX - 1).to_be_bytes();
 // db[NEXT_INDEX_KEY] = next_index;
 const NEXT_INDEX_KEY: DBKey = u64::MAX.to_be_bytes();
 
+// Default tree depth
+const DEFAULT_TREE_DEPTH: usize = 20;
+
 // Denotes keys (depth, index) in Merkle Tree. Can be converted to DBKey
 // TODO! Think about using hashing for that
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -27,7 +30,7 @@ where
     D: Database,
     H: Hasher,
 {
-    db: D,
+    pub db: D,
     depth: usize,
     next_index: usize,
     cache: Vec<H::Fr>,
@@ -90,14 +93,21 @@ where
         let db = D::load(db_config)?;
 
         // Load root
-        let root = H::deserialize(db.get(Key(0, 0).into())?.unwrap());
+        let root = match db.get(Key(0, 0).into())? {
+            Some(root) => H::deserialize(root),
+            None => H::default_leaf(),
+        };
 
         // Load depth & next_index values from db
-        let depth = db.get(DEPTH_KEY)?.unwrap().try_into().unwrap();
-        let depth = usize::from_be_bytes(depth);
+        let depth = match db.get(DEPTH_KEY)? {
+            Some(depth) => usize::from_be_bytes(depth.try_into().unwrap()),
+            None => DEFAULT_TREE_DEPTH,
+        };
 
-        let next_index = db.get(NEXT_INDEX_KEY)?.unwrap().try_into().unwrap();
-        let next_index = usize::from_be_bytes(next_index);
+        let next_index = match db.get(NEXT_INDEX_KEY)? {
+            Some(next_index) => usize::from_be_bytes(next_index.try_into().unwrap()),
+            None => 0,
+        };
 
         // Load cache vec
         let mut cache = vec![H::default_leaf(); depth + 1];
@@ -115,6 +125,11 @@ where
         })
     }
 
+    /// Closes the db connection
+    pub fn close(&mut self) -> PmtreeResult<()> {
+        self.db.close()
+    }
+
     /// Sets a leaf at the specified tree index
     pub fn set(&mut self, key: usize, leaf: H::Fr) -> PmtreeResult<()> {
         if key >= self.capacity() {

tests/memory_keccak.rs

@@ -37,6 +37,10 @@ impl Database for MemoryDB {
 
         Ok(())
     }
+
+    fn close(&mut self) -> PmtreeResult<()> {
+        Ok(())
+    }
 }
 
 impl Hasher for MyKeccak {