Merge pull request #13 from spacejam/tyler_comparator

initial comparator

src/comparator.rs

@@ -0,0 +1,79 @@
+/*
+   Copyright 2014 Tyler Neely
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+extern crate libc;
+use self::libc::{c_char, c_int, c_void, size_t};
+use std::ffi::CString;
+use std::mem;
+use std::ptr;
+use std::slice;
+
+use rocksdb_options::RocksDBOptions;
+use rocksdb::RocksDB;
+
+pub struct ComparatorCallback {
+    pub name: CString,
+    pub f: fn (&[u8], &[u8]) -> i32,
+}
+
+pub extern "C" fn destructor_callback(raw_cb: *mut c_void) {
+    // turn this back into a local variable so rust will reclaim it
+    let _: Box<ComparatorCallback> = unsafe {mem::transmute(raw_cb)};
+}
+
+pub extern "C" fn name_callback(raw_cb: *mut c_void) -> *const c_char {
+    unsafe {
+        let cb: &mut ComparatorCallback =
+            &mut *(raw_cb as *mut ComparatorCallback);
+        let ptr = cb.name.as_ptr();
+        ptr as *const c_char
+    }
+}
+
+pub extern "C" fn compare_callback(
+    raw_cb: *mut c_void,
+    a_raw: *const c_char, a_len: size_t,
+    b_raw: *const c_char, b_len: size_t) -> c_int {
+    unsafe {
+        let cb: &mut ComparatorCallback =
+            &mut *(raw_cb as *mut ComparatorCallback);
+        let a: &[u8] = slice::from_raw_parts(a_raw as *const u8, a_len as usize);
+        let b: &[u8] = slice::from_raw_parts(b_raw as *const u8, b_len as usize);
+        (cb.f)(a, b)
+    }
+}
+
+fn test_reverse_compare(a: &[u8], b: &[u8]) -> c_int {
+    if a < b {
+        1
+    } else if a > b {
+        -1
+    } else {
+        0
+    }
+}
+
+#[allow(dead_code)]
+#[test]
+fn compare_works() {
+    let path = "_rust_rocksdb_comparetest";
+    let opts = RocksDBOptions::new();
+    opts.create_if_missing(true);
+    opts.add_comparator("test comparator", test_reverse_compare);
+    let db = RocksDB::open(opts, path).unwrap();
+    // TODO add interesting test
+    db.close();
+    assert!(RocksDB::destroy(opts, path).is_ok());
+}

src/ffi.rs

@@ -53,6 +53,9 @@ pub struct RocksDBCFHandle(pub *const c_void);
 #[derive(Copy, Clone)]
 #[repr(C)]
 pub struct RocksDBWriteBatch(pub *const c_void);
+#[derive(Copy, Clone)]
+#[repr(C)]
+pub struct RocksDBComparator(pub *const c_void);
 
 pub fn new_bloom_filter(bits: c_int) -> RocksDBFilterPolicy {
     unsafe {
@@ -317,6 +320,20 @@ extern {
                               k: *const u8, klen: size_t));
     pub fn rocksdb_writebatch_data(batch: RocksDBWriteBatch,
                                    size: *mut size_t) -> *const u8;
+
+    // Comparator
+    pub fn rocksdb_options_set_comparator(options: RocksDBOptions,
+                                          cb: RocksDBComparator);
+    pub fn rocksdb_comparator_create(
+        state: *mut c_void,
+        destroy: extern fn(*mut c_void) -> (),
+        compare: extern fn (arg: *mut c_void,
+                            a: *const c_char, alen: size_t,
+                            b: *const c_char, blen: size_t
+                           ) -> c_int,
+        name_fn: extern fn(*mut c_void) -> *const c_char
+    ) -> RocksDBComparator;
+    pub fn rocksdb_comparator_destroy(cmp: RocksDBComparator);
 }
 
 #[allow(dead_code)]
@@ -324,8 +341,7 @@ extern {
 fn internal() {
     unsafe {
         let opts = rocksdb_options_create();
-        let RocksDBOptions(opt_ptr) = opts;
-        assert!(!opt_ptr.is_null());
+        assert!(!opts.0.is_null());
 
         rocksdb_options_increase_parallelism(opts, 0);
         rocksdb_options_optimize_level_style_compaction(opts, 0);
@@ -344,7 +360,6 @@ fn internal() {
 
         let key = b"name\x00";
         let val = b"spacejam\x00";
-
         rocksdb_put(db, writeopts.clone(), key.as_ptr(), 4, val.as_ptr(), 8, err);
         rocksdb_writeoptions_destroy(writeopts);
         assert!(err.is_null());

src/lib.rs

@@ -28,6 +28,7 @@ pub use ffi::{
     RocksDBCompactionStyle,
     RocksDBCompressionType,
     RocksDBSnapshot,
+    RocksDBComparator,
 };
 pub use rocksdb::{
     RocksDB,
@@ -46,3 +47,4 @@ pub mod rocksdb;
 pub mod ffi;
 pub mod rocksdb_options;
 pub mod merge_operator;
+pub mod comparator;

src/rocksdb_options.rs

@@ -19,8 +19,9 @@ use std::ffi::CString;
 use std::mem;
 
 use rocksdb_ffi;
-use merge_operator::{MergeOperatorCallback, MergeOperands, destructor_callback, full_merge_callback,
-              partial_merge_callback, name_callback};
+use merge_operator::{self, MergeOperatorCallback, MergeOperands, full_merge_callback,
+              partial_merge_callback};
+use comparator::{self, ComparatorCallback, compare_callback};
 
 #[derive(Copy, Clone)]
 pub struct RocksDBOptions {
@@ -77,15 +78,32 @@ impl RocksDBOptions {
         unsafe {
             let mo = rocksdb_ffi::rocksdb_mergeoperator_create(
                 mem::transmute(cb),
-                destructor_callback,
+                merge_operator::destructor_callback,
                 full_merge_callback,
                 partial_merge_callback,
                 None,
-                name_callback);
+                merge_operator::name_callback);
             rocksdb_ffi::rocksdb_options_set_merge_operator(self.inner, mo);
         }
     }
 
+    pub fn add_comparator<'a>(&self, name: &str, compare_fn: fn(&[u8], &[u8]) -> i32) {
+        let cb = Box::new(ComparatorCallback {
+            name: CString::new(name.as_bytes()).unwrap(),
+            f: compare_fn,
+        });
+
+        unsafe {
+            let cmp = rocksdb_ffi::rocksdb_comparator_create(
+                mem::transmute(cb),
+                comparator::destructor_callback,
+                compare_callback,
+                comparator::name_callback);
+            rocksdb_ffi::rocksdb_options_set_comparator(self.inner, cmp);
+        }
+    }
+
+
     pub fn set_block_size(&self, size: u64) {
         unsafe {
             rocksdb_ffi::rocksdb_block_based_options_set_block_size(