src/stdsimd/crates/assert-instr-macro/src/lib.rs

   1 //! Implementation of the `#[assert_instr]` macro
   2 //!
   3 //! This macro is used when testing the `stdsimd` crate and is used to generate
   4 //! test cases to assert that functions do indeed contain the instructions that
   5 //! we're expecting them to contain.
   6 //!
   7 //! The procedural macro here is relatively simple, it simply appends a
   8 //! `#[test]` function to the original token stream which asserts that the
   9 //! function itself contains the relevant instruction.
  10
  11 extern crate proc_macro;
  12 extern crate proc_macro2;
  13 #[macro_use]
  14 extern crate quote;
  15 extern crate syn;
  16
  17 use proc_macro2::TokenStream;
  18 use quote::ToTokens;
  19
  20 #[proc_macro_attribute]
  21 pub fn assert_instr(
  22     attr: proc_macro::TokenStream,
  23     item: proc_macro::TokenStream,
  24 ) -> proc_macro::TokenStream {
  25     let invoc = match syn::parse::<Invoc>(attr) {
  26         Ok(s) => s,
  27         Err(e) => return e.to_compile_error().into(),
  28     };
  29     let item = match syn::parse::<syn::Item>(item) {
  30         Ok(s) => s,
  31         Err(e) => return e.to_compile_error().into(),
  32     };
  33     let func = match item {
  34         syn::Item::Fn(ref f) => f,
  35         _ => panic!("must be attached to a function"),
  36     };
  37
  38     let instr = &invoc.instr;
  39     let name = &func.ident;
  40
  41     // Disable assert_instr for x86 targets compiled with avx enabled, which
  42     // causes LLVM to generate different intrinsics that the ones we are
  43     // testing for.
  44     let disable_assert_instr = std::env::var("STDSIMD_DISABLE_ASSERT_INSTR").is_ok();
  45
  46     let instr_str = instr
  47         .replace('.', "_")
  48         .replace('/', "_")
  49         .replace(':', "_")
  50         .replace(|c: char| c.is_whitespace(), "");
  51     let assert_name = syn::Ident::new(&format!("assert_{}_{}", name, instr_str), name.span());
  52     let shim_name = syn::Ident::new(&format!("{}_shim_{}", name, instr_str), name.span());
  53     let mut inputs = Vec::new();
  54     let mut input_vals = Vec::new();
  55     let ret = &func.decl.output;
  56     for arg in func.decl.inputs.iter() {
  57         let capture = match *arg {
  58             syn::FnArg::Captured(ref c) => c,
  59             ref v => panic!(
  60                 "arguments must not have patterns: `{:?}`",
  61                 v.clone().into_token_stream()
  62             ),
  63         };
  64         let ident = match capture.pat {
  65             syn::Pat::Ident(ref i) => &i.ident,
  66             _ => panic!("must have bare arguments"),
  67         };
  68         if let Some(&(_, ref tts)) = invoc.args.iter().find(|a| *ident == a.0) {
  69             input_vals.push(quote! { #tts });
  70         } else {
  71             inputs.push(capture);
  72             input_vals.push(quote! { #ident });
  73         }
  74     }
  75
  76     let attrs = func
  77         .attrs
  78         .iter()
  79         .filter(|attr| {
  80             attr.path
  81                 .segments
  82                 .first()
  83                 .expect("attr.path.segments.first() failed")
  84                 .value()
  85                 .ident
  86                 .to_string()
  87                 .starts_with("target")
  88         })
  89         .collect::<Vec<_>>();
  90     let attrs = Append(&attrs);
  91
  92     // Use an ABI on Windows that passes SIMD values in registers, like what
  93     // happens on Unix (I think?) by default.
  94     let abi = if cfg!(windows) {
  95         syn::LitStr::new("vectorcall", proc_macro2::Span::call_site())
  96     } else {
  97         syn::LitStr::new("C", proc_macro2::Span::call_site())
  98     };
  99     let shim_name_str = format!("{}{}", shim_name, assert_name);
 100     let to_test = quote! {
 101         #attrs
 102         #[no_mangle]
 103         unsafe extern #abi fn #shim_name(#(#inputs),*) #ret {
 104             // The compiler in optimized mode by default runs a pass called
 105             // "mergefunc" where it'll merge functions that look identical.
 106             // Turns out some intrinsics produce identical code and they're
 107             // folded together, meaning that one just jumps to another. This
 108             // messes up our inspection of the disassembly of this function and
 109             // we're not a huge fan of that.
 110             //
 111             // To thwart this pass and prevent functions from being merged we
 112             // generate some code that's hopefully very tight in terms of
 113             // codegen but is otherwise unique to prevent code from being
 114             // folded.
 115             ::stdsimd_test::_DONT_DEDUP = #shim_name_str;
 116             #name(#(#input_vals),*)
 117         }
 118     };
 119
 120     // If instruction tests are disabled avoid emitting this shim at all, just
 121     // return the original item without our attribute.
 122     if !cfg!(optimized) || disable_assert_instr {
 123         return (quote! { #item }).into();
 124     }
 125
 126     let tts: TokenStream = quote! {
 127         #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
 128         #[cfg_attr(not(target_arch = "wasm32"), test)]
 129         #[allow(non_snake_case)]
 130         fn #assert_name() {
 131             #to_test
 132
 133             ::stdsimd_test::assert(#shim_name as usize,
 134                                    stringify!(#shim_name),
 135                                    #instr);
 136         }
 137     };
 138     // why? necessary now to get tests to work?
 139     let tts: TokenStream = tts.to_string().parse().expect("cannot parse tokenstream");
 140
 141     let tts: TokenStream = quote! {
 142         #item
 143         #tts
 144     };
 145     tts.into()
 146 }
 147
 148 struct Invoc {
 149     instr: String,
 150     args: Vec<(syn::Ident, syn::Expr)>,
 151 }
 152
 153 impl syn::parse::Parse for Invoc {
 154     fn parse(input: syn::parse::ParseStream) -> syn::parse::Result<Self> {
 155         use syn::{ext::IdentExt, Token};
 156
 157         let mut instr = String::new();
 158         while !input.is_empty() {
 159             if input.parse::<Token![,]>().is_ok() {
 160                 break;
 161             }
 162             if let Ok(ident) = syn::Ident::parse_any(input) {
 163                 instr.push_str(&ident.to_string());
 164                 continue;
 165             }
 166             if input.parse::<Token![.]>().is_ok() {
 167                 instr.push_str(".");
 168                 continue;
 169             }
 170             if let Ok(s) = input.parse::<syn::LitStr>() {
 171                 instr.push_str(&s.value());
 172                 continue;
 173             }
 174             println!("{:?}", input.cursor().token_stream());
 175             return Err(input.error("expected an instruction"));
 176         }
 177         if instr.is_empty() {
 178             return Err(input.error("expected an instruction before comma"));
 179         }
 180         let mut args = Vec::new();
 181         while !input.is_empty() {
 182             let name = input.parse::<syn::Ident>()?;
 183             input.parse::<Token![=]>()?;
 184             let expr = input.parse::<syn::Expr>()?;
 185             args.push((name, expr));
 186
 187             if input.parse::<Token![,]>().is_err() {
 188                 if !input.is_empty() {
 189                     return Err(input.error("extra tokens at end"));
 190                 }
 191                 break;
 192             }
 193         }
 194         Ok(Self { instr, args })
 195     }
 196 }
 197
 198 struct Append<T>(T);
 199
 200 impl<T> quote::ToTokens for Append<T>
 201 where
 202     T: Clone + IntoIterator,
 203     T::Item: quote::ToTokens,
 204 {
 205     fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
 206         for item in self.0.clone() {
 207             item.to_tokens(tokens);
 208         }
 209     }
 210 }