From 2096f040922609d6d7c6856caf2b79f15e51fb5d Mon Sep 17 00:00:00 2001
From: Nicolas Gotchac <ngotchac@gmail.com>
Date: Tue, 29 Oct 2019 12:38:16 +0100
Subject: [PATCH] Add lifetimes for `DBIterator` return types  (#294)

---
 src/db.rs              | 16 ++++++++++------
 tests/test_iterator.rs | 18 ++++++++++++++++++
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index 6749d26..7a8a071 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -647,7 +647,7 @@ impl<'a> Snapshot<'a> {
         }
     }
 
-    pub fn iterator(&self, mode: IteratorMode) -> DBIterator {
+    pub fn iterator(&self, mode: IteratorMode) -> DBIterator<'a> {
         let readopts = ReadOptions::default();
         self.iterator_opt(mode, readopts)
     }
@@ -661,7 +661,7 @@ impl<'a> Snapshot<'a> {
         self.iterator_cf_opt(cf_handle, readopts, mode)
     }
 
-    pub fn iterator_opt(&self, mode: IteratorMode, mut readopts: ReadOptions) -> DBIterator {
+    pub fn iterator_opt(&self, mode: IteratorMode, mut readopts: ReadOptions) -> DBIterator<'a> {
         readopts.set_snapshot(self);
         DBIterator::new(self.db, &readopts, mode)
     }
@@ -1171,12 +1171,16 @@ impl DB {
         self.cfs.get(name)
     }
 
-    pub fn iterator(&self, mode: IteratorMode) -> DBIterator {
+    pub fn iterator<'a, 'b: 'a>(&'a self, mode: IteratorMode) -> DBIterator<'b> {
         let readopts = ReadOptions::default();
         self.iterator_opt(mode, &readopts)
     }
 
-    pub fn iterator_opt(&self, mode: IteratorMode, readopts: &ReadOptions) -> DBIterator {
+    pub fn iterator_opt<'a, 'b: 'a>(
+        &'a self,
+        mode: IteratorMode,
+        readopts: &ReadOptions,
+    ) -> DBIterator<'b> {
         DBIterator::new(self, &readopts, mode)
     }
 
@@ -1194,13 +1198,13 @@ impl DB {
     /// Opens an iterator 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 {
+    pub fn full_iterator<'a, 'b: 'a>(&'a self, mode: IteratorMode) -> DBIterator<'b> {
         let mut opts = ReadOptions::default();
         opts.set_total_order_seek(true);
         DBIterator::new(self, &opts, mode)
     }
 
-    pub fn prefix_iterator<P: AsRef<[u8]>>(&self, prefix: P) -> DBIterator {
+    pub fn prefix_iterator<'a, 'b: 'a, P: AsRef<[u8]>>(&'a self, prefix: P) -> DBIterator<'b> {
         let mut opts = ReadOptions::default();
         opts.set_prefix_same_as_start(true);
         DBIterator::new(
diff --git a/tests/test_iterator.rs b/tests/test_iterator.rs
index 4168957..cd3bd34 100644
--- a/tests/test_iterator.rs
+++ b/tests/test_iterator.rs
@@ -299,3 +299,21 @@ fn test_full_iterator() {
     let opts = Options::default();
     assert!(DB::destroy(&opts, path).is_ok());
 }
+
+fn custom_iter(db: &DB) -> impl Iterator<Item = usize> {
+    db.iterator(IteratorMode::Start)
+        .map(|(_, db_value)| db_value.len())
+}
+
+#[test]
+fn test_custom_iterator() {
+    let path = DBPath::new("_rust_rocksdb_customiterator_test");
+    {
+        let mut opts = Options::default();
+        opts.create_if_missing(true);
+        let db = DB::open(&opts, &path).unwrap();
+        let _data = custom_iter(&db).collect::<Vec<usize>>();
+    }
+    let opts = Options::default();
+    assert!(DB::destroy(&opts, path).is_ok());
+}