ckb_rocksdb/
merge_operator.rs

1// Copyright 2014 Tyler Neely
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//
15
16//! rustic merge operator
17//!
18//! ```
19//! use ckb_rocksdb::{prelude::*, MergeOperands};
20//! # use ckb_rocksdb::TemporaryDBPath;
21//!
22//! fn concat_merge(new_key: &[u8],
23//!                 existing_val: Option<&[u8]>,
24//!                 operands: &mut MergeOperands)
25//!                 -> Option<Vec<u8>> {
26//!
27//!    let mut result: Vec<u8> = Vec::with_capacity(operands.size_hint().0);
28//!    existing_val.map(|v| {
29//!        for e in v {
30//!            result.push(*e)
31//!        }
32//!    });
33//!    for op in operands {
34//!        for e in op {
35//!            result.push(*e)
36//!        }
37//!    }
38//!    Some(result)
39//! }
40//!
41//! fn main() {
42//!   let path = "_rust_path_to_rocksdb";
43//! # let path = TemporaryDBPath::new();
44//!   let mut opts = Options::default();
45//!   opts.create_if_missing(true);
46//!   opts.set_merge_operator_associative("test operator", concat_merge);
47//! # {
48
49//!   let db = DB::open(&opts, &path).unwrap();
50//!   let p = db.put(b"k1", b"a");
51//!   db.merge(b"k1", b"b");
52//!   db.merge(b"k1", b"c");
53//!   db.merge(b"k1", b"d");
54//!   db.merge(b"k1", b"efg");
55//!   let r = db.get(b"k1");
56//!   assert!(r.unwrap().unwrap().to_utf8().unwrap() == "abcdefg");
57
58//! # }
59//! }
60//! ```
61
62use libc::{self, c_char, c_int, c_void, size_t};
63use std::ffi::CString;
64use std::mem;
65use std::ptr;
66use std::slice;
67
68pub trait MergeFn:
69    Fn(&[u8], Option<&[u8]>, &mut MergeOperands) -> Option<Vec<u8>> + Send + Sync + 'static
70{
71}
72impl<F> MergeFn for F where
73    F: Fn(&[u8], Option<&[u8]>, &mut MergeOperands) -> Option<Vec<u8>> + Send + Sync + 'static
74{
75}
76
77pub struct MergeOperatorCallback<F: MergeFn, PF: MergeFn> {
78    pub name: CString,
79    pub full_merge_fn: F,
80    pub partial_merge_fn: PF,
81}
82
83pub unsafe extern "C" fn destructor_callback<F: MergeFn, PF: MergeFn>(raw_cb: *mut c_void) {
84    let _: Box<MergeOperatorCallback<F, PF>> =
85        Box::from_raw(raw_cb as *mut MergeOperatorCallback<F, PF>);
86}
87
88pub unsafe extern "C" fn delete_callback(
89    _raw_cb: *mut c_void,
90    value: *const c_char,
91    value_length: size_t,
92) {
93    if !value.is_null() {
94        let _ = Box::from_raw(slice::from_raw_parts_mut(value as *mut u8, value_length));
95    }
96}
97
98pub unsafe extern "C" fn name_callback<F: MergeFn, PF: MergeFn>(
99    raw_cb: *mut c_void,
100) -> *const c_char {
101    let cb = &mut *(raw_cb as *mut MergeOperatorCallback<F, PF>);
102    cb.name.as_ptr()
103}
104
105pub unsafe extern "C" fn full_merge_callback<F: MergeFn, PF: MergeFn>(
106    raw_cb: *mut c_void,
107    raw_key: *const c_char,
108    key_len: size_t,
109    existing_value: *const c_char,
110    existing_value_len: size_t,
111    operands_list: *const *const c_char,
112    operands_list_len: *const size_t,
113    num_operands: c_int,
114    success: *mut u8,
115    new_value_length: *mut size_t,
116) -> *mut c_char {
117    let cb = &mut *(raw_cb as *mut MergeOperatorCallback<F, PF>);
118    let operands = &mut MergeOperands::new(operands_list, operands_list_len, num_operands);
119    let key = slice::from_raw_parts(raw_key as *const u8, key_len);
120    let oldval = if existing_value.is_null() {
121        None
122    } else {
123        Some(slice::from_raw_parts(
124            existing_value as *const u8,
125            existing_value_len,
126        ))
127    };
128    (cb.full_merge_fn)(key, oldval, operands).map_or_else(
129        || {
130            *new_value_length = 0;
131            *success = 0_u8;
132            ptr::null_mut() as *mut c_char
133        },
134        |result| {
135            *new_value_length = result.len() as size_t;
136            *success = 1_u8;
137            Box::into_raw(result.into_boxed_slice()) as *mut c_char
138        },
139    )
140}
141
142pub unsafe extern "C" fn partial_merge_callback<F: MergeFn, PF: MergeFn>(
143    raw_cb: *mut c_void,
144    raw_key: *const c_char,
145    key_len: size_t,
146    operands_list: *const *const c_char,
147    operands_list_len: *const size_t,
148    num_operands: c_int,
149    success: *mut u8,
150    new_value_length: *mut size_t,
151) -> *mut c_char {
152    let cb = &mut *(raw_cb as *mut MergeOperatorCallback<F, PF>);
153    let operands = &mut MergeOperands::new(operands_list, operands_list_len, num_operands);
154    let key = slice::from_raw_parts(raw_key as *const u8, key_len);
155    (cb.partial_merge_fn)(key, None, operands).map_or_else(
156        || {
157            *new_value_length = 0;
158            *success = 0_u8;
159            ptr::null_mut::<c_char>()
160        },
161        |result| {
162            *new_value_length = result.len() as size_t;
163            *success = 1_u8;
164            Box::into_raw(result.into_boxed_slice()) as *mut c_char
165        },
166    )
167}
168
169pub struct MergeOperands {
170    operands_list: *const *const c_char,
171    operands_list_len: *const size_t,
172    num_operands: usize,
173    cursor: usize,
174}
175
176impl MergeOperands {
177    fn new(
178        operands_list: *const *const c_char,
179        operands_list_len: *const size_t,
180        num_operands: c_int,
181    ) -> MergeOperands {
182        assert!(num_operands >= 0);
183        MergeOperands {
184            operands_list,
185            operands_list_len,
186            num_operands: num_operands as usize,
187            cursor: 0,
188        }
189    }
190}
191
192impl<'a> Iterator for &'a mut MergeOperands {
193    type Item = &'a [u8];
194
195    fn next(&mut self) -> Option<&'a [u8]> {
196        if self.cursor == self.num_operands {
197            None
198        } else {
199            unsafe {
200                let base = self.operands_list as usize;
201                let base_len = self.operands_list_len as usize;
202                let spacing = mem::size_of::<*const *const u8>();
203                let spacing_len = mem::size_of::<*const size_t>();
204                let len_ptr = (base_len + (spacing_len * self.cursor)) as *const size_t;
205                let len = *len_ptr;
206                let ptr = base + (spacing * self.cursor);
207                self.cursor += 1;
208                Some(slice::from_raw_parts(
209                    *(ptr as *const *const u8) as *const u8,
210                    len,
211                ))
212            }
213        }
214    }
215
216    fn size_hint(&self) -> (usize, Option<usize>) {
217        let remaining = self.num_operands - self.cursor;
218        (remaining, Some(remaining))
219    }
220}