From f7a1c70089bf453554f5ae787da4865caf0ee5c4 Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Fri, 3 Jun 2022 16:18:56 +0200
Subject: [PATCH] Implement iterator for LMDB

---
 src/db/lmdb_adapter.rs | 106 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 98 insertions(+), 8 deletions(-)

diff --git a/src/db/lmdb_adapter.rs b/src/db/lmdb_adapter.rs
index ec1a5444..d568b5c5 100644
--- a/src/db/lmdb_adapter.rs
+++ b/src/db/lmdb_adapter.rs
@@ -3,11 +3,12 @@ use core::ops::Bound;
 use core::pin::Pin;
 use core::ptr::NonNull;
 
-use std::cell::RefCell;
 use std::collections::HashMap;
+use std::convert::TryInto;
 use std::sync::{Arc, RwLock};
 
-use heed::{Env, RoTxn, RwTxn, UntypedDatabase as Database};
+use heed::types::ByteSlice;
+use heed::{BytesDecode, Env, RoTxn, RwTxn, UntypedDatabase as Database};
 
 use crate::{
 	Db, Error, IDb, ITx, ITxFn, IValue, Result, TxError, TxFnResult, TxResult, Value, ValueIter,
@@ -101,11 +102,17 @@ impl IDb for LmdbDb {
 	}
 
 	fn remove(&self, tree: usize, key: &[u8]) -> Result<bool> {
-		unimplemented!()
+		let tree = self.get_tree(tree)?;
+		let mut tx = self.db.write_txn()?;
+		let deleted = tree.delete(&mut tx, &key)?;
+		tx.commit()?;
+		Ok(deleted)
 	}
 
 	fn len(&self, tree: usize) -> Result<usize> {
-		unimplemented!()
+		let tree = self.get_tree(tree)?;
+		let tx = self.db.read_txn()?;
+		Ok(tree.len(&tx)?.try_into().unwrap())
 	}
 
 	fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<()> {
@@ -117,11 +124,15 @@ impl IDb for LmdbDb {
 	}
 
 	fn iter(&self, tree: usize) -> Result<ValueIter<'_>> {
-		unimplemented!()
+		let tree = self.get_tree(tree)?;
+		let tx = self.db.read_txn()?;
+		TxAndIterator::make(tx, |tx| Ok(tree.iter(tx)?))
 	}
 
 	fn iter_rev(&self, tree: usize) -> Result<ValueIter<'_>> {
-		unimplemented!()
+		let tree = self.get_tree(tree)?;
+		let tx = self.db.read_txn()?;
+		TxAndIterator::make(tx, |tx| Ok(tree.rev_iter(tx)?))
 	}
 
 	fn range<'r>(
@@ -130,7 +141,9 @@ impl IDb for LmdbDb {
 		low: Bound<&'r [u8]>,
 		high: Bound<&'r [u8]>,
 	) -> Result<ValueIter<'_>> {
-		unimplemented!()
+		let tree = self.get_tree(tree)?;
+		let tx = self.db.read_txn()?;
+		TxAndIterator::make(tx, |tx| Ok(tree.range(tx, &(low, high))?))
 	}
 	fn range_rev<'r>(
 		&self,
@@ -138,7 +151,9 @@ impl IDb for LmdbDb {
 		low: Bound<&'r [u8]>,
 		high: Bound<&'r [u8]>,
 	) -> Result<ValueIter<'_>> {
-		unimplemented!()
+		let tree = self.get_tree(tree)?;
+		let tx = self.db.read_txn()?;
+		TxAndIterator::make(tx, |tx| Ok(tree.rev_range(tx, &(low, high))?))
 	}
 
 	// ----
@@ -262,3 +277,78 @@ impl<'a> std::borrow::Borrow<[u8]> for TxAndValuePin<'a> {
 		self.as_ref()
 	}
 }
+
+// ----
+
+type IteratorItem<'a> = heed::Result<(
+	<ByteSlice as BytesDecode<'a>>::DItem,
+	<ByteSlice as BytesDecode<'a>>::DItem,
+)>;
+
+struct TxAndIterator<'a, I>
+where
+	I: Iterator<Item = IteratorItem<'a>> + 'a,
+{
+	tx: RoTxn<'a>,
+	iter: Option<I>,
+	_pin: PhantomPinned,
+}
+
+impl<'a, I> TxAndIterator<'a, I>
+where
+	I: Iterator<Item = IteratorItem<'a>> + 'a,
+{
+	fn make<F>(tx: RoTxn<'a>, iterfun: F) -> Result<ValueIter<'a>>
+	where
+		F: FnOnce(&'a RoTxn<'a>) -> Result<I>,
+	{
+		let res = TxAndIterator {
+			tx,
+			iter: None,
+			_pin: PhantomPinned,
+		};
+		let mut boxed = Box::pin(res);
+
+		unsafe {
+			let tx = NonNull::from(&boxed.tx);
+			let iter = iterfun(tx.as_ref())?;
+
+			let mut_ref: Pin<&mut TxAndIterator<'a, I>> = Pin::as_mut(&mut boxed);
+			Pin::get_unchecked_mut(mut_ref).iter = Some(iter);
+		}
+
+		Ok(Box::new(TxAndIteratorPin(boxed)))
+	}
+}
+
+impl<'a, I> Drop for TxAndIterator<'a, I>
+where
+	I: Iterator<Item = IteratorItem<'a>> + 'a,
+{
+	fn drop(&mut self) {
+		drop(self.iter.take());
+	}
+}
+
+struct TxAndIteratorPin<'a, I: Iterator<Item = IteratorItem<'a>> + 'a>(
+	Pin<Box<TxAndIterator<'a, I>>>,
+);
+
+impl<'a, I> Iterator for TxAndIteratorPin<'a, I>
+where
+	I: Iterator<Item = IteratorItem<'a>> + 'a,
+{
+	type Item = Result<(Value<'a>, Value<'a>)>;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		let iter_ref = unsafe {
+			let mut_ref: Pin<&mut TxAndIterator<'a, I>> = Pin::as_mut(&mut self.0);
+			Pin::get_unchecked_mut(mut_ref).iter.as_mut()
+		};
+		match iter_ref.unwrap().next() {
+			None => None,
+			Some(Err(e)) => Some(Err(e.into())),
+			Some(Ok((k, v))) => Some(Ok((k.into(), v.into()))),
+		}
+	}
+}