feat(stapsdt): Allow compiler to choose argument registers

tests/does-it-work/src/main.rs

@@ -242,21 +242,70 @@ mod tests {
                 semaphore_address
             );
 
-            // Verify the argument types
+            // Verify the argument types; the format of the line is eg.
+            // "Arguments: size@%reg size@%reg", but we cannot know which
+            // registers the compiler has chosen to use for arguments passing,
+            // so we have to instead split out the two arguments we expect to
+            // find for this probe specifically, check their sizes, and match
+            // the register strings ("%reg") against a list of possible values.
             let line = lines.next().expect("Expected a line containing arguments");
             let line = line.trim();
-            let arguments_line = if cfg!(target_arch = "x86_64") {
-                "Arguments: 1@%dil 8@%rsi"
-            } else if cfg!(target_arch = "aarch64") {
-                "Arguments: 1@x0 8@x1"
-            } else {
-                unreachable!("Unsupported Linux target architecture")
-            };
+            let (args_header, args_rest) = line
+                .split_once(" ")
+                .expect("Expected arguments line to have one space");
+            let (first_arg, second_arg) = args_rest
+                .split_once(" ")
+                .expect("Expected arguments line to have two spaces");
+            let (first_arg_size, first_arg_register) = first_arg
+                .split_once("@")
+                .expect("Expected first argument to have @ sign");
+            let (second_arg_size, second_arg_register) = second_arg
+                .split_once("@")
+                .expect("Expected first argument to have @ sign");
+            assert_eq!(
+                first_arg_size, "1",
+                "First argument size appears incorrect: {}",
+                first_arg_size
+            );
             assert_eq!(
-                line, arguments_line,
+                second_arg_size, "8",
+                "Second argument size appears incorrect: {}",
+                second_arg_size
+            );
+            assert_eq!(
+                args_header, "Arguments:",
                 "Arguments line appears incorrect: {}",
                 line
             );
+            if cfg!(target_arch = "x86_64") {
+                assert!(
+                    ["%ah", "%al", "%bh", "%bl", "%ch", "%cl", "%dh", "%dl", "%dil", "%sil"]
+                        .contains(&first_arg_register),
+                    "First argument register appears incorrect: {}",
+                    first_arg_register
+                );
+                assert!(
+                    ["%rax", "%rbx", "%rcx", "%rdx", "%rdi", "%rsi"].contains(&second_arg_register),
+                    "Second argument register appears incorrect: {}",
+                    second_arg_register
+                );
+            } else if cfg!(target_arch = "aarch64") {
+                fn is_valid_register(reg: &str) -> bool {
+                    reg.starts_with('x') && str::parse::<u8>(&reg[1..]).is_ok_and(|val| val <= 30)
+                }
+                assert!(
+                    is_valid_register(first_arg_register),
+                    "First argument register appears incorrect: {}",
+                    first_arg_register
+                );
+                assert!(
+                    is_valid_register(second_arg_register),
+                    "Second argument register appears incorrect: {}",
+                    second_arg_register
+                );
+            } else {
+                unreachable!("Unsupported Linux target architecture")
+            }
 
             thr.join().expect("Failed to join test runner thread");
         }

usdt-impl/src/common.rs

@@ -101,6 +101,26 @@ fn shared_slice_elem_type(reference: &syn::TypeReference) -> Option<&syn::Type>
     }
 }
 
