add cstr support

Author: conradludgate

dtrace-parser/src/lib.rs

@@ -156,6 +156,7 @@ impl Integer {
 pub enum DataType {
     Integer(Integer),
     Pointer(Integer),
+    CString,
     String,
 }
 
@@ -246,7 +247,7 @@ impl DataType {
         match self {
             DataType::Integer(int) => int.to_c_type(),
             DataType::Pointer(int) => format!("{}*", int.to_c_type()),
-            DataType::String => String::from("char*"),
+            DataType::CString | DataType::String => String::from("char*"),
         }
     }
 
@@ -255,7 +256,7 @@ impl DataType {
         match self {
             DataType::Integer(int) => int.to_rust_ffi_type(),
             DataType::Pointer(int) => format!("*const {}", int.to_rust_ffi_type()),
-            DataType::String => format!("*const {RUST_TYPE_PREFIX}char"),
+            DataType::CString | DataType::String => format!("*const {RUST_TYPE_PREFIX}char"),
         }
     }
 
@@ -265,6 +266,7 @@ impl DataType {
             DataType::Integer(int) => int.to_rust_type(),
             DataType::Pointer(int) => format!("*const {}", int.to_rust_type()),
             DataType::String => String::from("&str"),
+            DataType::CString => String::from("&::core::ffi::CStr"),
         }
     }
 }

probe-test-attr/src/main.rs

@@ -82,6 +82,9 @@ mod test {
 
     /// Constant pointers to integer types are also supported
     fn work_with_pointer(_buffer: *const u8, _: u64) {}
+
+    /// Avoid expensive string allocation with cstring support
+    fn cstring(_: &CStr, _: CString) {}
 }
 
 fn main() {
@@ -106,5 +109,6 @@ fn main() {
         test::arg_as_tuple!(|| (arg.x, &arg.buffer[..]));
         test::not_json_serializable!(|| Whoops::NoBueno(0));
         test::work_with_pointer!(|| (buffer.as_ptr(), buffer.len() as u64));
+        test::cstring!(|| (c"hello world", c"and when owned".to_owned()));
     }
 }

usdt-attr-macro/src/lib.rs

@@ -303,6 +303,8 @@ fn is_simple_type(ident: &syn::Ident) -> bool {
             | "i64"
             | "String"
             | "str"
+            | "CString"
+            | "CStr"
             | "usize"
             | "isize"
     )
@@ -363,6 +365,8 @@ fn data_type_from_path(path: &syn::Path, pointer: bool) -> DataType {
         }))
     } else if path.is_ident("String") || path.is_ident("str") {
         DataType::Native(DType::String)
+    } else if path.is_ident("CString") || path.is_ident("CStr") {
+        DataType::Native(DType::CString)
     } else if path.is_ident("isize") {
         DataType::Native(variant(Integer {
             sign: Sign::Signed,
@@ -462,6 +466,10 @@ mod tests {
     #[case("String", DType::String)]
     #[case("&&str", DType::String)]
     #[case("&String", DType::String)]
+    #[case("&CStr", DType::CString)]
+    #[case("CString", DType::CString)]
+    #[case("&&CStr", DType::CString)]
+    #[case("&CString", DType::CString)]
     fn test_parse_probe_argument_native(#[case] name: &str, #[case] ty: dtrace_parser::DataType) {
         let arg = syn::parse_str(name).unwrap();
         let out = parse_probe_argument(&arg, 0, 0).unwrap();

usdt-impl/src/common.rs

@@ -67,6 +67,9 @@ pub fn construct_type_check(
                 }
             }
             DataType::Native(dtrace_parser::DataType::String) => quote! { _: impl AsRef<str> },
+            DataType::Native(dtrace_parser::DataType::CString) => {
+                quote! { _: impl AsRef<::core::ffi::CStr> }
+            }
             _ => {
                 let arg = typ.to_rust_type();
                 quote! { _: impl ::std::borrow::Borrow<#arg> }
@@ -181,6 +184,10 @@ fn asm_type_convert(typ: &DataType, input: TokenStream) -> (TokenStream, TokenSt
             },
             quote! { .as_ptr() as usize },
         ),
+        DataType::Native(dtrace_parser::DataType::CString) => (
+            quote! { #input.as_ref() as &::core::ffi::CStr },
+            quote! { .as_ptr() as usize },
+        ),
         DataType::Native(dtrace_parser::DataType::String) => (
             quote! {
                 [(#input.as_ref() as &str).as_bytes(), &[0_u8]].concat()
@@ -397,5 +404,15 @@ mod tests {
             quote! { [(foo.as_ref() as &str).as_bytes(), &[0_u8]].concat() }.to_string()
         );
         assert_eq!(post.to_string(), quote! { .as_ptr() as usize }.to_string());
+
+        let (out, post) = asm_type_convert(
+            &DataType::Native(dtrace_parser::DataType::CString),
+            TokenStream::from_str("foo").unwrap(),
+        );
+        assert_eq!(
+            out.to_string(),
+            quote! { foo.as_ref() as &::core::ffi::CStr }.to_string()
+        );
+        assert_eq!(post.to_string(), quote! { .as_ptr() as usize }.to_string());
     }
 }