1 files changed, 127 insertions, 0 deletions
diff --git a/src/db/counted_tree_hack.rs b/src/db/counted_tree_hack.rs
new file mode 100644
index 00000000..bbe943a2
--- /dev/null
+++ b/src/db/counted_tree_hack.rs
@@ -0,0 +1,127 @@
+//! This hack allows a db tree to keep in RAM a counter of the number of entries
+//! it contains, which is used to call .len() on it.  This is usefull only for
+//! the sled backend where .len() otherwise would have to traverse the whole
+//! tree to count items.  For sqlite and lmdb, this is mostly useless (but
+//! hopefully not harmfull!). Note that a CountedTree cannot be part of a
+//! transaction.
+
+use std::sync::{
+	atomic::{AtomicUsize, Ordering},
+	Arc,
+};
+
+use crate::{Result, Tree, TxError, Value, ValueIter};
+
+#[derive(Clone)]
+pub struct CountedTree(Arc<CountedTreeInternal>);
+
+struct CountedTreeInternal {
+	tree: Tree,
+	len: AtomicUsize,
+}
+
+impl CountedTree {
+	pub fn new(tree: Tree) -> Result<Self> {
+		let len = tree.len()?;
+		Ok(Self(Arc::new(CountedTreeInternal {
+			tree,
+			len: AtomicUsize::new(len),
+		})))
+	}
+
+	pub fn len(&self) -> usize {
+		self.0.len.load(Ordering::SeqCst)
+	}
+
+	pub fn is_empty(&self) -> bool {
+		self.len() == 0
+	}
+
+	pub fn get<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<Value>> {
+		self.0.tree.get(key)
+	}
+
+	pub fn first(&self) -> Result<Option<(Value, Value)>> {
+		self.0.tree.first()
+	}
+
+	pub fn iter(&self) -> Result<ValueIter<'_>> {
+		self.0.tree.iter()
+	}
+
+	// ---- writing functions ----
+
+	pub fn insert<K, V>(&self, key: K, value: V) -> Result<Option<Value>>
+	where
+		K: AsRef<[u8]>,
+		V: AsRef<[u8]>,
+	{
+		let old_val = self.0.tree.insert(key, value)?;
+		if old_val.is_none() {
+			self.0.len.fetch_add(1, Ordering::SeqCst);
+		}
+		Ok(old_val)
+	}
+
+	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::SeqCst);
+		}
+		Ok(old_val)
+	}
+
+	pub fn compare_and_swap<K, OV, NV>(
+		&self,
+		key: K,
+		expected_old: Option<OV>,
+		new: Option<NV>,
+	) -> Result<bool>
+	where
+		K: AsRef<[u8]>,
+		OV: AsRef<[u8]>,
+		NV: AsRef<[u8]>,
+	{
+		let old_some = expected_old.is_some();
+		let new_some = new.is_some();
+
+		let tx_res = self.0.tree.db().transaction(|mut tx| {
+			let old_val = tx.get(&self.0.tree, &key)?;
+			let is_same = match (&old_val, &expected_old) {
+				(None, None) => true,
+				(Some(x), Some(y)) if x == y.as_ref() => true,
+				_ => false,
+			};
+			if is_same {
+				match &new {
+					Some(v) => {
+						tx.insert(&self.0.tree, &key, v)?;
+					}
+					None => {
+						tx.remove(&self.0.tree, &key)?;
+					}
+				}
+				tx.commit(())
+			} else {
+				tx.abort(())
+			}
+		});
+
+		match tx_res {
+			Ok(()) => {
+				match (old_some, new_some) {
+					(false, true) => {
+						self.0.len.fetch_add(1, Ordering::SeqCst);
+					}
+					(true, false) => {
+						self.0.len.fetch_sub(1, Ordering::SeqCst);
+					}
+					_ => (),
+				}
+				Ok(true)
+			}
+			Err(TxError::Abort(())) => Ok(false),
+			Err(TxError::Db(e)) => Err(e),
+		}
+	}
+}