From bbed0ba2e6f24e849b674daee17042078ac318c8 Mon Sep 17 00:00:00 2001
From: palkeo <contact@palkeo.com>
Date: Thu, 7 Jan 2021 23:26:18 +0000
Subject: [PATCH] Add size module/trait to estimate size of a key range.

---
 src/database/mod.rs  |  1 +
 src/database/size.rs | 35 +++++++++++++++++++++++++++++++++++
 src/lib.rs           |  1 +
 tests/size.rs        | 22 ++++++++++++++++++++++
 tests/tests.rs       |  3 ++-
 5 files changed, 61 insertions(+), 1 deletion(-)
 create mode 100644 src/database/size.rs
 create mode 100644 tests/size.rs

diff --git a/src/database/mod.rs b/src/database/mod.rs
index cd17f00..5452761 100644
--- a/src/database/mod.rs
+++ b/src/database/mod.rs
@@ -27,6 +27,7 @@ pub mod kv;
 pub mod batch;
 pub mod management;
 pub mod compaction;
+pub mod size;
 pub mod bytes;
 
 #[allow(missing_docs)]
diff --git a/src/database/size.rs b/src/database/size.rs
new file mode 100644
index 0000000..b005f38
--- /dev/null
+++ b/src/database/size.rs
@@ -0,0 +1,35 @@
+//! Compaction
+use super::Database;
+use super::key::Key;
+use leveldb_sys::leveldb_approximate_sizes;
+use libc::{c_char, size_t};
+
+pub trait Size<'a, K: Key + 'a> {
+    fn approximate_size(&self, start: &'a K, stop: &'a K) -> u64;
+}
+
+impl<'a, K: Key + 'a> Size<'a, K> for Database<K> {
+    fn approximate_size(&self, start: &'a K, stop: &'a K) -> u64 {
+        // TODO: Find a way to get rid of these un-necessary copies, while keeping
+        // the code simple. If I put the unsafe call inside the closures, the borrow
+        // checker complains that size is borrowed inside the closures.
+        let start = start.as_slice(|s| Vec::from(s));
+        let stop = stop.as_slice(|s| Vec::from(s));
+
+        let start_ptr = start.as_ptr() as *const c_char;
+        let stop_ptr = stop.as_ptr() as *const c_char;
+
+        let mut size: u64 = 0;
+        unsafe {
+            leveldb_approximate_sizes(
+                self.database.ptr,
+                1,
+                &start_ptr as *const *const c_char,
+                &start.len() as *const size_t,
+                &stop_ptr as *const *const c_char,
+                &stop.len() as *const size_t,
+                &mut size as *mut u64);
+        }
+        size
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index d6672f2..7fb9f96 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -56,6 +56,7 @@ pub use database::kv;
 pub use database::batch;
 pub use database::management;
 pub use database::compaction;
+pub use database::size;
 
 #[allow(missing_docs)]
 pub mod database;
diff --git a/tests/size.rs b/tests/size.rs
new file mode 100644
index 0000000..984af05
--- /dev/null
+++ b/tests/size.rs
@@ -0,0 +1,22 @@
+#[cfg(test)]
+mod size {
+     use utils::{open_database,tmpdir,db_put_simple};
+     use leveldb::size::Size;
+
+    #[test]
+    fn test_iterator_from_to() {
+        let tmp = tmpdir("size");
+        let database = &mut open_database(tmp.path(), true);
+
+        assert_eq!(database.approximate_size(&0, &1_000_000), 0);
+
+        // We need to write a lot of data so it gets written to disk and size approximation works.
+        for i in 0..1_000_000 {
+            db_put_simple(database, i, &[0]);
+        }
+
+        assert!(database.approximate_size(&0, &1_000_000) > 1_000_000);
+        assert!(database.approximate_size(&1, &1000) < 1_000_000);
+        assert!(database.approximate_size(&1, &1000) > 1_000);
+    }
+}
diff --git a/tests/tests.rs b/tests/tests.rs
index 3d6f386..4843977 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -13,4 +13,5 @@ mod cache;
 mod writebatch;
 mod management;
 mod compaction;
-mod concurrent_access;
\ No newline at end of file
+mod concurrent_access;
+mod size;