diff --git a/src/comparator.rs b/src/comparator.rs new file mode 100644 index 0000000..8440f9a --- /dev/null +++ b/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 = 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()); +} diff --git a/src/ffi.rs b/src/ffi.rs index 4dfc518..a704eaf 100644 --- a/src/ffi.rs +++ b/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,15 +320,28 @@ 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)] #[test] 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()); diff --git a/src/lib.rs b/src/lib.rs index 6510a7c..8abf55a 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/rocksdb_options.rs b/src/rocksdb_options.rs index 34f8c84..86c37c9 100644 --- a/src/rocksdb_options.rs +++ b/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 { @@ -67,7 +68,7 @@ impl RocksDBOptions { } } - pub fn add_merge_operator<'a>( &self, name: &str, + pub fn add_merge_operator<'a>(&self, name: &str, merge_fn: fn (&[u8], Option<&[u8]>, &mut MergeOperands) -> Vec) { let cb = Box::new(MergeOperatorCallback { name: CString::new(name.as_bytes()).unwrap(), @@ -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(