From 37470d341cad23b9b56c02be0ad677a69f50ca47 Mon Sep 17 00:00:00 2001
From: Martin Ek <mail@ekmartin.com>
Date: Fri, 6 Apr 2018 14:56:36 +0200
Subject: [PATCH] Add db.full_iterator()

---
 src/db.rs              | 19 ++++++++++++++++
 tests/test_iterator.rs | 50 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/src/db.rs b/src/db.rs
index a26fae9..78fd8c7 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -888,6 +888,15 @@ impl DB {
         DBIterator::new(self, &opts, mode)
     }
 
+    /// Opens an interator with `set_total_order_seek` enabled.
+    /// This must be used to iterate across prefixes when `set_memtable_factory` has been called
+    /// with a Hash-based implementation.
+    pub fn full_iterator(&self, mode: IteratorMode) -> DBIterator {
+        let mut opts = ReadOptions::default();
+        opts.set_total_order_seek(true);
+        DBIterator::new(self, &opts, mode)
+    }
+
     pub fn prefix_iterator<'a>(&self, prefix: &'a [u8]) -> DBIterator {
         let mut opts = ReadOptions::default();
         opts.set_prefix_same_as_start(true);
@@ -903,6 +912,16 @@ impl DB {
         DBIterator::new_cf(self, cf_handle, &opts, mode)
     }
 
+    pub fn full_iterator_cf(
+        &self,
+        cf_handle: ColumnFamily,
+        mode: IteratorMode,
+    ) -> Result<DBIterator, Error> {
+        let mut opts = ReadOptions::default();
+        opts.set_total_order_seek(true);
+        DBIterator::new_cf(self, cf_handle, &opts, mode)
+    }
+
     pub fn prefix_iterator_cf<'a>(
         &self,
         cf_handle: ColumnFamily,
diff --git a/tests/test_iterator.rs b/tests/test_iterator.rs
index 3cb368f..5f8b1bb 100644
--- a/tests/test_iterator.rs
+++ b/tests/test_iterator.rs
@@ -14,7 +14,7 @@
 //
 extern crate rocksdb;
 
-use rocksdb::{DB, Direction, IteratorMode, Options};
+use rocksdb::{DB, Direction, IteratorMode, MemtableFactory, Options};
 
 fn cba(input: &Box<[u8]>) -> Box<[u8]> {
     input.iter().cloned().collect::<Vec<_>>().into_boxed_slice()
@@ -194,3 +194,51 @@ pub fn test_prefix_iterator() {
         }
     }
 }
+
+#[test]
+pub fn test_full_iterator() {
+    let path = "_rust_rocksdb_fulliteratortest";
+    {
+        let a1: Box<[u8]> = key(b"aaa1");
+        let a2: Box<[u8]> = key(b"aaa2");
+        let b1: Box<[u8]> = key(b"bbb1");
+        let b2: Box<[u8]> = key(b"bbb2");
+
+        let prefix_extractor = rocksdb::SliceTransform::create_fixed_prefix(3);
+        let factory = MemtableFactory::HashSkipList {
+            bucket_count: 1_000_000,
+            height: 4,
+            branching_factor: 4,
+        };
+
+        let mut opts = Options::default();
+        opts.create_if_missing(true);
+        opts.set_prefix_extractor(prefix_extractor);
+        opts.set_allow_concurrent_memtable_write(false);
+        opts.set_memtable_factory(factory);
+
+        let db = DB::open(&opts, path).unwrap();
+
+        assert!(db.put(&*a1, &*a1).is_ok());
+        assert!(db.put(&*a2, &*a2).is_ok());
+        assert!(db.put(&*b1, &*b1).is_ok());
+        assert!(db.put(&*b2, &*b2).is_ok());
+
+        // A normal iterator won't work here since we're using a HashSkipList for our memtable
+        // implementation (which buckets keys based on their prefix):
+        let bad_iterator = db.iterator(IteratorMode::Start);
+        assert_eq!(bad_iterator.collect::<Vec<_>>(), vec![]);
+
+        let expected = vec![
+            (cba(&a1), cba(&a1)),
+            (cba(&a2), cba(&a2)),
+            (cba(&b1), cba(&b1)),
+            (cba(&b2), cba(&b2)),
+        ];
+
+        let a_iterator = db.full_iterator(IteratorMode::Start);
+        assert_eq!(a_iterator.collect::<Vec<_>>(), expected)
+    }
+    let opts = Options::default();
+    assert!(DB::destroy(&opts, path).is_ok());
+}