Allow passing &CStr arguments (#641)

src/db.rs

@@ -19,7 +19,7 @@ use crate::{
     column_family::UnboundColumnFamily,
     db_options::OptionsMustOutliveDB,
     ffi,
-    ffi_util::{from_cstr, opt_bytes_to_ptr, raw_data, to_cpath},
+    ffi_util::{from_cstr, opt_bytes_to_ptr, raw_data, to_cpath, CStrLike},
     ColumnFamily, ColumnFamilyDescriptor, CompactOptions, DBIteratorWithThreadMode,
     DBPinnableSlice, DBRawIteratorWithThreadMode, DBWALIterator, Direction, Error, FlushOptions,
     IngestExternalFileOptions, IteratorMode, Options, ReadOptions, SnapshotWithThreadMode,
@@ -1080,16 +1080,14 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
 
     fn create_inner_cf_handle(
         &self,
-        name: &str,
+        name: impl CStrLike,
         opts: &Options,
     ) -> Result<*mut ffi::rocksdb_column_family_handle_t, Error> {
-        let cf_name = if let Ok(c) = CString::new(name.as_bytes()) {
-            c
-        } else {
-            return Err(Error::new(
-                "Failed to convert path to CString when creating cf".to_owned(),
-            ));
-        };
+        let cf_name = name.bake().map_err(|err| {
+            Error::new(format!(
+                "Failed to convert path to CString when creating cf: {err}"
+            ))
+        })?;
         Ok(unsafe {
             ffi_try!(ffi::rocksdb_create_column_family(
                 self.inner,
@@ -1567,11 +1565,11 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
     /// the end. That string is parsed using `parse` callback which produces
     /// the returned result.
     fn property_value_impl<R>(
-        name: &str,
+        name: impl CStrLike,
         get_property: impl FnOnce(*const c_char) -> *mut c_char,
         parse: impl FnOnce(&str) -> Result<R, Error>,
     ) -> Result<Option<R>, Error> {
-        let value = match CString::new(name) {
+        let value = match name.bake() {
             Ok(prop_name) => get_property(prop_name.as_ptr()),
             Err(e) => {
                 return Err(Error::new(format!(
@@ -1600,7 +1598,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
     ///
     /// Full list of properties could be find
     /// [here](https://github.com/facebook/rocksdb/blob/08809f5e6cd9cc4bc3958dd4d59457ae78c76660/include/rocksdb/db.h#L428-L634).
-    pub fn property_value(&self, name: &str) -> Result<Option<String>, Error> {
+    pub fn property_value(&self, name: impl CStrLike) -> Result<Option<String>, Error> {
         Self::property_value_impl(
             name,
             |prop_name| unsafe { ffi::rocksdb_property_value(self.inner, prop_name) },
@@ -1615,7 +1613,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
     pub fn property_value_cf(
         &self,
         cf: &impl AsColumnFamilyRef,
-        name: &str,
+        name: impl CStrLike,
     ) -> Result<Option<String>, Error> {
         Self::property_value_impl(
             name,
@@ -1639,7 +1637,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
     ///
     /// Full list of properties that return int values could be find
     /// [here](https://github.com/facebook/rocksdb/blob/08809f5e6cd9cc4bc3958dd4d59457ae78c76660/include/rocksdb/db.h#L654-L689).
-    pub fn property_int_value(&self, name: &str) -> Result<Option<u64>, Error> {
+    pub fn property_int_value(&self, name: impl CStrLike) -> Result<Option<u64>, Error> {
         Self::property_value_impl(
             name,
             |prop_name| unsafe { ffi::rocksdb_property_value(self.inner, prop_name) },
@@ -1654,7 +1652,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
     pub fn property_int_value_cf(
         &self,
         cf: &impl AsColumnFamilyRef,
-        name: &str,
+        name: impl CStrLike,
     ) -> Result<Option<u64>, Error> {
         Self::property_value_impl(
             name,

src/db_options.rs

@@ -24,6 +24,7 @@ use crate::{
     comparator::{self, ComparatorCallback, CompareFn},
     db::DBAccess,
     ffi,
+    ffi_util::CStrLike,
     merge_operator::{
         self, full_merge_callback, partial_merge_callback, MergeFn, MergeOperatorCallback,
     },
@@ -1274,11 +1275,11 @@ impl Options {
 
     pub fn set_merge_operator_associative<F: MergeFn + Clone>(
         &mut self,
-        name: &str,
+        name: impl CStrLike,
         full_merge_fn: F,
     ) {
         let cb = Box::new(MergeOperatorCallback {
-            name: CString::new(name.as_bytes()).unwrap(),
+            name: name.into_c_string().unwrap(),
             full_merge_fn: full_merge_fn.clone(),
             partial_merge_fn: full_merge_fn,
         });
@@ -1298,12 +1299,12 @@ impl Options {
 
     pub fn set_merge_operator<F: MergeFn, PF: MergeFn>(
         &mut self,
-        name: &str,
+        name: impl CStrLike,
         full_merge_fn: F,
         partial_merge_fn: PF,
     ) {
         let cb = Box::new(MergeOperatorCallback {
-            name: CString::new(name.as_bytes()).unwrap(),
+            name: name.into_c_string().unwrap(),
             full_merge_fn,
             partial_merge_fn,
         });
@@ -1339,12 +1340,12 @@ impl Options {
     ///
     /// If multi-threaded compaction is used, `filter_fn` may be called multiple times
     /// simultaneously.
-    pub fn set_compaction_filter<F>(&mut self, name: &str, filter_fn: F)
+    pub fn set_compaction_filter<F>(&mut self, name: impl CStrLike, filter_fn: F)
     where
         F: CompactionFilterFn + Send + 'static,
     {
         let cb = Box::new(CompactionFilterCallback {
-            name: CString::new(name.as_bytes()).unwrap(),
+            name: name.into_c_string().unwrap(),
             filter_fn,
         });
 
@@ -1391,9 +1392,9 @@ impl Options {
     /// The client must ensure that the comparator supplied here has the same
     /// name and orders keys *exactly* the same as the comparator provided to
     /// previous open calls on the same DB.
-    pub fn set_comparator(&mut self, name: &str, compare_fn: CompareFn) {
+    pub fn set_comparator(&mut self, name: impl CStrLike, compare_fn: CompareFn) {
         let cb = Box::new(ComparatorCallback {
-            name: CString::new(name.as_bytes()).unwrap(),
+            name: name.into_c_string().unwrap(),
             f: compare_fn,
         });
 

src/ffi_util.rs

@@ -80,3 +80,155 @@ macro_rules! ffi_try_impl {
         result
     }};
 }
+
+/// Value which can be converted into a C string.
+///
+/// The trait is used as argument to functions which wish to accept either
+/// [`&str`] or [`&CStr`] arguments while internally need to interact with
+/// C APIs.  Accepting [`&str`] may be more convenient for users but requires
+/// conversion into [`CString`] internally which requires allocation.  With this
+/// trait, latency-conscious users may choose to prepare `CStr` in advance and
+/// then pass it directly without having to incur the conversion cost.
+///
+/// To use the trait, function should accept `impl CStrLike` and after baking
+/// the argument (with [`CStrLike::bake`] method) it can use it as a `&CStr`
+/// (since the baked result dereferences into `CStr`).
+///
+/// # Example
+///
+/// ```
+/// use std::ffi::{CStr, CString};
+/// use rocksdb::CStrLike;
+///
+/// fn strlen(arg: impl CStrLike) -> std::result::Result<usize, String> {
+///     let baked = arg.bake().map_err(|err| err.to_string())?;
+///     Ok(unsafe { libc::strlen(baked.as_ptr()) })
+/// }
+///
+/// const FOO: &str = "foo";
+/// const BAR: &CStr = unsafe { CStr::from_bytes_with_nul_unchecked(b"bar\0") };
+///
+/// assert_eq!(Ok(3), strlen(FOO));
+/// assert_eq!(Ok(3), strlen(BAR));
+/// ```
+pub trait CStrLike {
+    type Baked: std::ops::Deref<Target = CStr>;
+    type Error: std::fmt::Debug + std::fmt::Display;
+
+    /// Bakes self into value which can be freely converted into [`&CStr`].
+    ///
+    /// This may require allocation and may fail if `self` has invalid value.
+    fn bake(self) -> Result<Self::Baked, Self::Error>;
+
+    /// Consumers and converts value into an owned [`CString`].
+    ///
+    /// If `Self` is already a `CString` simply returns it; if it’s a reference
+    /// to a `CString` then the value is cloned.  In other cases this may
+    /// require allocation and may fail if `self` has invalid value.
+    fn into_c_string(self) -> Result<CString, Self::Error>;
+}
+
+impl CStrLike for &str {
+    type Baked = CString;
+    type Error = std::ffi::NulError;
+
+    fn bake(self) -> Result<Self::Baked, Self::Error> {
+        CString::new(self)
+    }
+    fn into_c_string(self) -> Result<CString, Self::Error> {
+        CString::new(self)
+    }
+}
+
+// This is redundant for the most part and exists so that `foo(&string)` (where
+// `string: String` works just as if `foo` took `arg: &str` argument.
+impl CStrLike for &String {
+    type Baked = CString;
+    type Error = std::ffi::NulError;
+
+    fn bake(self) -> Result<Self::Baked, Self::Error> {
+        CString::new(self.as_bytes())
+    }
+    fn into_c_string(self) -> Result<CString, Self::Error> {
+        CString::new(self.as_bytes())
+    }
+}
+
+impl CStrLike for &CStr {
+    type Baked = Self;
+    type Error = std::convert::Infallible;
+
+    fn bake(self) -> Result<Self::Baked, Self::Error> {
+        Ok(self)
+    }
+    fn into_c_string(self) -> Result<CString, Self::Error> {
+        Ok(self.to_owned())
+    }
+}
+
+// This exists so that if caller constructs a `CString` they can pass it into
+// the function accepting `CStrLike` argument.  Some of such functions may take
+// the argument whereas otherwise they would need to allocated a new owned
+// object.
+impl CStrLike for CString {
+    type Baked = CString;
+    type Error = std::convert::Infallible;
+
+    fn bake(self) -> Result<Self::Baked, Self::Error> {
+        Ok(self)
+    }
+    fn into_c_string(self) -> Result<CString, Self::Error> {
+        Ok(self)
+    }
+}
+
+// This is redundant for the most part and exists so that `foo(&cstring)` (where
+// `string: CString` works just as if `foo` took `arg: &CStr` argument.
+impl<'a> CStrLike for &'a CString {
+    type Baked = &'a CStr;
+    type Error = std::convert::Infallible;
+
+    fn bake(self) -> Result<Self::Baked, Self::Error> {
+        Ok(self)
+    }
+    fn into_c_string(self) -> Result<CString, Self::Error> {
+        Ok(self.clone())
+    }
+}
+
+#[test]
+fn test_c_str_like_bake() {
+    fn test<S: CStrLike>(value: S) -> Result<usize, S::Error> {
+        value
+            .bake()
+            .map(|value| unsafe { libc::strlen(value.as_ptr()) })
+    }
+
+    assert_eq!(Ok(3), test("foo")); // &str
+    assert_eq!(Ok(3), test(&String::from("foo"))); // String
+    assert_eq!(Ok(3), test(CString::new("foo").unwrap().as_ref())); // &CStr
+    assert_eq!(Ok(3), test(&CString::new("foo").unwrap())); // &CString
+    assert_eq!(Ok(3), test(CString::new("foo").unwrap())); // CString
+
+    assert_eq!(3, test("foo\0bar").err().unwrap().nul_position());
+}
+
+#[test]
+fn test_c_str_like_into() {
+    fn test<S: CStrLike>(value: S) -> Result<CString, S::Error> {
+        value.into_c_string()
+    }
+
+    let want = CString::new("foo").unwrap();
+
+    assert_eq!(Ok(want.clone()), test("foo")); // &str
+    assert_eq!(Ok(want.clone()), test(&String::from("foo"))); // &String
+    assert_eq!(
+        Ok(want.clone()),
+        test(CString::new("foo").unwrap().as_ref())
+    ); // &CStr
+    assert_eq!(Ok(want.clone()), test(&CString::new("foo").unwrap())); // &CString
+    assert_eq!(Ok(want), test(CString::new("foo").unwrap())); // CString
+
+    assert_eq!(3, test("foo\0bar").err().unwrap().nul_position());
+}

src/lib.rs

@@ -112,6 +112,7 @@ pub use crate::{
         UniversalCompactOptions, UniversalCompactionStopStyle, WriteOptions,
     },
     db_pinnable_slice::DBPinnableSlice,
+    ffi_util::CStrLike,
     merge_operator::MergeOperands,
     perf::{PerfContext, PerfMetric, PerfStatsLevel},
     slice_transform::SliceTransform,

src/slice_transform.rs

@@ -17,7 +17,7 @@ use std::slice;
 
 use libc::{c_char, c_void, size_t};
 
-use crate::ffi;
+use crate::{ffi, ffi_util::CStrLike};
 
 /// A `SliceTransform` is a generic pluggable way of transforming one string
 /// to another. Its primary use-case is in configuring rocksdb
@@ -35,12 +35,12 @@ pub struct SliceTransform {
 
 impl SliceTransform {
     pub fn create(
-        name: &str,
+        name: impl CStrLike,
         transform_fn: TransformFn,
         in_domain_fn: Option<InDomainFn>,
     ) -> SliceTransform {
         let cb = Box::into_raw(Box::new(TransformCallback {
-            name: CString::new(name.as_bytes()).unwrap(),
+            name: name.into_c_string().unwrap(),
             transform_fn,
             in_domain_fn,
         }));

tests/test_property.rs

@@ -25,7 +25,20 @@ fn property_test() {
     {
         let db = DB::open_default(&n).unwrap();
         let value = db.property_value(properties::STATS).unwrap().unwrap();
+        assert!(value.contains("Stats"));
+    }
 
+    {
+        let db = DB::open_default(&n).unwrap();
+        let prop_name = std::ffi::CString::new(properties::STATS).unwrap();
+        let value = db.property_value(prop_name).unwrap().unwrap();
+        assert!(value.contains("Stats"));
+    }
+
+    {
+        let db = DB::open_default(&n).unwrap();
+        let prop_name = std::ffi::CString::new(properties::STATS).unwrap();
+        let value = db.property_value(prop_name.as_ref()).unwrap().unwrap();
         assert!(value.contains("Stats"));
     }
 }