feat: subtree pairs iterator

grovedb/src/lib.rs

@@ -10,7 +10,7 @@ use std::{
 
 pub use merk::proofs::{query::QueryItem, Query};
 use merk::{self, proofs::query::Map, Merk};
-use rs_merkle::{algorithms::Sha256, MerkleProof, MerkleTree};
+use rs_merkle::{algorithms::Sha256, Hasher, MerkleProof, MerkleTree};
 use storage::{
     rocksdb_storage::{PrefixedRocksDbStorage, PrefixedRocksDbStorageError},
     Storage,
@@ -143,13 +143,12 @@ impl GroveDb {
         key: Vec<u8>,
         mut element: subtree::Element,
     ) -> Result<(), Error> {
-        let compressed_path = Self::compress_path(path, None);
         match &mut element {
             Element::Tree(subtree_root_hash) => {
                 // Helper closure to create a new subtree under path + key
                 let create_subtree_merk =
                     || -> Result<(Vec<u8>, Merk<PrefixedRocksDbStorage>), Error> {
-                        let compressed_path_subtree = Self::compress_path(path, Some(&key));
+                        let compressed_path_subtree = Self::compress_subtree_key(&path, Some(&key));
                         Ok((
                             compressed_path_subtree.clone(),
                             Merk::open(PrefixedRocksDbStorage::new(
@@ -175,6 +174,7 @@ impl GroveDb {
                     }
                     self.propagate_changes(&[&key])?;
                 } else {
+                    let compressed_path = Self::compress_subtree_key(path, None);
                     // Add subtree to another subtree.
                     // First, check if a subtree exists to create a new subtree under it
                     self.subtrees
@@ -206,7 +206,7 @@ impl GroveDb {
                 // Get a Merk by a path
                 let mut merk = self
                     .subtrees
-                    .get_mut(&compressed_path)
+                    .get_mut(&Self::compress_subtree_key(path, None))
                     .ok_or(Error::InvalidPath("no subtree found under that path"))?;
                 element.insert(&mut merk, key)?;
                 self.propagate_changes(path)?;
@@ -237,11 +237,19 @@ impl GroveDb {
         }
     }
 
+    pub fn elements_iterator(&self, path: &[&[u8]]) -> Result<subtree::ElementsIterator, Error> {
+        let merk = self
+            .subtrees
+            .get(&Self::compress_subtree_key(path, None))
+            .ok_or(Error::InvalidPath("no subtree found under that path"))?;
+        Ok(Element::iterator(merk.raw_iter()))
+    }
+
     /// Get tree item without following references
     fn get_raw(&self, path: &[&[u8]], key: &[u8]) -> Result<subtree::Element, Error> {
         let merk = self
             .subtrees
-            .get(&Self::compress_path(path, None))
+            .get(&Self::compress_subtree_key(path, None))
             .ok_or(Error::InvalidPath("no subtree found under that path"))?;
         Element::get(&merk, key)
     }
@@ -290,7 +298,7 @@ impl GroveDb {
                 // Get proof for root tree at current key
                 let root_key_index = self
                     .root_leaf_keys
-                    .get(*key)
+                    .get(&Self::compress_subtree_key(&[key], None))
                     .ok_or(Error::InvalidPath("root key not found"))?;
                 proofs.push(self.root_tree.proof(&[*root_key_index]).to_bytes());
             } else {
@@ -313,7 +321,7 @@ impl GroveDb {
     fn prove_item(&self, path: &[&[u8]], proof_query: Query) -> Result<Vec<u8>, Error> {
         let merk = self
             .subtrees
-            .get(&Self::compress_path(path, None))
+            .get(&Self::compress_subtree_key(path, None))
             .ok_or(Error::InvalidPath("no subtree found under that path"))?;
 
         let proof_result = merk
@@ -417,8 +425,8 @@ impl GroveDb {
                 self.root_tree = Self::build_root_tree(&self.subtrees, &self.root_leaf_keys);
                 break;
             } else {
-                let compressed_path_upper_tree = Self::compress_path(path_slice, None);
-                let compressed_path_subtree = Self::compress_path(path_slice, Some(key));
+                let compressed_path_upper_tree = Self::compress_subtree_key(path_slice, None);
+                let compressed_path_subtree = Self::compress_subtree_key(path_slice, Some(key));
                 let subtree = self
                     .subtrees
                     .get(&compressed_path_subtree)
@@ -437,14 +445,26 @@ impl GroveDb {
 
     /// A helper method to build a prefix to rocksdb keys or identify a subtree
     /// in `subtrees` map by tree path;
-    fn compress_path(path: &[&[u8]], key: Option<&[u8]>) -> Vec<u8> {
-        let mut res = path.iter().fold(Vec::<u8>::new(), |mut acc, p| {
+    fn compress_subtree_key(path: &[&[u8]], key: Option<&[u8]>) -> Vec<u8> {
+        let segments_iter = path.into_iter().map(|x| *x).chain(key.into_iter());
+        let mut segments_count = path.len();
+        if key.is_some() {
+            segments_count += 1;
+        }
+        let mut res = segments_iter.fold(Vec::<u8>::new(), |mut acc, p| {
             acc.extend(p.into_iter());
             acc
         });
-        if let Some(k) = key {
-            res.extend_from_slice(k);
-        }
+
+        res.extend(segments_count.to_ne_bytes());
+        path.into_iter()
+            .map(|x| *x)
+            .chain(key.into_iter())
+            .fold(&mut res, |acc, p| {
+                acc.extend(p.len().to_ne_bytes());
+                acc
+            });
+        res = Sha256::hash(&res).to_vec();
         res
     }
 }

grovedb/src/subtree.rs

@@ -1,9 +1,12 @@
 //! Module for subtrees handling.
 //! Subtrees handling is isolated so basically this module is about adapting
 //! Merk API to GroveDB needs.
-use merk::Op;
+use merk::{tree::Tree, Op};
 use serde::{Deserialize, Serialize};
-use storage::rocksdb_storage::PrefixedRocksDbStorage;
+use storage::{
+    rocksdb_storage::{PrefixedRocksDbStorage, RawPrefixedIterator},
+    RawIterator, Store,
+};
 
 use crate::{Error, Merk};
 
@@ -56,6 +59,36 @@ impl Element {
         merk.apply(&batch, &[])
             .map_err(|e| Error::CorruptedData(e.to_string()))
     }
+
+    pub fn iterator(mut raw_iter: RawPrefixedIterator) -> ElementsIterator {
+        raw_iter.seek_to_first();
+        ElementsIterator { raw_iter }
+    }
+}
+
+pub struct ElementsIterator<'a> {
+    raw_iter: RawPrefixedIterator<'a>,
+}
+
+impl<'a> ElementsIterator<'a> {
+    pub fn next(&mut self) -> Result<Option<(Vec<u8>, Element)>, Error> {
+        Ok(if self.raw_iter.valid() {
+            if let Some((key, value)) = self.raw_iter.key().zip(self.raw_iter.value()) {
+                let tree = <Tree as Store>::decode(value)
+                    .map_err(|e| Error::CorruptedData(e.to_string()))?;
+                let element: Element = bincode::deserialize(tree.value()).map_err(|_| {
+                    Error::CorruptedData(String::from("unable to deserialize element"))
+                })?;
+                let key = key.to_vec();
+                self.raw_iter.next();
+                Some((key, element))
+            } else {
+                None
+            }
+        } else {
+            None
+        })
+    }
 }
 
 #[cfg(test)]

grovedb/src/tests.rs

@@ -235,8 +235,8 @@ fn test_tree_structure_is_presistent() {
 fn test_root_tree_leafs_are_noted() {
     let db = make_grovedb();
     let mut hm = HashMap::new();
-    hm.insert(TEST_LEAF.to_vec(), 0);
-    hm.insert(ANOTHER_TEST_LEAF.to_vec(), 1);
+    hm.insert(GroveDb::compress_subtree_key(&[TEST_LEAF], None), 0);
+    hm.insert(GroveDb::compress_subtree_key(&[ANOTHER_TEST_LEAF], None), 1);
     assert_eq!(db.root_leaf_keys, hm);
     assert_eq!(db.root_tree.leaves_len(), 2);
 }
@@ -428,7 +428,6 @@ fn test_checkpoint() {
         .expect("cannot insert a subtree 2 into GroveDB");
     db.insert(&[b"key1", b"key2"], b"key3".to_vec(), element1.clone())
         .expect("cannot insert an item into GroveDB");
-
     assert_eq!(
         db.get(&[b"key1", b"key2"], b"key3")
             .expect("cannot get from grovedb"),
@@ -517,3 +516,80 @@ fn test_insert_if_not_exists() {
     );
     assert!(matches!(result, Err(Error::InvalidPath(_))));
 }
+
+#[test]
+fn test_subtree_pairs_iterator() {
+    let mut db = make_grovedb();
+    let element = Element::Item(b"ayy".to_vec());
+    let element2 = Element::Item(b"lmao".to_vec());
+
+    // Insert some nested subtrees
+    db.insert(&[TEST_LEAF], b"subtree1".to_vec(), Element::empty_tree())
+        .expect("successful subtree 1 insert");
+    db.insert(
+        &[TEST_LEAF, b"subtree1"],
+        b"subtree11".to_vec(),
+        Element::empty_tree(),
+    )
+    .expect("successful subtree 2 insert");
+    // Insert an element into subtree
+    db.insert(
+        &[TEST_LEAF, b"subtree1", b"subtree11"],
+        b"key1".to_vec(),
+        element.clone(),
+    )
+    .expect("successful value insert");
+    assert_eq!(
+        db.get(&[TEST_LEAF, b"subtree1", b"subtree11"], b"key1")
+            .expect("succesful get 1"),
+        element
+    );
+    db.insert(
+        &[TEST_LEAF, b"subtree1", b"subtree11"],
+        b"key0".to_vec(),
+        element.clone(),
+    )
+    .expect("successful value insert");
+    db.insert(
+        &[TEST_LEAF, b"subtree1"],
+        b"subtree12".to_vec(),
+        Element::empty_tree(),
+    )
+    .expect("successful subtree 3 insert");
+    db.insert(&[TEST_LEAF, b"subtree1"], b"key1".to_vec(), element.clone())
+        .expect("succesful value insert");
+    db.insert(
+        &[TEST_LEAF, b"subtree1"],
+        b"key2".to_vec(),
+        element2.clone(),
+    )
+    .expect("succesful value insert");
+
+    // Iterate over subtree1 to see if keys of other subtrees messed up
+    let mut iter = db
+        .elements_iterator(&[TEST_LEAF, b"subtree1"])
+        .expect("cannot create iterator");
+    assert_eq!(iter.next().unwrap(), Some((b"key1".to_vec(), element)));
+    assert_eq!(iter.next().unwrap(), Some((b"key2".to_vec(), element2)));
+    let subtree_element = iter.next().unwrap().unwrap();
+    assert_eq!(subtree_element.0, b"subtree11".to_vec());
+    assert!(matches!(subtree_element.1, Element::Tree(_)));
+    let subtree_element = iter.next().unwrap().unwrap();
+    assert_eq!(subtree_element.0, b"subtree12".to_vec());
+    assert!(matches!(subtree_element.1, Element::Tree(_)));
+    assert!(matches!(iter.next(), Ok(None)));
+}
+
+#[test]
+fn test_compress_path_not_possible_collision() {
+    let path_a = [b"aa".as_ref(), b"b"];
+    let path_b = [b"a".as_ref(), b"ab"];
+    assert_ne!(
+        GroveDb::compress_subtree_key(&path_a, None),
+        GroveDb::compress_subtree_key(&path_b, None)
+    );
+    assert_eq!(
+        GroveDb::compress_subtree_key(&path_a, None),
+        GroveDb::compress_subtree_key(&path_a, None),
+    );
+}

merk/src/merk/mod.rs

@@ -372,7 +372,10 @@ impl Commit for MerkCommitter {
 
 #[cfg(test)]
 mod test {
-    use storage::rocksdb_storage::{default_rocksdb, PrefixedRocksDbStorage};
+    use storage::{
+        rocksdb_storage::{default_rocksdb, PrefixedRocksDbStorage},
+        RawIterator,
+    };
     use tempdir::TempDir;
 
     use super::{Merk, MerkSource, RefWalker};
@@ -539,7 +542,7 @@ mod test {
 
     #[test]
     fn reopen_iter() {
-        fn collect(iter: &mut rocksdb::DBRawIterator, nodes: &mut Vec<(Vec<u8>, Vec<u8>)>) {
+        fn collect(iter: &mut impl RawIterator, nodes: &mut Vec<(Vec<u8>, Vec<u8>)>) {
             while iter.valid() {
                 nodes.push((iter.key().unwrap().to_vec(), iter.value().unwrap().to_vec()));
                 iter.next();