db abstraction: make .insert() and .remove() return the old value

6 files changed, 127 insertions, 91 deletions

diff --git a/src/db/counted_tree_hack.rs b/src/db/counted_tree_hack.rs
index 52893b41..354a9dad 100644
--- a/src/db/counted_tree_hack.rs
+++ b/src/db/counted_tree_hack.rs
@@ -51,36 +51,26 @@ impl CountedTree {
 
 	// ---- writing functions ----
 
-	pub fn insert<K, V>(&self, key: K, value: V) -> Result<bool>
+	pub fn insert<K, V>(&self, key: K, value: V) -> Result<Option<Value>>
 	where
 		K: AsRef<[u8]>,
 		V: AsRef<[u8]>,
 	{
-		let inserted = self.0.tree.insert(key, value)?;
-		if inserted {
+		let old_val = self.0.tree.insert(key, value)?;
+		if old_val.is_none() {
 			self.0.len.fetch_add(1, Ordering::Relaxed);
 		}
-		Ok(inserted)
+		Ok(old_val)
 	}
 
-	pub fn remove<K: AsRef<[u8]>>(&self, key: K) -> Result<bool> {
-		let removed = self.0.tree.remove(key)?;
-		if removed {
+	pub fn remove<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<Value>> {
+		let old_val = self.0.tree.remove(key)?;
+		if old_val.is_some() {
 			self.0.len.fetch_sub(1, Ordering::Relaxed);
 		}
-		Ok(removed)
+		Ok(old_val)
 	}
 
-	/*
-	pub fn pop_min(&self) -> Result<Option<(Value, Value)>> {
-		let res = self.0.tree.pop_min();
-		if let Ok(Some(_)) = &res {
-			self.0.len.fetch_sub(1, Ordering::Relaxed);
-		};
-		res
-	}
-	*/
-
 	pub fn compare_and_swap<K, OV, NV>(
 		&self,
 		key: K,
diff --git a/src/db/lib.rs b/src/db/lib.rs
index 4543e53c..ed1bdb75 100644
--- a/src/db/lib.rs
+++ b/src/db/lib.rs
@@ -166,15 +166,18 @@ impl Tree {
 			.transpose()
 	}
 
-	/// True if item didn't exist before, false if item already existed
-	/// and was replaced.
+	/// Returns the old value if there was one
 	#[inline]
-	pub fn insert<T: AsRef<[u8]>, U: AsRef<[u8]>>(&self, key: T, value: U) -> Result<bool> {
+	pub fn insert<T: AsRef<[u8]>, U: AsRef<[u8]>>(
+		&self,
+		key: T,
+		value: U,
+	) -> Result<Option<Value>> {
 		self.0.insert(self.1, key.as_ref(), value.as_ref())
 	}
-	/// True if item was removed, false if item already didn't exist
+	/// Returns the old value if there was one
 	#[inline]
-	pub fn remove<T: AsRef<[u8]>>(&self, key: T) -> Result<bool> {
+	pub fn remove<T: AsRef<[u8]>>(&self, key: T) -> Result<Option<Value>> {
 		self.0.remove(self.1, key.as_ref())
 	}
 
@@ -220,20 +223,19 @@ impl<'a> Transaction<'a> {
 		self.0.len(tree.1)
 	}
 
-	/// True if item didn't exist before, false if item already existed
-	/// and was replaced.
+	/// Returns the old value if there was one
 	#[inline]
 	pub fn insert<T: AsRef<[u8]>, U: AsRef<[u8]>>(
 		&mut self,
 		tree: &Tree,
 		key: T,
 		value: U,
-	) -> Result<bool> {
+	) -> Result<Option<Value>> {
 		self.0.insert(tree.1, key.as_ref(), value.as_ref())
 	}
-	/// True if item was removed, false if item already didn't exist
+	/// Returns the old value if there was one
 	#[inline]
-	pub fn remove<T: AsRef<[u8]>>(&mut self, tree: &Tree, key: T) -> Result<bool> {
+	pub fn remove<T: AsRef<[u8]>>(&mut self, tree: &Tree, key: T) -> Result<Option<Value>> {
 		self.0.remove(tree.1, key.as_ref())
 	}
 
@@ -289,8 +291,8 @@ pub(crate) trait IDb: Send + Sync {
 	fn get(&self, tree: usize, key: &[u8]) -> Result<Option<Value>>;
 	fn len(&self, tree: usize) -> Result<usize>;
 
-	fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<bool>;
-	fn remove(&self, tree: usize, key: &[u8]) -> Result<bool>;
+	fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<Option<Value>>;
+	fn remove(&self, tree: usize, key: &[u8]) -> Result<Option<Value>>;
 
 	fn iter(&self, tree: usize) -> Result<ValueIter<'_>>;
 	fn iter_rev(&self, tree: usize) -> Result<ValueIter<'_>>;
@@ -315,8 +317,8 @@ pub(crate) trait ITx {
 	fn get(&self, tree: usize, key: &[u8]) -> Result<Option<Value>>;
 	fn len(&self, tree: usize) -> Result<usize>;
 
-	fn insert(&mut self, tree: usize, key: &[u8], value: &[u8]) -> Result<bool>;
-	fn remove(&mut self, tree: usize, key: &[u8]) -> Result<bool>;
+	fn insert(&mut self, tree: usize, key: &[u8], value: &[u8]) -> Result<Option<Value>>;
+	fn remove(&mut self, tree: usize, key: &[u8]) -> Result<Option<Value>>;
 
 	fn iter(&self, tree: usize) -> Result<ValueIter<'_>>;
 	fn iter_rev(&self, tree: usize) -> Result<ValueIter<'_>>;
diff --git a/src/db/lmdb_adapter.rs b/src/db/lmdb_adapter.rs
index d8e5bcd3..9e4306c8 100644
--- a/src/db/lmdb_adapter.rs
+++ b/src/db/lmdb_adapter.rs
@@ -108,27 +108,28 @@ impl IDb for LmdbDb {
 		}
 	}
 
-	fn remove(&self, tree: usize, key: &[u8]) -> Result<bool> {
-		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> {
 		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<bool> {
+	fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<Option<Value>> {
 		let tree = self.get_tree(tree)?;
 		let mut tx = self.db.write_txn()?;
 		let old_val = tree.get(&tx, key)?.map(Vec::from);
 		tree.put(&mut tx, key, value)?;
 		tx.commit()?;
-		Ok(old_val.is_none())
+		Ok(old_val)
+	}
+
+	fn remove(&self, tree: usize, key: &[u8]) -> Result<Option<Value>> {
+		let tree = self.get_tree(tree)?;
+		let mut tx = self.db.write_txn()?;
+		let old_val = tree.get(&tx, key)?.map(Vec::from);
+		tree.delete(&mut tx, key)?;
+		tx.commit()?;
+		Ok(old_val)
 	}
 
 	fn iter(&self, tree: usize) -> Result<ValueIter<'_>> {
@@ -222,16 +223,17 @@ impl<'a, 'db> ITx for LmdbTx<'a, 'db> {
 		unimplemented!(".len() in transaction not supported with LMDB backend")
 	}
 
-	fn insert(&mut self, tree: usize, key: &[u8], value: &[u8]) -> Result<bool> {
+	fn insert(&mut self, tree: usize, key: &[u8], value: &[u8]) -> Result<Option<Value>> {
 		let tree = *self.get_tree(tree)?;
 		let old_val = tree.get(&self.tx, key)?.map(Vec::from);
 		tree.put(&mut self.tx, key, value)?;
-		Ok(old_val.is_none())
+		Ok(old_val)
 	}
-	fn remove(&mut self, tree: usize, key: &[u8]) -> Result<bool> {
+	fn remove(&mut self, tree: usize, key: &[u8]) -> Result<Option<Value>> {
 		let tree = *self.get_tree(tree)?;
-		let deleted = tree.delete(&mut self.tx, key)?;
-		Ok(deleted)
+		let old_val = tree.get(&self.tx, key)?.map(Vec::from);
+		tree.delete(&mut self.tx, key)?;
+		Ok(old_val)
 	}
 
 	fn iter(&self, _tree: usize) -> Result<ValueIter<'_>> {
diff --git a/src/db/sled_adapter.rs b/src/db/sled_adapter.rs
index 18f457c8..b07401c9 100644
--- a/src/db/sled_adapter.rs
+++ b/src/db/sled_adapter.rs
@@ -83,20 +83,21 @@ impl IDb for SledDb {
 		Ok(val.map(|x| x.to_vec()))
 	}
 
-	fn remove(&self, tree: usize, key: &[u8]) -> Result<bool> {
-		let tree = self.get_tree(tree)?;
-		Ok(tree.remove(key)?.is_some())
-	}
-
 	fn len(&self, tree: usize) -> Result<usize> {
 		let tree = self.get_tree(tree)?;
 		Ok(tree.len())
 	}
 
-	fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<bool> {
+	fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<Option<Value>> {
 		let tree = self.get_tree(tree)?;
 		let old_val = tree.insert(key, value)?;
-		Ok(old_val.is_none())
+		Ok(old_val.map(|x| x.to_vec()))
+	}
+
+	fn remove(&self, tree: usize, key: &[u8]) -> Result<Option<Value>> {
+		let tree = self.get_tree(tree)?;
+		let old_val = tree.remove(key)?;
+		Ok(old_val.map(|x| x.to_vec()))
 	}
 
 	fn iter(&self, tree: usize) -> Result<ValueIter<'_>> {
@@ -206,14 +207,15 @@ impl<'a> ITx for SledTx<'a> {
 		unimplemented!(".len() in transaction not supported with Sled backend")
 	}
 
-	fn insert(&mut self, tree: usize, key: &[u8], value: &[u8]) -> Result<bool> {
+	fn insert(&mut self, tree: usize, key: &[u8], value: &[u8]) -> Result<Option<Value>> {
 		let tree = self.get_tree(tree)?;
 		let old_val = self.save_error(tree.insert(key, value))?;
-		Ok(old_val.is_none())
+		Ok(old_val.map(|x| x.to_vec()))
 	}
-	fn remove(&mut self, tree: usize, key: &[u8]) -> Result<bool> {
+	fn remove(&mut self, tree: usize, key: &[u8]) -> Result<Option<Value>> {
 		let tree = self.get_tree(tree)?;
-		Ok(self.save_error(tree.remove(key))?.is_some())
+		let old_val = self.save_error(tree.remove(key))?;
+		Ok(old_val.map(|x| x.to_vec()))
 	}
 
 	fn iter(&self, _tree: usize) -> Result<ValueIter<'_>> {
diff --git a/src/db/sqlite_adapter.rs b/src/db/sqlite_adapter.rs
index 32557a53..b23885bd 100644
--- a/src/db/sqlite_adapter.rs
+++ b/src/db/sqlite_adapter.rs
@@ -125,18 +125,6 @@ impl IDb for SqliteDb {
 		this.internal_get(tree, key)
 	}
 
-	fn remove(&self, tree: usize, key: &[u8]) -> Result<bool> {
-		trace!("remove {}: lock db", tree);
-		let this = self.0.lock().unwrap();
-		trace!("remove {}: lock acquired", tree);
-
-		let tree = this.get_tree(tree)?;
-		let res = this
-			.db
-			.execute(&format!("DELETE FROM {} WHERE k = ?1", tree), params![key])?;
-		Ok(res > 0)
-	}
-
 	fn len(&self, tree: usize) -> Result<usize> {
 		trace!("len {}: lock db", tree);
 		let this = self.0.lock().unwrap();
@@ -151,18 +139,50 @@ impl IDb for SqliteDb {
 		}
 	}
 
-	fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<bool> {
+	fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<Option<Value>> {
 		trace!("insert {}: lock db", tree);
 		let this = self.0.lock().unwrap();
 		trace!("insert {}: lock acquired", tree);
 
 		let tree = this.get_tree(tree)?;
 		let old_val = this.internal_get(tree, key)?;
-		this.db.execute(
-			&format!("INSERT OR REPLACE INTO {} (k, v) VALUES (?1, ?2)", tree),
-			params![key, value],
-		)?;
-		Ok(old_val.is_none())
+
+		match &old_val {
+			Some(_) => {
+				let n = this.db.execute(
+					&format!("UPDATE {} SET v = ?2 WHERE k = ?1", tree),
+					params![key, value],
+				)?;
+				assert_eq!(n, 1);
+			}
+			None => {
+				let n = this.db.execute(
+					&format!("INSERT INTO {} (k, v) VALUES (?1, ?2)", tree),
+					params![key, value],
+				)?;
+				assert_eq!(n, 1);
+			}
+		}
+
+		Ok(old_val)
+	}
+
+	fn remove(&self, tree: usize, key: &[u8]) -> Result<Option<Value>> {
+		trace!("remove {}: lock db", tree);
+		let this = self.0.lock().unwrap();
+		trace!("remove {}: lock acquired", tree);
+
+		let tree = this.get_tree(tree)?;
+		let old_val = this.internal_get(tree, key)?;
+
+		if old_val.is_some() {
+			let n = this
+				.db
+				.execute(&format!("DELETE FROM {} WHERE k = ?1", tree), params![key])?;
+			assert_eq!(n, 1);
+		}
+
+		Ok(old_val)
 	}
 
 	fn iter(&self, tree: usize) -> Result<ValueIter<'_>> {
@@ -308,21 +328,41 @@ impl<'a> ITx for SqliteTx<'a> {
 		}
 	}
 
-	fn insert(&mut self, tree: usize, key: &[u8], value: &[u8]) -> Result<bool> {
+	fn insert(&mut self, tree: usize, key: &[u8], value: &[u8]) -> Result<Option<Value>> {
 		let tree = self.get_tree(tree)?;
 		let old_val = self.internal_get(tree, key)?;
-		self.tx.execute(
-			&format!("INSERT OR REPLACE INTO {} (k, v) VALUES (?1, ?2)", tree),
-			params![key, value],
-		)?;
-		Ok(old_val.is_none())
+
+		match &old_val {
+			Some(_) => {
+				let n = self.tx.execute(
+					&format!("UPDATE {} SET v = ?2 WHERE k = ?1", tree),
+					params![key, value],
+				)?;
+				assert_eq!(n, 1);
+			}
+			None => {
+				let n = self.tx.execute(
+					&format!("INSERT INTO {} (k, v) VALUES (?1, ?2)", tree),
+					params![key, value],
+				)?;
+				assert_eq!(n, 1);
+			}
+		}
+
+		Ok(old_val)
 	}
-	fn remove(&mut self, tree: usize, key: &[u8]) -> Result<bool> {
+	fn remove(&mut self, tree: usize, key: &[u8]) -> Result<Option<Value>> {
 		let tree = self.get_tree(tree)?;
-		let res = self
-			.tx
-			.execute(&format!("DELETE FROM {} WHERE k = ?1", tree), params![key])?;
-		Ok(res > 0)
+		let old_val = self.internal_get(tree, key)?;
+
+		if old_val.is_some() {
+			let n = self
+				.tx
+				.execute(&format!("DELETE FROM {} WHERE k = ?1", tree), params![key])?;
+			assert_eq!(n, 1);
+		}
+
+		Ok(old_val)
 	}
 
 	fn iter(&self, _tree: usize) -> Result<ValueIter<'_>> {
diff --git a/src/db/test.rs b/src/db/test.rs
index 31fea1f0..cfcee643 100644
--- a/src/db/test.rs
+++ b/src/db/test.rs
@@ -14,13 +14,13 @@ fn test_suite(db: Db) {
 	let vb: &[u8] = &b"plip"[..];
 	let vc: &[u8] = &b"plup"[..];
 
-	tree.insert(ka, va).unwrap();
+	assert!(tree.insert(ka, va).unwrap().is_none());
 	assert_eq!(tree.get(ka).unwrap().unwrap(), va);
 
 	let res = db.transaction::<_, (), _>(|mut tx| {
 		assert_eq!(tx.get(&tree, ka).unwrap().unwrap(), va);
 
-		tx.insert(&tree, ka, vb).unwrap();
+		assert_eq!(tx.insert(&tree, ka, vb).unwrap().unwrap(), va);
 
 		assert_eq!(tx.get(&tree, ka).unwrap().unwrap(), vb);
 
@@ -32,7 +32,7 @@ fn test_suite(db: Db) {
 	let res = db.transaction::<(), _, _>(|mut tx| {
 		assert_eq!(tx.get(&tree, ka).unwrap().unwrap(), vb);
 
-		tx.insert(&tree, ka, vc).unwrap();
+		assert_eq!(tx.insert(&tree, ka, vc).unwrap().unwrap(), vb);
 
 		assert_eq!(tx.get(&tree, ka).unwrap().unwrap(), vc);
 
@@ -47,7 +47,7 @@ fn test_suite(db: Db) {
 	assert!(iter.next().is_none());
 	drop(iter);
 
-	tree.insert(kb, vc).unwrap();
+	assert!(tree.insert(kb, vc).unwrap().is_none());
 	assert_eq!(tree.get(kb).unwrap().unwrap(), vc);
 
 	let mut iter = tree.iter().unwrap();