From 54051c6f426e7caca9ede0d86e47710d899048a4 Mon Sep 17 00:00:00 2001
From: Christopher Berner <me@cberner.com>
Date: Thu, 21 Dec 2023 09:18:33 -0800
Subject: [PATCH] Add static lifetime versions of range() and multimap table
 get()

These methods return reference counted iterators from read only
transactions
---
 src/db.rs               |  3 +--
 src/multimap_table.rs   | 29 +++++++++++++++++++++
 src/table.rs            | 11 ++++++++
 tests/basic_tests.rs    | 24 +++++++++++++++++
 tests/multimap_tests.rs | 58 +++++++++++++++++++++++++++++++++++++++++
 5 files changed, 123 insertions(+), 2 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index e6606a3e..ab70a1da 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -6,8 +6,7 @@ use crate::tree_store::{
 };
 use crate::types::{RedbKey, RedbValue};
 use crate::{
-    CompactionError, DatabaseError, Durability, ReadOnlyTable, ReadableTable, SavepointError,
-    StorageError,
+    CompactionError, DatabaseError, Durability, ReadOnlyTable, SavepointError, StorageError,
 };
 use crate::{ReadTransaction, Result, WriteTransaction};
 use std::fmt::{Debug, Display, Formatter};
diff --git a/src/multimap_table.rs b/src/multimap_table.rs
index 5eb6cc5c..cf62121e 100644
--- a/src/multimap_table.rs
+++ b/src/multimap_table.rs
@@ -1296,6 +1296,35 @@ impl<'txn, K: RedbKey + 'static, V: RedbKey + 'static> ReadOnlyMultimapTable<'tx
             _lifetime: Default::default(),
         })
     }
+
+    /// This method is like [`ReadableMultimapTable::get()`], but the iterator is reference counted and keeps the transaction
+    /// alive until it is dropped.
+    pub fn get<'a>(&self, key: impl Borrow<K::SelfType<'a>>) -> Result<MultimapValue<'static, V>> {
+        let iter = if let Some(collection) = self.tree.get(key.borrow())? {
+            DynamicCollection::iter(collection, self.transaction_guard.clone(), self.mem.clone())?
+        } else {
+            MultimapValue::new_subtree(
+                BtreeRangeIter::new::<RangeFull, &V::SelfType<'_>>(&(..), None, self.mem.clone())?,
+                self.transaction_guard.clone(),
+            )
+        };
+
+        Ok(iter)
+    }
+
+    /// This method is like [`ReadableMultimapTable::range()`], but the iterator is reference counted and keeps the transaction
+    /// alive until it is dropped.
+    pub fn range<'a, KR>(&self, range: impl RangeBounds<KR>) -> Result<MultimapRange<'static, K, V>>
+    where
+        KR: Borrow<K::SelfType<'a>>,
+    {
+        let inner = self.tree.range(&range)?;
+        Ok(MultimapRange::new(
+            inner,
+            self.transaction_guard.clone(),
+            self.mem.clone(),
+        ))
+    }
 }
 
 impl<'txn, K: RedbKey + 'static, V: RedbKey + 'static> ReadableMultimapTable<K, V>
diff --git a/src/table.rs b/src/table.rs
index 2541df77..061dc69b 100644
--- a/src/table.rs
+++ b/src/table.rs
@@ -439,6 +439,17 @@ impl<'txn, K: RedbKey + 'static, V: RedbValue + 'static> ReadOnlyTable<'txn, K,
             _lifetime: Default::default(),
         })
     }
+
+    /// This method is like [`ReadableTable::range()`], but the iterator is reference counted and keeps the transaction
+    /// alive until it is dropped.
+    pub fn range<'a, KR>(&self, range: impl RangeBounds<KR>) -> Result<Range<'static, K, V>>
+    where
+        KR: Borrow<K::SelfType<'a>>,
+    {
+        self.tree
+            .range(&range)
+            .map(|x| Range::new(x, self.transaction_guard.clone()))
+    }
 }
 
 impl<'txn, K: RedbKey + 'static, V: RedbValue + 'static> ReadableTable<K, V>
diff --git a/tests/basic_tests.rs b/tests/basic_tests.rs
index 4138c832..490b0bdb 100644
--- a/tests/basic_tests.rs
+++ b/tests/basic_tests.rs
@@ -1246,6 +1246,30 @@ fn range_lifetime() {
     assert!(iter.next().is_none());
 }
 
+#[test]
+fn range_arc() {
+    let tmpfile = create_tempfile();
+    let db = Database::create(tmpfile.path()).unwrap();
+
+    let definition: TableDefinition<&str, &str> = TableDefinition::new("x");
+
+    let write_txn = db.begin_write().unwrap();
+    {
+        let mut table = write_txn.open_table(definition).unwrap();
+        table.insert("hello", "world").unwrap();
+    }
+    write_txn.commit().unwrap();
+
+    let mut iter = {
+        let read_txn = db.begin_read().unwrap();
+        let table = read_txn.open_table(definition).unwrap();
+        let start = "hello".to_string();
+        table.range::<&str>(start.as_str()..).unwrap()
+    };
+    assert_eq!(iter.next().unwrap().unwrap().1.value(), "world");
+    assert!(iter.next().is_none());
+}
+
 #[test]
 fn drain_lifetime() {
     let tmpfile = create_tempfile();
diff --git a/tests/multimap_tests.rs b/tests/multimap_tests.rs
index 317096b8..45ac8586 100644
--- a/tests/multimap_tests.rs
+++ b/tests/multimap_tests.rs
@@ -172,6 +172,64 @@ fn range_lifetime() {
     assert!(iter.next().is_none());
 }
 
+#[test]
+fn range_arc_lifetime() {
+    let tmpfile = create_tempfile();
+    let db = Database::create(tmpfile.path()).unwrap();
+
+    let definition: MultimapTableDefinition<&str, &str> = MultimapTableDefinition::new("x");
+
+    let write_txn = db.begin_write().unwrap();
+    {
+        let mut table = write_txn.open_multimap_table(definition).unwrap();
+        table.insert("hello", "world").unwrap();
+    }
+    write_txn.commit().unwrap();
+
+    let mut iter = {
+        let read_txn = db.begin_read().unwrap();
+        let table = read_txn.open_multimap_table(definition).unwrap();
+        let start = "hello".to_string();
+        table.range::<&str>(start.as_str()..).unwrap()
+    };
+    assert_eq!(
+        iter.next()
+            .unwrap()
+            .unwrap()
+            .1
+            .next()
+            .unwrap()
+            .unwrap()
+            .value(),
+        "world"
+    );
+    assert!(iter.next().is_none());
+}
+
+#[test]
+fn get_arc_lifetime() {
+    let tmpfile = create_tempfile();
+    let db = Database::create(tmpfile.path()).unwrap();
+
+    let definition: MultimapTableDefinition<&str, &str> = MultimapTableDefinition::new("x");
+
+    let write_txn = db.begin_write().unwrap();
+    {
+        let mut table = write_txn.open_multimap_table(definition).unwrap();
+        table.insert("hello", "world").unwrap();
+    }
+    write_txn.commit().unwrap();
+
+    let mut iter = {
+        let read_txn = db.begin_read().unwrap();
+        let table = read_txn.open_multimap_table(definition).unwrap();
+        let start = "hello".to_string();
+        table.get(start.as_str()).unwrap()
+    };
+    assert_eq!(iter.next().unwrap().unwrap().value(), "world");
+    assert!(iter.next().is_none());
+}
+
 #[test]
 fn delete() {
     let tmpfile = create_tempfile();