Skip to content

Commit

Permalink
Remove Clone requirement for argument types
Browse files Browse the repository at this point in the history 
- 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
  • Loading branch information
@bnaecker
bnaecker committed Dec 17, 2024
1 parent 7a168c6 commit f941a0c
Show file tree
Hide file tree
Showing 4 changed files with 50 additions and 102 deletions.
18 changes: 0 additions & 18 deletions tests/compile-errors/src/different-serializable-type.stderr
View file
Original file line number Diff line number Diff line change
@@ -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
|
Expand Down
12 changes: 4 additions & 8 deletions usdt-attr-macro/src/lib.rs
View file
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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)*
};
}
Expand Down Expand Up @@ -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>()
}
}
}
Expand Down
117 changes: 42 additions & 75 deletions usdt-impl/src/common.rs
View file
Original file line number Diff line number Diff line change
@@ -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.
Expand All @@ -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) => {
Expand Down Expand Up @@ -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),*); };
}
}

Expand Down Expand Up @@ -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)(); },
}
}

Expand Down Expand Up @@ -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 {
Expand All @@ -243,7 +223,6 @@ pub(crate) fn build_probe_macro(
};
($args_lambda:expr) => {
{
#type_check_block
#impl_block
}
};
Expand All @@ -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 = &[
Expand All @@ -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());
}

Expand All @@ -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();
};
Expand Down
5 changes: 4 additions & 1 deletion usdt-impl/src/linker.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}" };
Expand Down Expand Up @@ -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,
Expand All @@ -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)]
Expand Down

0 comments on commit f941a0c

Please sign in to comment.