Remove allocations from SliceTransform's callback

master
Martin Ek 6 years ago
parent 559d7a7eab
commit 91a536d03b
  1. 40
      src/slice_transform.rs
  2. 4
      tests/test_slice_transform.rs

@ -13,11 +13,9 @@
// limitations under the License. // limitations under the License.
use std::ffi::CString; use std::ffi::CString;
use std::mem;
use std::ptr;
use std::slice; use std::slice;
use libc::{self, c_char, c_void, size_t}; use libc::{c_char, c_void, size_t};
use ffi; use ffi;
@ -41,15 +39,15 @@ impl SliceTransform {
transform_fn: TransformFn, transform_fn: TransformFn,
in_domain_fn: Option<InDomainFn>, in_domain_fn: Option<InDomainFn>,
) -> SliceTransform { ) -> SliceTransform {
let cb = Box::new(TransformCallback { let cb = Box::into_raw(Box::new(TransformCallback {
name: CString::new(name.as_bytes()).unwrap(), name: CString::new(name.as_bytes()).unwrap(),
transform_fn, transform_fn,
in_domain_fn, in_domain_fn,
}); }));
let st = unsafe { let st = unsafe {
ffi::rocksdb_slicetransform_create( ffi::rocksdb_slicetransform_create(
mem::transmute(cb), cb as *mut c_void,
Some(slice_transform_destructor_callback), Some(slice_transform_destructor_callback),
Some(transform_callback), Some(transform_callback),
// this is ugly, but I can't get the compiler // this is ugly, but I can't get the compiler
@ -82,18 +80,17 @@ impl SliceTransform {
} }
} }
pub type TransformFn = fn(&[u8]) -> Vec<u8>; pub type TransformFn<'a> = fn(&'a [u8]) -> &'a [u8];
pub type InDomainFn = fn(&[u8]) -> bool; pub type InDomainFn = fn(&[u8]) -> bool;
pub struct TransformCallback { pub struct TransformCallback<'a> {
pub name: CString, pub name: CString,
pub transform_fn: TransformFn, pub transform_fn: TransformFn<'a>,
pub in_domain_fn: Option<InDomainFn>, pub in_domain_fn: Option<InDomainFn>,
} }
pub unsafe extern "C" fn slice_transform_destructor_callback(raw_cb: *mut c_void) { pub unsafe extern "C" fn slice_transform_destructor_callback(raw_cb: *mut c_void) {
let transform: Box<TransformCallback> = mem::transmute(raw_cb); Box::from_raw(raw_cb as *mut TransformCallback);
drop(transform);
} }
pub unsafe extern "C" fn slice_transform_name_callback(raw_cb: *mut c_void) -> *const c_char { pub unsafe extern "C" fn slice_transform_name_callback(raw_cb: *mut c_void) -> *const c_char {
@ -109,16 +106,9 @@ pub unsafe extern "C" fn transform_callback(
) -> *mut c_char { ) -> *mut c_char {
let cb = &mut *(raw_cb as *mut TransformCallback); let cb = &mut *(raw_cb as *mut TransformCallback);
let key = slice::from_raw_parts(raw_key as *const u8, key_len as usize); let key = slice::from_raw_parts(raw_key as *const u8, key_len as usize);
let mut result = (cb.transform_fn)(key); let prefix = (cb.transform_fn)(key);
result.shrink_to_fit(); *dst_length = prefix.len() as size_t;
prefix.as_ptr() as *mut c_char
// copy the result into a C++ destroyable buffer
let buf = libc::malloc(result.len() as size_t);
assert!(!buf.is_null());
ptr::copy(result.as_ptr() as *mut c_void, &mut *buf, result.len());
*dst_length = result.len() as size_t;
buf as *mut c_char
} }
pub unsafe extern "C" fn in_domain_callback( pub unsafe extern "C" fn in_domain_callback(
@ -128,10 +118,6 @@ pub unsafe extern "C" fn in_domain_callback(
) -> u8 { ) -> u8 {
let cb = &mut *(raw_cb as *mut TransformCallback); let cb = &mut *(raw_cb as *mut TransformCallback);
let key = slice::from_raw_parts(raw_key as *const u8, key_len as usize); let key = slice::from_raw_parts(raw_key as *const u8, key_len as usize);
let in_domain = cb.in_domain_fn.unwrap();
if (cb.in_domain_fn.unwrap())(key) { in_domain(key) as u8
1
} else {
0
}
} }

@ -13,8 +13,8 @@ pub fn test_slice_transform() {
let b1: Box<[u8]> = key(b"bbb1"); let b1: Box<[u8]> = key(b"bbb1");
let b2: Box<[u8]> = key(b"bbb2"); let b2: Box<[u8]> = key(b"bbb2");
fn first_three(k: &[u8]) -> Vec<u8> { fn first_three<'a>(k: &'a [u8]) -> &'a [u8] {
k.iter().take(3).cloned().collect() &k[..3]
} }
let prefix_extractor = SliceTransform::create("first_three", first_three, None); let prefix_extractor = SliceTransform::create("first_three", first_three, None);

Loading…
Cancel
Save