+/// Returns true if the DataType requires the special "reg_byte" register class
+/// to be used for the `asm!` macro arguments passing. This only happens for
+/// STAPSDT probes on x86.
+#[cfg(usdt_backend_stapsdt)]
+fn use_reg_byte(typ: &DataType) -> bool {
+    #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
+    {
+        use dtrace_parser::{BitWidth, Integer};
+        matches!(
+            typ,
+            DataType::Native(dtrace_parser::DataType::Integer(Integer {
+                sign: _,
+                width: BitWidth::Bit8
+            }))
+        )
+    }
+    #[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))]
+    false
+}
+
 // Return code to destructure a probe arguments into identifiers, and to pass those to ASM
 // registers.
 pub fn construct_probe_args(types: &[DataType]) -> (TokenStream, TokenStream) {
@@ -135,8 +155,28 @@ pub fn construct_probe_args(types: &[DataType]) -> (TokenStream, TokenStream) {
                 let #arg = #value;
             };
             // Here, we convert the argument to store it within a register.
+            #[cfg(usdt_backend_stapsdt)]
+            let register_arg = {
+                // In SystemTap probes, the arguments can be passed freely in
+                // any registers without regard to standard function call ABIs.
+                // We thus do not need the register names in the STAPSDT
+                // backend. Intead, we use the "name = in(reg) arg" pattern to
+                // bind each argument into a local name (we shadow the original
+                // argument name here): this name can then be used by the
+                // `asm!` macro to refer to the argument "by register"; the
+                // compiler will fill in the actual register name at each
+                // reference point. This does away with a need for the compiler
+                // to generate code moving arugments around in registers before
+                // a probe site.
+                let _ = reg;
+                if use_reg_byte(typ) {
+                    quote! { #arg = in(reg_byte) (#arg #at_use) }
+                } else {
+                    quote! { #arg = in(reg) (#arg #at_use) }
+                }
+            };
+            #[cfg(not(usdt_backend_stapsdt))]
             let register_arg = quote! { in(#reg) (#arg #at_use) };
-
             (destructured_arg, register_arg)
         })
         .unzip();
@@ -164,7 +204,7 @@ pub fn call_argument_closure(types: &[DataType]) -> TokenStream {
 // Convert a supported data type to 1. a type to store for the duration of the
 // probe invocation and 2. a transformation for compatibility with an asm
 // register.
-fn asm_type_convert(typ: &DataType, input: TokenStream) -> (TokenStream, TokenStream) {
+pub(crate) fn asm_type_convert(typ: &DataType, input: TokenStream) -> (TokenStream, TokenStream) {
     match typ {
         DataType::Serializable(_) => (
             // Convert the input to JSON. This is a fallible operation, however, so we wrap the
@@ -189,10 +229,15 @@ fn asm_type_convert(typ: &DataType, input: TokenStream) -> (TokenStream, TokenSt
         ),
         DataType::Native(_) => {
             let ty = typ.to_rust_type();
-            (
-                quote! { (*<_ as ::std::borrow::Borrow<#ty>>::borrow(&#input) as usize) },
-                quote! {},
-            )
+            #[cfg(not(usdt_backend_stapsdt))]
+            let value = quote! { (*<_ as ::std::borrow::Borrow<#ty>>::borrow(&#input) as usize) };
+            // For STAPSDT probes, we cannot "widen" the data type at the
+            // interface, as we've left the register choice up to the
+            // compiler and the compiler doesn't like mismatched register
+            // classes and types for single-byte values (reg_byte vs usize).
+            #[cfg(usdt_backend_stapsdt)]
+            let value = quote! { (*<_ as ::std::borrow::Borrow<#ty>>::borrow(&#input)) };
+            (value, quote! {})
         }
         DataType::UniqueId => (quote! { #input.as_u64() as usize }, quote! {}),
     }
@@ -349,20 +394,34 @@ mod tests {
         #[cfg(target_arch = "aarch64")]
         let registers = ["x0", "x1"];
         let (args, regs) = construct_probe_args(types);
+        #[cfg(not(usdt_backend_stapsdt))]
         let expected = quote! {
             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();
         };
+        #[cfg(usdt_backend_stapsdt)]
+        let expected = quote! {
+            let args = ($args_lambda)();
+            let arg_0 = (*<_ as ::std::borrow::Borrow<*const u8>>::borrow(&args.0));
+            let arg_1 = [(args.1.as_ref() as &str).as_bytes(), &[0_u8]].concat();
+        };
         assert_eq!(args.to_string(), expected.to_string());
 
         for (i, (expected, actual)) in registers
             .iter()
             .zip(regs.to_string().split(','))
             .enumerate()
         {
+            // We don't need the register names for STAPSDT probes.
+            #[cfg(usdt_backend_stapsdt)]
+            let _ = expected;
+
             let reg = actual.replace(' ', "");
+            #[cfg(not(usdt_backend_stapsdt))]
             let expected = format!("in(\"{}\")(arg_{}", expected, i);
+            #[cfg(usdt_backend_stapsdt)]
+            let expected = format!("arg_{i}=in(reg)(arg_{i}");
             assert!(
                 reg.starts_with(&expected),
                 "reg: {}; expected {}",
@@ -382,10 +441,12 @@ mod tests {
             })),
             TokenStream::from_str("foo").unwrap(),
         );
-        assert_eq!(
-            out.to_string(),
-            quote! {(*<_ as ::std::borrow::Borrow<u8>>::borrow(&foo) as usize)}.to_string()
-        );
+        #[cfg(usdt_backend_stapsdt)]
+        let out_expected = quote! {(*<_ as ::std::borrow::Borrow<u8>>::borrow(&foo))}.to_string();
+        #[cfg(not(usdt_backend_stapsdt))]
+        let out_expected =
+            quote! {(*<_ as ::std::borrow::Borrow<u8>>::borrow(&foo) as usize)}.to_string();
+        assert_eq!(out.to_string(), out_expected);
         assert_eq!(post.to_string(), quote! {}.to_string());
 
         let (out, post) = asm_type_convert(

usdt-impl/src/stapsdt.rs

@@ -22,6 +22,7 @@
 #[path = "stapsdt/args.rs"]
 mod args;
 
+use crate::common::construct_probe_args;
 use crate::{common, DataType};
 use crate::{Probe, Provider};
 use args::format_argument;
@@ -192,7 +193,7 @@ fn compile_probe(
     probe: &Probe,
     config: &crate::CompileProvidersConfig,
 ) -> TokenStream {
-    let (unpacked_args, in_regs) = common::construct_probe_args(&probe.types);
+    let (unpacked_args, in_regs) = construct_probe_args(&probe.types);
     let probe_rec = emit_probe_record(&provider.name, &probe.name, Some(&probe.types));
     let type_check_fn = common::construct_type_check(
         &provider.name,
@@ -202,6 +203,33 @@ fn compile_probe(
     );
 
     let sema_name = format_ident!("__usdt_sema_{}_{}", provider.name, probe.name);
+    let options = if cfg!(target_arch = "x86_64") || cfg!(target_arch = "x86") {
+        // STAPSDT probes contain an "arguments" string which contains the
+        // size, type, and location of each argument. This string is expected
+        // to be in the AT&T syntax: we change the syntax for x86 only, as only
+        // there the syntax effects register naming. The rest of our inline
+        // assembly here is the same in both AT&T and Intel syntax, so we can
+        // freely change the syntax without changing the generating code.
+        // The arguments string on x86 looks like this:
+        // * "2@%ax -4@%edi 8f@%rsi"
+        // and in our inline assembly it is as follows:
+        // * "2@{arg0} -4@{arg1} 8f@{arg2}"
+        // The argument size and type is explicitly named on the left side of
+        // the "@" sign, but the register is given by argument name instead of
+        // explicitly naming eg. "%ax". This gives the compiler the freedom to
+        // choose for itself where it wants to place the arguments. The only
+        // thing we need to make sure of is that the argument register strings
+        // are in the AT&T syntax: in Intel syntax the the "%" character would
+        // be missing from the register names.
+        // Note that we could manually fill in the "%" character and still use
+        // Intel syntax, but that will break down if Rust's inline assembly
+        // ever gets memory operands. Memory operands in the syntax look like:
+        // * "8@0x18(%rsp)"
+        // for "8-byte unsigned integer at stack + 0x18".
+        quote! { options(att_syntax, nomem, nostack, preserves_flags) }
+    } else {
+        quote! { options(nomem, nostack, preserves_flags) }
+    };
     let impl_block = quote! {
         unsafe extern "C" {
             // Note: C libraries use a struct containing an unsigned short
@@ -226,7 +254,7 @@ fn compile_probe(
                     "990:   nop",
                     #probe_rec,
                     #in_regs
-                    options(nomem, nostack, preserves_flags)
+                    #options
                 );
             }
         }