Merge pull request #6 from dbpunk-labs/feat/cf_it

feat: add iterator

src/merk/mod.rs

@@ -370,6 +370,22 @@ impl Merk {
     pub fn raw_iter(&self) -> rocksdb::DBRawIterator {
         self.db.raw_iterator()
     }
+    pub fn iter_opt(
+        &self,
+        mode: rocksdb::IteratorMode,
+        readopts: rocksdb::ReadOptions,
+    ) -> rocksdb::DBIterator {
+        self.db.iterator_opt(mode, readopts)
+    }
+
+    pub fn iter_opt_aux(
+        &self,
+        mode: rocksdb::IteratorMode,
+        readopts: rocksdb::ReadOptions,
+    ) -> rocksdb::DBIterator {
+        let aux_cf = self.db.cf_handle(AUX_CF_NAME).unwrap();
+        self.db.iterator_cf_opt(aux_cf, readopts, mode)
+    }
 
     pub fn checkpoint<P: AsRef<Path>>(&self, path: P) -> Result<Merk> {
         Checkpoint::new(&self.db)?.create_checkpoint(&path)?;
@@ -509,10 +525,10 @@ fn load_root(db: &DB) -> Result<Option<Tree>> {
 
 #[cfg(test)]
 mod test {
-    use super::{Merk, MerkSource, RefWalker};
+    use super::{Merk, MerkSource, Op, RefWalker};
     use crate::proofs::query::Query;
     use crate::test_utils::*;
-    use crate::Op;
+    use crate::tree;
     use std::ops::Range;
     use std::thread;
     use tempdir::TempDir;
@@ -866,4 +882,63 @@ mod test {
 
         std::fs::remove_dir_all(&path).unwrap();
     }
+    #[test]
+    fn iter_opt_range() {
+        let path = thread::current().name().unwrap().to_owned();
+        let mut merk = TempMerk::open(path).expect("failed to open merk");
+
+        // apply data and aux
+        merk.apply(
+            &[
+                (b"k10".to_vec(), Op::Put(b"v10".to_vec())),
+                (b"k20".to_vec(), Op::Put(b"v20".to_vec())),
+                (b"k30".to_vec(), Op::Put(b"v30".to_vec())),
+                (b"k40".to_vec(), Op::Put(b"v40".to_vec())),
+                (b"k50".to_vec(), Op::Put(b"v50".to_vec())),
+            ],
+            &[
+                (b"k11".to_vec(), Op::Put(b"v11".to_vec())),
+                (b"k21".to_vec(), Op::Put(b"v21".to_vec())),
+                (b"k31".to_vec(), Op::Put(b"v31".to_vec())),
+                (b"k41".to_vec(), Op::Put(b"v41".to_vec())),
+                (b"k51".to_vec(), Op::Put(b"v51".to_vec())),
+            ],
+        )
+        .unwrap();
+
+        // iterate on range ["k20", "k50")
+        let mut readopts = rocksdb::ReadOptions::default();
+        readopts.set_iterate_lower_bound(b"k20".to_vec());
+        readopts.set_iterate_upper_bound(b"k50".to_vec());
+
+        let iter = merk.iter_opt(rocksdb::IteratorMode::Start, readopts);
+        let expected = vec![
+            (b"k20".to_vec(), b"v20".to_vec()),
+            (b"k30".to_vec(), b"v30".to_vec()),
+            (b"k40".to_vec(), b"v40".to_vec()),
+        ];
+        let actual = iter
+            .map(|(k, v)| {
+                let kv = tree::Tree::decode(k.to_vec(), &v);
+                (kv.key().to_vec(), kv.value().to_vec())
+            })
+            .collect::<Vec<_>>();
+        assert_eq!(expected, actual);
+
+        // iterate aux on range ["k21", "k51")
+        let mut readopts_aux = rocksdb::ReadOptions::default();
+        readopts_aux.set_iterate_lower_bound(b"k21".to_vec());
+        readopts_aux.set_iterate_upper_bound(b"k51".to_vec());
+
+        let iter_aux = merk.iter_opt_aux(rocksdb::IteratorMode::Start, readopts_aux);
+        let expected_aux = vec![
+            (b"k21".to_vec(), b"v21".to_vec()),
+            (b"k31".to_vec(), b"v31".to_vec()),
+            (b"k41".to_vec(), b"v41".to_vec()),
+        ];
+        let actual_aux = iter_aux
+            .map(|(k, v)| (k.to_vec(), v.to_vec()))
+            .collect::<Vec<_>>();
+        assert_eq!(expected_aux, actual_aux);
+    }
 }