#![doc = include_str!("../README.md")]
use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use syn::spanned::Spanned;
macro_rules! unsupported_if_some {
($value:expr) => {
if let Some(value) = $value {
return syn::Error::new(value.span(), "unsupported by #[unsafe_target_feature(...)]")
.into_compile_error()
.into();
}
};
}
macro_rules! unsupported {
($value: expr) => {
return syn::Error::new(
$value.span(),
"unsupported by #[unsafe_target_feature(...)]",
)
.into_compile_error()
.into()
};
}
mod kw {
syn::custom_keyword!(conditional);
}
enum SpecializeArg {
LitStr(syn::LitStr),
Conditional(Conditional),
}
impl SpecializeArg {
fn lit(&self) -> &syn::LitStr {
match self {
SpecializeArg::LitStr(lit) => lit,
SpecializeArg::Conditional(conditional) => &conditional.lit,
}
}
fn condition(&self) -> Option<&TokenStream2> {
match self {
SpecializeArg::LitStr(..) => None,
SpecializeArg::Conditional(conditional) => Some(&conditional.attr),
}
}
}
struct Conditional {
lit: syn::LitStr,
attr: TokenStream2,
}
impl syn::parse::Parse for Conditional {
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
let lit = input.parse()?;
input.parse::<syn::Token![,]>()?;
let attr = input.parse()?;
Ok(Conditional { lit, attr })
}
}
impl syn::parse::Parse for SpecializeArg {
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
let lookahead = input.lookahead1();
if lookahead.peek(kw::conditional) {
input.parse::<kw::conditional>()?;
let content;
syn::parenthesized!(content in input);
let conditional = content.parse()?;
Ok(SpecializeArg::Conditional(conditional))
} else {
Ok(SpecializeArg::LitStr(input.parse()?))
}
}
}
struct SpecializeArgs(syn::punctuated::Punctuated<SpecializeArg, syn::Token![,]>);
impl syn::parse::Parse for SpecializeArgs {
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
Ok(Self(syn::punctuated::Punctuated::parse_terminated(input)?))
}
}
#[proc_macro_attribute]
pub fn unsafe_target_feature(attributes: TokenStream, input: TokenStream) -> TokenStream {
let attributes = syn::parse_macro_input!(attributes as syn::LitStr);
let item = syn::parse_macro_input!(input as syn::Item);
process_item(&attributes, item, true)
}
#[proc_macro_attribute]
pub fn unsafe_target_feature_specialize(
attributes: TokenStream,
input: TokenStream,
) -> TokenStream {
let attributes = syn::parse_macro_input!(attributes as SpecializeArgs);
let item_mod = syn::parse_macro_input!(input as syn::ItemMod);
let mut out = Vec::new();
for attributes in attributes.0 {
let features: Vec<_> = attributes
.lit()
.value()
.split(',')
.map(|feature| feature.replace(' ', ""))
.collect();
let name = format!("{}_{}", item_mod.ident, features.join("_"));
let ident = syn::Ident::new(&name, item_mod.ident.span());
let mut attrs = item_mod.attrs.clone();
if let Some(condition) = attributes.condition() {
attrs.push(syn::Attribute {
pound_token: Default::default(),
style: syn::AttrStyle::Outer,
bracket_token: Default::default(),
meta: syn::Meta::List(syn::MetaList {
path: syn::Ident::new("cfg", attributes.lit().span()).into(),
delimiter: syn::MacroDelimiter::Paren(Default::default()),
tokens: condition.clone(),
}),
});
}
let item_mod = process_mod(
attributes.lit(),
syn::ItemMod {
attrs,
ident,
..item_mod.clone()
},
Some(features),
);
out.push(item_mod);
}
quote::quote! {
#(#out)*
}
.into()
}
fn process_item(attributes: &syn::LitStr, item: syn::Item, strict: bool) -> TokenStream {
match item {
syn::Item::Fn(function) => process_function(attributes, function, None),
syn::Item::Impl(item_impl) => process_impl(attributes, item_impl),
syn::Item::Mod(item_mod) => process_mod(attributes, item_mod, None).into(),
item => {
if strict {
unsupported!(item)
} else {
quote::quote! { #item }.into()
}
}
}
}
fn process_mod(
attributes: &syn::LitStr,
mut item_mod: syn::ItemMod,
spec_features: Option<Vec<String>>,
) -> TokenStream2 {
if let Some((_, ref mut content)) = item_mod.content {
'next_item: for item in content {
if let Some(ref spec_features) = spec_features {
match item {
syn::Item::Const(syn::ItemConst { ref mut attrs, .. })
| syn::Item::Enum(syn::ItemEnum { ref mut attrs, .. })
| syn::Item::ExternCrate(syn::ItemExternCrate { ref mut attrs, .. })
| syn::Item::Fn(syn::ItemFn { ref mut attrs, .. })
| syn::Item::ForeignMod(syn::ItemForeignMod { ref mut attrs, .. })
| syn::Item::Impl(syn::ItemImpl { ref mut attrs, .. })
| syn::Item::Macro(syn::ItemMacro { ref mut attrs, .. })
| syn::Item::Mod(syn::ItemMod { ref mut attrs, .. })
| syn::Item::Static(syn::ItemStatic { ref mut attrs, .. })
| syn::Item::Struct(syn::ItemStruct { ref mut attrs, .. })
| syn::Item::Trait(syn::ItemTrait { ref mut attrs, .. })
| syn::Item::TraitAlias(syn::ItemTraitAlias { ref mut attrs, .. })
| syn::Item::Type(syn::ItemType { ref mut attrs, .. })
| syn::Item::Union(syn::ItemUnion { ref mut attrs, .. })
| syn::Item::Use(syn::ItemUse { ref mut attrs, .. }) => {
let mut index = 0;
while index < attrs.len() {
let attr = &attrs[index];
if matches!(attr.style, syn::AttrStyle::Outer) {
match attr.meta {
syn::Meta::List(ref list)
if is_path_eq(&list.path, "for_target_feature") =>
{
let feature: syn::LitStr = match list.parse_args() {
Ok(feature) => feature,
Err(error) => {
return error.into_compile_error();
}
};
let feature = feature.value();
if !spec_features
.iter()
.any(|enabled_feature| feature == *enabled_feature)
{
*item = syn::Item::Verbatim(Default::default());
continue 'next_item;
}
attrs.remove(index);
continue;
}
_ => {}
}
}
index += 1;
continue;
}
}
_ => {
unsupported!(item_mod);
}
}
}
*item = syn::Item::Verbatim(
process_item(
attributes,
std::mem::replace(item, syn::Item::Verbatim(Default::default())),
false,
)
.into(),
);
}
}
quote::quote! {
#item_mod
}
}
fn process_impl(attributes: &syn::LitStr, mut item_impl: syn::ItemImpl) -> TokenStream {
unsupported_if_some!(item_impl.defaultness);
unsupported_if_some!(item_impl.unsafety);
let mut items = Vec::new();
for item in item_impl.items.drain(..) {
match item {
syn::ImplItem::Fn(function) => {
unsupported_if_some!(function.defaultness);
let function = syn::ItemFn {
attrs: function.attrs,
vis: function.vis,
sig: function.sig,
block: Box::new(function.block),
};
let output_item = process_function(
attributes,
function,
Some((item_impl.generics.clone(), item_impl.self_ty.clone())),
);
items.push(syn::ImplItem::Verbatim(output_item.into()));
}
item => items.push(item),
}
}
item_impl.items = items;
quote::quote! {
#item_impl
}
.into()
}
fn is_path_eq(path: &syn::Path, ident: &str) -> bool {
let segments: Vec<_> = ident.split("::").collect();
path.segments.len() == segments.len()
&& path
.segments
.iter()
.zip(segments.iter())
.all(|(segment, expected)| segment.ident == expected && segment.arguments.is_none())
}
fn process_function(
attributes: &syn::LitStr,
function: syn::ItemFn,
outer: Option<(syn::Generics, Box<syn::Type>)>,
) -> TokenStream {
if function.sig.unsafety.is_some() {
return quote::quote! {
#[target_feature(enable = #attributes)]
#function
}
.into();
}
unsupported_if_some!(function.sig.constness);
unsupported_if_some!(function.sig.asyncness);
unsupported_if_some!(function.sig.abi);
unsupported_if_some!(function.sig.variadic);
let function_visibility = function.vis;
let function_name = function.sig.ident;
let function_return = function.sig.output;
let function_inner_name =
syn::Ident::new(&format!("_impl_{}", function_name), function_name.span());
let function_args = function.sig.inputs;
let function_body = function.block;
let mut function_call_args = Vec::new();
let mut function_args_outer = Vec::new();
let mut function_args_inner = Vec::new();
for (index, arg) in function_args.iter().enumerate() {
match arg {
syn::FnArg::Receiver(receiver) => {
unsupported_if_some!(receiver.attrs.first());
unsupported_if_some!(receiver.colon_token);
if outer.is_none() {
return syn::Error::new(receiver.span(), "unsupported by #[unsafe_target_feature(...)]; put the attribute on the outer `impl`").into_compile_error().into();
}
function_args_inner.push(syn::FnArg::Receiver(receiver.clone()));
function_args_outer.push(syn::FnArg::Receiver(receiver.clone()));
function_call_args.push(syn::Ident::new("self", receiver.self_token.span()));
}
syn::FnArg::Typed(ty) => {
unsupported_if_some!(ty.attrs.first());
match &*ty.pat {
syn::Pat::Ident(pat_ident) => {
unsupported_if_some!(pat_ident.attrs.first());
function_args_inner.push(arg.clone());
function_args_outer.push(syn::FnArg::Typed(syn::PatType {
attrs: Vec::new(),
pat: Box::new(syn::Pat::Ident(syn::PatIdent {
attrs: Vec::new(),
by_ref: None,
mutability: None,
ident: pat_ident.ident.clone(),
subpat: None,
})),
colon_token: ty.colon_token,
ty: ty.ty.clone(),
}));
function_call_args.push(pat_ident.ident.clone());
}
syn::Pat::Wild(pat_wild) => {
unsupported_if_some!(pat_wild.attrs.first());
let ident = syn::Ident::new(
&format!("__arg_{}__", index),
pat_wild.underscore_token.span(),
);
function_args_inner.push(arg.clone());
function_args_outer.push(syn::FnArg::Typed(syn::PatType {
attrs: Vec::new(),
pat: Box::new(syn::Pat::Ident(syn::PatIdent {
attrs: Vec::new(),
by_ref: None,
mutability: None,
ident: ident.clone(),
subpat: None,
})),
colon_token: ty.colon_token,
ty: ty.ty.clone(),
}));
function_call_args.push(ident);
}
_ => unsupported!(arg),
}
}
}
}
let mut maybe_inline = quote::quote! {};
let mut maybe_outer_attributes = Vec::new();
let mut maybe_cfg = quote::quote! {};
for attribute in function.attrs {
match &attribute.meta {
syn::Meta::Path(path) if is_path_eq(path, "inline") => {
maybe_inline = quote::quote! { #[inline] };
}
syn::Meta::Path(path) if is_path_eq(path, "test") => {
maybe_outer_attributes.push(attribute);
maybe_cfg = quote::quote! { #[cfg(target_feature = #attributes)] };
}
syn::Meta::List(syn::MetaList { path, tokens, .. })
if is_path_eq(path, "inline") && tokens.to_string() == "always" =>
{
maybe_inline = quote::quote! { #[inline] };
}
syn::Meta::NameValue(syn::MetaNameValue { path, .. }) if is_path_eq(path, "doc") => {
maybe_outer_attributes.push(attribute);
}
syn::Meta::List(syn::MetaList { path, .. })
if is_path_eq(path, "cfg")
|| is_path_eq(path, "allow")
|| is_path_eq(path, "deny") =>
{
maybe_outer_attributes.push(attribute);
}
syn::Meta::Path(path) if is_path_eq(path, "rustfmt::skip") => {
maybe_outer_attributes.push(attribute);
}
_ => unsupported!(attribute),
}
}
let (fn_impl_generics, fn_ty_generics, fn_where_clause) =
function.sig.generics.split_for_impl();
let fn_call_generics = fn_ty_generics.as_turbofish();
if let Some((generics, self_ty)) = outer {
let (outer_impl_generics, outer_ty_generics, outer_where_clause) =
generics.split_for_impl();
let trait_ident =
syn::Ident::new(&format!("__Impl_{}__", function_name), function_name.span());
let item_trait = quote::quote! {
#[allow(non_camel_case_types)]
trait #trait_ident #outer_impl_generics #outer_where_clause {
unsafe fn #function_inner_name #fn_impl_generics (#(#function_args_outer),*) #function_return #fn_where_clause;
}
};
let item_trait_impl = quote::quote! {
impl #outer_impl_generics #trait_ident #outer_ty_generics for #self_ty #outer_where_clause {
#[target_feature(enable = #attributes)]
#maybe_inline
unsafe fn #function_inner_name #fn_impl_generics (#(#function_args_inner),*) #function_return #fn_where_clause #function_body
}
};
quote::quote! {
#[inline(always)]
#(#maybe_outer_attributes)*
#function_visibility fn #function_name #fn_impl_generics (#(#function_args_outer),*) #function_return #fn_where_clause {
#item_trait
#item_trait_impl
unsafe {
<Self as #trait_ident #outer_ty_generics> ::#function_inner_name #fn_call_generics (#(#function_call_args),*)
}
}
}.into()
} else {
quote::quote! {
#[inline(always)]
#maybe_cfg
#(#maybe_outer_attributes)*
#function_visibility fn #function_name #fn_impl_generics (#(#function_args_outer),*) #function_return #fn_where_clause {
#[target_feature(enable = #attributes)]
#maybe_inline
unsafe fn #function_inner_name #fn_impl_generics (#(#function_args_inner),*) #function_return #fn_where_clause #function_body
unsafe {
#function_inner_name #fn_call_generics (#(#function_call_args),*)
}
}
}.into()
}
}