From 32aa2463001c0af9f87633a1ff78858dd4157eb9 Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Fri, 8 Mar 2024 17:39:17 +0100
Subject: [PATCH] [rm-sled] Make proper use of pinning in LMDB adapter +
 comment unsafe

---
 src/db/lmdb_adapter.rs   | 28 +++++++++++++++++++------
 src/db/sqlite_adapter.rs | 44 +++++++++++++++++++++++-----------------
 2 files changed, 47 insertions(+), 25 deletions(-)

diff --git a/src/db/lmdb_adapter.rs b/src/db/lmdb_adapter.rs
index ddfb6ed5..5ce7d3e3 100644
--- a/src/db/lmdb_adapter.rs
+++ b/src/db/lmdb_adapter.rs
@@ -3,6 +3,7 @@ use core::ptr::NonNull;
 
 use std::collections::HashMap;
 use std::convert::TryInto;
+use std::pin::Pin;
 use std::sync::{Arc, RwLock};
 
 use heed::types::ByteSlice;
@@ -319,12 +320,20 @@ where
 	where
 		F: FnOnce(&'a RoTxn<'a>) -> Result<I>,
 	{
-		let mut res = TxAndIterator { tx, iter: None };
+		let res = TxAndIterator { tx, iter: None };
+		let mut boxed = Box::pin(res);
 
-		let tx = unsafe { NonNull::from(&res.tx).as_ref() };
-		res.iter = Some(iterfun(tx)?);
+		// This unsafe allows us to bypass lifetime checks
+		let tx = unsafe { NonNull::from(&boxed.tx).as_ref() };
+		let iter = iterfun(tx)?;
 
-		Ok(Box::new(res))
+		let mut_ref = Pin::as_mut(&mut boxed);
+		// This unsafe allows us to write in a field of the pinned struct
+		unsafe {
+			Pin::get_unchecked_mut(mut_ref).iter = Some(iter);
+		}
+
+		Ok(Box::new(TxAndIteratorPin(boxed)))
 	}
 }
 
@@ -338,14 +347,21 @@ where
 	}
 }
 
-impl<'a, I> Iterator for TxAndIterator<'a, I>
+struct TxAndIteratorPin<'a, I>(Pin<Box<TxAndIterator<'a, I>>>)
+where
+	I: Iterator<Item = IteratorItem<'a>> + 'a;
+
+impl<'a, I> Iterator for TxAndIteratorPin<'a, I>
 where
 	I: Iterator<Item = IteratorItem<'a>> + 'a,
 {
 	type Item = Result<(Value, Value)>;
 
 	fn next(&mut self) -> Option<Self::Item> {
-		match self.iter.as_mut().unwrap().next() {
+		let mut_ref = Pin::as_mut(&mut self.0);
+		// This unsafe allows us to mutably access the iterator field
+		let next = unsafe { Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next() };
+		match next {
 			None => None,
 			Some(Err(e)) => Some(Err(e.into())),
 			Some(Ok((k, v))) => Some(Ok((k.to_vec(), v.to_vec()))),
diff --git a/src/db/sqlite_adapter.rs b/src/db/sqlite_adapter.rs
index 077c1f1b..2c6a4159 100644
--- a/src/db/sqlite_adapter.rs
+++ b/src/db/sqlite_adapter.rs
@@ -444,17 +444,23 @@ impl<'a> DbValueIterator<'a> {
 		let mut boxed = Box::pin(res);
 		trace!("make iterator with sql: {}", sql);
 
+		// This unsafe allows us to bypass lifetime checks
+		let db = unsafe { NonNull::from(&boxed.db).as_ref() };
+		let stmt = db.db.prepare(sql)?;
+
+		let mut_ref = Pin::as_mut(&mut boxed);
+		// This unsafe allows us to write in a field of the pinned struct
 		unsafe {
-			let db = NonNull::from(&boxed.db);
-			let stmt = db.as_ref().db.prepare(sql)?;
-
-			let mut_ref: Pin<&mut DbValueIterator<'a>> = Pin::as_mut(&mut boxed);
 			Pin::get_unchecked_mut(mut_ref).stmt = Some(stmt);
+		}
 
-			let mut stmt = NonNull::from(&boxed.stmt);
-			let iter = stmt.as_mut().as_mut().unwrap().query(args)?;
+		// This unsafe allows us to bypass lifetime checks
+		let stmt = unsafe { NonNull::from(&boxed.stmt).as_mut() };
+		let iter = stmt.as_mut().unwrap().query(args)?;
 
-			let mut_ref: Pin<&mut DbValueIterator<'a>> = Pin::as_mut(&mut boxed);
+		let mut_ref = Pin::as_mut(&mut boxed);
+		// This unsafe allows us to write in a field of the pinned struct
+		unsafe {
 			Pin::get_unchecked_mut(mut_ref).iter = Some(iter);
 		}
 
@@ -476,10 +482,9 @@ impl<'a> Iterator for DbValueIteratorPin<'a> {
 	type Item = Result<(Value, Value)>;
 
 	fn next(&mut self) -> Option<Self::Item> {
-		let next = unsafe {
-			let mut_ref: Pin<&mut DbValueIterator<'a>> = Pin::as_mut(&mut self.0);
-			Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next()
-		};
+		let mut_ref = Pin::as_mut(&mut self.0);
+		// This unsafe allows us to mutably access the iterator field
+		let next = unsafe { Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next() };
 		let row = match next {
 			Err(e) => return Some(Err(e.into())),
 			Ok(None) => return None,
@@ -522,11 +527,13 @@ impl<'a> TxValueIterator<'a> {
 		let mut boxed = Box::pin(res);
 		trace!("make iterator with sql: {}", sql);
 
-		unsafe {
-			let mut stmt = NonNull::from(&boxed.stmt);
-			let iter = stmt.as_mut().query(args)?;
+		// This unsafe allows us to bypass lifetime checks
+		let stmt = unsafe { NonNull::from(&boxed.stmt).as_mut() };
+		let iter = stmt.query(args)?;
 
-			let mut_ref: Pin<&mut TxValueIterator<'a>> = Pin::as_mut(&mut boxed);
+		let mut_ref = Pin::as_mut(&mut boxed);
+		// This unsafe allows us to write in a field of the pinned struct
+		unsafe {
 			Pin::get_unchecked_mut(mut_ref).iter = Some(iter);
 		}
 
@@ -547,10 +554,9 @@ impl<'a> Iterator for TxValueIteratorPin<'a> {
 	type Item = TxOpResult<(Value, Value)>;
 
 	fn next(&mut self) -> Option<Self::Item> {
-		let next = unsafe {
-			let mut_ref: Pin<&mut TxValueIterator<'a>> = Pin::as_mut(&mut self.0);
-			Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next()
-		};
+		let mut_ref = Pin::as_mut(&mut self.0);
+		// This unsafe allows us to mutably access the iterator field
+		let next = unsafe { Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next() };
 		let row = match next {
 			Err(e) => return Some(Err(e.into())),
 			Ok(None) => return None,