Remove Clone requirement for argument types

- Rework construction of the internal type-checking function to remove
  the clone requirement. Construct and call the type-check-function
  right at the site we already expand the probe argument closure, rather
  than outside. This means we don't need to expand things twice, and
  don't need to clone the closure, and so don't need the argument types
  to be cloneable either.
- Fixup tests for new method, add some comments.
- Update trybuild tests
- Closes #136

Author: bnaecker

.github/workflows/rust.yml

@@ -81,7 +81,6 @@ jobs:
           --exclude does-it-work
           --exclude test-json
           --exclude test-unique-id
-          --exclude compile-errors
 
   stable-test-no-op:
     name: Test with probes disabled

tests/compile-errors/src/different-serializable-type.stderr

@@ -1,21 +1,3 @@
-error[E0277]: the trait bound `Expected: Clone` is not satisfied
-  --> src/different-serializable-type.rs:31:20
-   |
-31 |     fn my_probe(_: Expected) {}
-   |                    ^^^^^^^^ the trait `Clone` is not implemented for `Expected`
-   |
-note: required by a bound in `usdt_types_must_be_clone_and_serialize`
-  --> src/different-serializable-type.rs:28:1
-   |
-28 | #[usdt::provider]
-   | ^^^^^^^^^^^^^^^^^ required by this bound in `usdt_types_must_be_clone_and_serialize`
-   = note: this error originates in the attribute macro `usdt::provider` (in Nightly builds, run with -Z macro-backtrace for more info)
-help: consider annotating `Expected` with `#[derive(Clone)]`
-   |
-19 + #[derive(Clone)]
-20 | struct Expected {
-   |
-
 error[E0277]: the trait bound `Different: Borrow<Expected>` is not satisfied
   --> src/different-serializable-type.rs:28:1
    |

usdt-attr-macro/src/lib.rs

@@ -1,6 +1,6 @@
 //! Generate USDT probes from an attribute macro
 
-// Copyright 2021 Oxide Computer Company
+// Copyright 2024 Oxide Computer Company
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -144,7 +144,7 @@ fn generate_provider_item(
         quote! {
             const _: fn() = || {
                 #(#use_statements)*
-                fn usdt_types_must_be_clone_and_serialize<T: ?Sized + Clone + ::serde::Serialize>() {}
+                fn usdt_types_must_be_serialize<T: ?Sized + ::serde::Serialize>() {}
                 #(#check_fns)*
             };
         }
@@ -269,16 +269,12 @@ fn build_serializable_check_function<T>(ident: &T, fn_index: usize, arg_index: u
 where
     T: quote::ToTokens,
 {
-    let fn_name = quote::format_ident!(
-        "usdt_types_must_be_clone_and_serialize_{}_{}",
-        fn_index,
-        arg_index
-    );
+    let fn_name = quote::format_ident!("usdt_types_must_be_serialize_{}_{}", fn_index, arg_index);
     quote! {
         fn #fn_name() {
             // #ident must be in scope here, because this function is defined in the same module as
             // the actual probe functions, and thus shares any imports the consumer wants.
-            usdt_types_must_be_clone_and_serialize::<#ident>()
+            usdt_types_must_be_serialize::<#ident>()
         }
     }
 }

usdt-impl/src/common.rs

@@ -1,6 +1,6 @@
 //! Shared code used in both the linker and no-linker implementations of this crate.
 
-// Copyright 2021 Oxide Computer Company
+// Copyright 2024 Oxide Computer Company
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -14,39 +14,29 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::{DataType, Provider};
+use crate::DataType;
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 
-// Construct function call that is used internally in the UDST-generated macros, to allow
-// compile-time type checking of the lambda arguments.
-pub fn generate_type_check(
+/// Construct a function to type-check the argument closure.
+///
+/// This constructs a function that is never called, but is used to ensure that
+/// the closure provided to each probe macro returns arguments of the right
+/// type.
+pub fn construct_type_check(
     provider_name: &str,
-    use_statements: &[syn::ItemUse],
     probe_name: &str,
+    use_statements: &[syn::ItemUse],
     types: &[DataType],
 ) -> TokenStream {
-    // If the probe has zero arguments, verify that the result of calling the closure is `()`
-    // Note that there's no need to clone the closure here, since () is Copy.
+    // If there are zero arguments, we need to make sure we can assign the
+    // result of the closure to ().
     if types.is_empty() {
         return quote! {
-            let __usdt_private_args_lambda = $args_lambda;
-            let _ = || {
-                let _: () = __usdt_private_args_lambda();
-            };
+            let _: () = ($args_lambda)();
         };
     }
-
-    // For one or more arguments, verify that we can unpack the closure into a tuple of type
-    // `(arg0, arg1, ...)`. We verify that we can pass those arguments to a type-check function
-    // that is _similar to_, but not exactly the probe function signature. In particular, we try to
-    // support passing things by value or reference, and take some form of reference to that thing.
-    // The mapping is generally:
-    //
-    // T or &T -> Borrow<T>
-    // Strings -> AsRef<str>
-    // [T; N] or &[T] -> AsRef<[T]>
-    let type_check_args = types
+    let type_check_params = types
         .iter()
         .map(|typ| match typ {
             DataType::Serializable(ty) => {
@@ -84,28 +74,22 @@ pub fn generate_type_check(
         })
         .collect::<Vec<_>>();
 
-    // Unpack the tuple from the closure to `args.0, args.1, ...`.
-    let expanded_lambda_args = (0..types.len())
+    // Create a list of arguments `arg.0`, `arg.1`, ... to pass to the check
+    // function.
+    let type_check_args = (0..types.len())
         .map(|i| {
             let index = syn::Index::from(i);
             quote! { args.#index }
         })
         .collect::<Vec<_>>();
 
-    let preamble = unpack_argument_lambda(types, /* clone = */ true);
-
-    let type_check_function =
-        format_ident!("__usdt_private_{}_{}_type_check", provider_name, probe_name);
+    let type_check_fn = format_ident!("__usdt_private_{}_{}_type_check", provider_name, probe_name);
     quote! {
-        let __usdt_private_args_lambda = $args_lambda;
         #[allow(unused_imports)]
         #(#use_statements)*
         #[allow(non_snake_case)]
-        fn #type_check_function (#(#type_check_args),*) { }
-        let _ = || {
-            #preamble
-            #type_check_function(#(#expanded_lambda_args),*);
-        };
+        fn #type_check_fn(#(#type_check_params),*) {}
+        let _ = || { #type_check_fn(#(#type_check_args),*); };
     }
 }
 
@@ -156,28 +140,23 @@ pub fn construct_probe_args(types: &[DataType]) -> (TokenStream, TokenStream) {
             (destructured_arg, register_arg)
         })
         .unzip();
-    let preamble = unpack_argument_lambda(types, /* clone = */ false);
+    let arg_lambda = unpack_argument_lambda(types);
     let unpacked_args = quote! {
-        #preamble
+        #arg_lambda
         #(#unpacked_args)*
     };
     let in_regs = quote! { #(#in_regs,)* };
     (unpacked_args, in_regs)
 }
 
-fn unpack_argument_lambda(types: &[DataType], clone: bool) -> TokenStream {
-    let maybe_clone = if clone {
-        quote! { .clone() }
-    } else {
-        quote! {}
-    };
+fn unpack_argument_lambda(types: &[DataType]) -> TokenStream {
     match types.len() {
-        // Don't bother with arguments if there are none.
-        0 => quote! { __usdt_private_args_lambda #maybe_clone (); },
+        // Don't bother with any closure if there are no arguments.
+        0 => quote! {},
         // Wrap a single argument in a tuple.
-        1 => quote! { let args = (__usdt_private_args_lambda #maybe_clone (),); },
+        1 => quote! { let args = (($args_lambda)(),); },
         // General case.
-        _ => quote! { let args = __usdt_private_args_lambda #maybe_clone (); },
+        _ => quote! { let args = ($args_lambda)(); },
     }
 }
 
@@ -218,17 +197,18 @@ fn asm_type_convert(typ: &DataType, input: TokenStream) -> (TokenStream, TokenSt
     }
 }
 
+/// Create the top-level probe macro.
+///
+/// This takes the implementation block constructed elsewhere, and builds out
+/// the actual macro users call in their code to fire the probe.
 pub(crate) fn build_probe_macro(
     config: &crate::CompileProvidersConfig,
-    provider: &Provider,
     probe_name: &str,
     types: &[DataType],
     impl_block: TokenStream,
 ) -> TokenStream {
     let module = config.module_ident();
     let macro_name = config.probe_ident(probe_name);
-    let type_check_block =
-        generate_type_check(&provider.name, &provider.use_statements, probe_name, types);
     let no_args_match = if types.is_empty() {
         quote! { () => { crate::#module::#macro_name!(|| ()) }; }
     } else {
@@ -243,7 +223,6 @@ pub(crate) fn build_probe_macro(
             };
             ($args_lambda:expr) => {
                 {
-                    #type_check_block
                     #impl_block
                 }
             };
@@ -263,20 +242,16 @@ mod tests {
     use dtrace_parser::Sign;
 
     #[test]
-    fn test_generate_type_check_empty() {
-        let types = &[];
+    fn test_construct_type_check_empty() {
         let expected = quote! {
-            let __usdt_private_args_lambda = $args_lambda;
-            let _ = || {
-                let _: () = __usdt_private_args_lambda();
-            };
+            let _ : () = ($args_lambda)();
         };
-        let block = generate_type_check("", &[], "", types);
+        let block = construct_type_check("", "", &[], &[]);
         assert_eq!(block.to_string(), expected.to_string());
     }
 
     #[test]
-    fn test_generate_type_check_native() {
+    fn test_construct_type_check_native() {
         let provider = "provider";
         let probe = "probe";
         let types = &[
@@ -290,80 +265,72 @@ mod tests {
             })),
         ];
         let expected = quote! {
-            let __usdt_private_args_lambda = $args_lambda;
             #[allow(unused_imports)]
             #[allow(non_snake_case)]
             fn __usdt_private_provider_probe_type_check(
                 _: impl ::std::borrow::Borrow<u8>,
                 _: impl ::std::borrow::Borrow<i64>
             ) { }
             let _ = || {
-                let args = __usdt_private_args_lambda.clone()();
                 __usdt_private_provider_probe_type_check(args.0, args.1);
             };
         };
-        let block = generate_type_check(provider, &[], probe, types);
+        let block = construct_type_check(provider, probe, &[], types);
         assert_eq!(block.to_string(), expected.to_string());
     }
 
     #[test]
-    fn test_generate_type_check_with_string() {
+    fn test_construct_type_check_with_string() {
         let provider = "provider";
         let probe = "probe";
         let types = &[DataType::Native(dtrace_parser::DataType::String)];
         let use_statements = vec![];
         let expected = quote! {
-            let __usdt_private_args_lambda = $args_lambda;
             #[allow(unused_imports)]
             #[allow(non_snake_case)]
             fn __usdt_private_provider_probe_type_check(_: impl AsRef<str>) { }
             let _ = || {
-                let args = (__usdt_private_args_lambda.clone()(),);
                 __usdt_private_provider_probe_type_check(args.0);
             };
         };
-        let block = generate_type_check(provider, &use_statements, probe, types);
+        let block = construct_type_check(provider, probe, &use_statements, types);
         assert_eq!(block.to_string(), expected.to_string());
     }
 
     #[test]
-    fn test_generate_type_check_with_shared_slice() {
+    fn test_construct_type_check_with_shared_slice() {
         let provider = "provider";
         let probe = "probe";
         let types = &[DataType::Serializable(syn::parse_str("&[u8]").unwrap())];
         let use_statements = vec![];
         let expected = quote! {
-            let __usdt_private_args_lambda = $args_lambda;
             #[allow(unused_imports)]
             #[allow(non_snake_case)]
             fn __usdt_private_provider_probe_type_check(_: impl AsRef<[u8]>) { }
             let _ = || {
-                let args = (__usdt_private_args_lambda.clone()(),);
                 __usdt_private_provider_probe_type_check(args.0);
             };
         };
-        let block = generate_type_check(provider, &use_statements, probe, types);
+        let block = construct_type_check(provider, probe, &use_statements, types);
         assert_eq!(block.to_string(), expected.to_string());
     }
 
     #[test]
-    fn test_generate_type_check_with_custom_type() {
+    fn test_construct_type_check_with_custom_type() {
         let provider = "provider";
         let probe = "probe";
         let types = &[DataType::Serializable(syn::parse_str("MyType").unwrap())];
         let use_statements = vec![syn::parse2(quote! { use my_module::MyType; }).unwrap()];
         let expected = quote! {
-            let __usdt_private_args_lambda = $args_lambda;
             #[allow(unused_imports)]
             use my_module::MyType;
             #[allow(non_snake_case)]
             fn __usdt_private_provider_probe_type_check(_: impl ::std::borrow::Borrow<MyType>) { }
             let _ = || {
-                let args = (__usdt_private_args_lambda.clone()(),);
                 __usdt_private_provider_probe_type_check(args.0);
             };
         };
-        let block = generate_type_check(provider, &use_statements, probe, types);
+        let block = construct_type_check(provider, probe, &use_statements, types);
         assert_eq!(block.to_string(), expected.to_string());
     }
 
@@ -382,7 +349,7 @@ mod tests {
         let registers = ["x0", "x1"];
         let (args, regs) = construct_probe_args(types);
         let expected = quote! {
-            let args = __usdt_private_args_lambda();
+            let args = ($args_lambda)();
             let arg_0 = (*<_ as ::std::borrow::Borrow<*const u8>>::borrow(&args.0) as usize);
             let arg_1 = [(args.1.as_ref() as &str).as_bytes(), &[0_u8]].concat();
         };

usdt-impl/src/empty.rs

@@ -76,8 +76,17 @@ fn compile_probe(
     probe: &Probe,
     config: &crate::CompileProvidersConfig,
 ) -> TokenStream {
-    let impl_block = quote! { let _ = || (__usdt_private_args_lambda()) ; };
-    common::build_probe_macro(config, provider, &probe.name, &probe.types, impl_block)
+    let type_check_fn = common::construct_type_check(
+        &provider.name,
+        &probe.name,
+        &provider.use_statements,
+        &probe.types,
+    );
+    let impl_block = quote! {
+        let args = || { ($args_lambda)() };
+        #type_check_fn
+    };
+    common::build_probe_macro(config, &probe.name, &probe.types, impl_block)
 }
 
 pub fn register_probes() -> Result<(), crate::Error> {

usdt-impl/src/linker.rs

@@ -170,6 +170,8 @@ fn compile_probe(
         syn::parse2::<syn::FnArg>(quote! { _: #ty }).unwrap()
     });
     let (unpacked_args, in_regs) = common::construct_probe_args(types);
+    let type_check_fn =
+        common::construct_type_check(&provider.name, probe_name, &provider.use_statements, types);
 
     #[cfg(target_arch = "x86_64")]
     let call_instruction = quote! { "call {extern_probe_fn}" };
@@ -199,6 +201,7 @@ fn compile_probe(
         unsafe {
             if #is_enabled_fn() != 0 {
                 #unpacked_args
+                #type_check_fn
                 ::std::arch::asm!(
                     ".reference {typedefs}",
                     #call_instruction,
@@ -213,7 +216,7 @@ fn compile_probe(
         }
     };
 
-    common::build_probe_macro(config, provider, probe_name, types, impl_block)
+    common::build_probe_macro(config, probe_name, types, impl_block)
 }
 
 #[derive(Debug, Default, Clone)]