From e3fa07077c0a726d42958f1102f9075f4d8cd9e1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Thu, 5 Sep 2019 14:53:09 +0200
Subject: [PATCH] feat: Introduce the CfTree and CfIter types

---
 meilidb-data/Cargo.toml     |  2 +-
 meilidb-data/src/cf_tree.rs | 72 +++++++++++++++++++++++++++++++++++++
 meilidb-data/src/lib.rs     |  5 ++-
 3 files changed, 77 insertions(+), 2 deletions(-)
 create mode 100644 meilidb-data/src/cf_tree.rs

diff --git a/meilidb-data/Cargo.toml b/meilidb-data/Cargo.toml
index 1491cba31..afec1b87e 100644
--- a/meilidb-data/Cargo.toml
+++ b/meilidb-data/Cargo.toml
@@ -14,11 +14,11 @@ meilidb-core = { path = "../meilidb-core", version = "0.1.0" }
 meilidb-schema = { path = "../meilidb-schema", version = "0.1.0" }
 meilidb-tokenizer = { path = "../meilidb-tokenizer", version = "0.1.0" }
 ordered-float = { version = "1.0.2", features = ["serde"] }
+rocksdb = "0.12.3"
 sdset = "0.3.2"
 serde = { version = "1.0.99", features = ["derive"] }
 serde_json = "1.0.40"
 siphasher = "0.3.0"
-sled = "0.26.0"
 zerocopy = "0.2.8"
 
 [dependencies.rmp-serde]
diff --git a/meilidb-data/src/cf_tree.rs b/meilidb-data/src/cf_tree.rs
new file mode 100644
index 000000000..853ca2a53
--- /dev/null
+++ b/meilidb-data/src/cf_tree.rs
@@ -0,0 +1,72 @@
+use std::sync::Arc;
+use rocksdb::{DBVector, IteratorMode, Direction};
+use crate::RocksDbResult;
+
+#[derive(Clone)]
+pub struct CfTree(Arc<CfTreeInner>);
+
+struct CfTreeInner {
+    db: rocksdb::DB,
+    name: String,
+}
+
+impl CfTree {
+    pub fn insert<K, V>(&self, key: K, value: V) -> RocksDbResult<()>
+    where K: AsRef<[u8]>,
+          V: AsRef<[u8]>,
+    {
+        let cf = self.0.db.cf_handle(&self.0.name).unwrap();
+        self.0.db.put_cf(cf, key, value)
+    }
+
+    pub fn get<K>(&self, key: K) -> RocksDbResult<Option<DBVector>>
+    where K: AsRef<[u8]>,
+    {
+        let cf = self.0.db.cf_handle(&self.0.name).unwrap();
+        self.0.db.get_cf(cf, key)
+    }
+
+    pub fn remove<K>(&self, key: K) -> RocksDbResult<()>
+    where K: AsRef<[u8]>
+    {
+        let cf = self.0.db.cf_handle(&self.0.name).unwrap();
+        self.0.db.delete_cf(cf, key)
+    }
+
+    /// Start and end key range is inclusive on both bounds.
+    pub fn range<KS, KE>(&self, start: KS, end: KE) -> RocksDbResult<CfIter>
+    where KS: AsRef<[u8]>,
+          KE: AsRef<[u8]>,
+    {
+        let cf = self.0.db.cf_handle(&self.0.name).unwrap();
+
+        let mut iter = self.0.db.iterator_cf(cf, IteratorMode::Start)?;
+        iter.set_mode(IteratorMode::From(start.as_ref(), Direction::Forward));
+
+        let end_bound = Box::from(end.as_ref());
+        Ok(CfIter { iter, end_bound: Some(end_bound) })
+    }
+
+    pub fn iter(&self) -> RocksDbResult<CfIter> {
+        let cf = self.0.db.cf_handle(&self.0.name).unwrap();
+        let iter = self.0.db.iterator_cf(cf, IteratorMode::Start)?;
+        Ok(CfIter { iter, end_bound: None })
+    }
+}
+
+pub struct CfIter<'a> {
+    iter: rocksdb::DBIterator<'a>,
+    end_bound: Option<Box<[u8]>>,
+}
+
+impl Iterator for CfIter<'_> {
+    type Item = (Box<[u8]>, Box<[u8]>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match (self.iter.next(), self.end_bound) {
+            (Some((key, _)), Some(end_bound)) if key > end_bound => None,
+            (Some(entry), _) => Some(entry),
+            (None, _) => None,
+        }
+    }
+}
diff --git a/meilidb-data/src/lib.rs b/meilidb-data/src/lib.rs
index 05c6041ed..9124dcc77 100644
--- a/meilidb-data/src/lib.rs
+++ b/meilidb-data/src/lib.rs
@@ -1,3 +1,4 @@
+mod cf_tree;
 mod database;
 mod document_attr_key;
 mod indexer;
@@ -5,8 +6,10 @@ mod number;
 mod ranked_map;
 mod serde;
 
-pub use sled;
+pub use self::cf_tree::{CfTree, CfIter};
 pub use self::database::{Database, Index, CustomSettingsIndex};
 pub use self::number::Number;
 pub use self::ranked_map::RankedMap;
 pub use self::serde::{compute_document_id, extract_document_id, value_to_string};
+
+pub type RocksDbResult<T> = Result<T, rocksdb::Error>;