Emit smaller dispatcher code

Author: TethysSvensson

src/dispatcher.rs

@@ -11,7 +11,7 @@ pub(crate) struct Dispatcher {
 
 impl ToTokens for Dispatcher {
     fn to_tokens(&self, tokens: &mut TokenStream) {
-        let function_arguments = &self.signature.inputs.iter();
+        let function_arguments = self.signature.inputs.iter();
         let argument_names = &self
             .signature
             .inputs
@@ -48,50 +48,41 @@ impl ToTokens for Dispatcher {
             let return_if_detected = self.functions.iter().map(|(target, function)| {
                 let target_arch = target.target_arch();
                 let features_detected = target.features_detected();
-                let function_arguments = function_arguments.clone();
-                let argument_names = argument_names.clone();
                 quote! {
                     #target_arch
                     {
-                        // wrap the function in an unsafe function with the same signature
-                        unsafe fn unsafe_wrapper(#(#function_arguments),*) #returns {
-                            #function(#(#argument_names),*)
-                        }
                         if #features_detected {
-                            return unsafe_wrapper
+                            return #function
                         }
                     }
                 }
             });
             let default = &self.default;
-            let function_arguments = function_arguments.clone();
-            let argument_names = argument_names.clone();
             quote! {
                 #(#return_if_detected)*
                 #[cfg(not(any(#(target_arch = #defaulted_arches),*)))]
                 {
-                    unsafe fn unsafe_wrapper(#(#function_arguments),*) #returns {
-                        #default(#(#argument_names),*)
-                    }
-                    return unsafe_wrapper
+                    return #default
                 }
             }
         };
         tokens.extend(quote! {
-            use std::sync::atomic::{AtomicUsize, Ordering};
+            use std::sync::atomic::{AtomicPtr, Ordering};
             type __fn_ty = unsafe fn (#(#argument_ty),*) #returns;
             #[cold]
             fn __get_fn() -> __fn_ty {
                 #feature_detection
             }
-            static __DISPATCHED_FN: AtomicUsize = AtomicUsize::new(0usize);
-            let mut __current_ptr = __DISPATCHED_FN.load(Ordering::SeqCst);
-            if __current_ptr == 0 {
-                __current_ptr = unsafe { std::mem::transmute(__get_fn()) };
-                __DISPATCHED_FN.store(__current_ptr, Ordering::SeqCst);
+            #[cold]
+            unsafe fn __resolver_fn (#(#function_arguments),*) #returns {
+                let __current_fn = __get_fn();
+                __DISPATCHED_FN.store(__current_fn as *mut (), Ordering::SeqCst);
+                __current_fn(#(#argument_names),*)
             }
+            static __DISPATCHED_FN: AtomicPtr<()> = AtomicPtr::new(__resolver_fn as *mut ());
+            let __current_ptr = __DISPATCHED_FN.load(Ordering::SeqCst);
             unsafe {
-                let __current_fn = std::mem::transmute::<usize, __fn_ty>(__current_ptr);
+                let __current_fn = std::mem::transmute::<*mut (), __fn_ty>(__current_ptr);
                 __current_fn(#(#argument_names),*)
             }
         });

src/target_clones.rs

@@ -6,7 +6,7 @@ use proc_macro2::{Span, TokenStream};
 use quote::{quote, ToTokens};
 use syn::parse::{Parse, ParseStream, Result};
 use syn::punctuated::Punctuated;
-use syn::{token, Attribute, Block, FnArg, Ident, ItemFn, LitStr, Signature, Token, Visibility};
+use syn::{token, Attribute, Block, Ident, ItemFn, LitStr, Signature, Token, Visibility};
 
 pub(crate) struct Config {
     targets: Vec<Target>,
@@ -39,36 +39,12 @@ impl ToTokens for FunctionClone<'_> {
             .as_ref()
             .map_or(TokenStream::new(), |x| x.target_features());
         let signature = &self.signature;
-        let body = &self.body;
-        tokens.extend(if signature.unsafety.is_some() || self.target.is_none() {
-            quote! { #target_arch #target_features #signature #body }
-        } else {
-            let mut unsafe_signature = signature.clone();
-            unsafe_signature.unsafety = Some(token::Unsafe {
-                span: Span::call_site(),
-            });
-            unsafe_signature.ident = Ident::new("__unsafe_fn", Span::call_site());
-            let argument_names = &signature
-                .inputs
-                .iter()
-                .map(|x| {
-                    if let FnArg::Typed(p) = x {
-                        p.pat.as_ref()
-                    } else {
-                        unimplemented!("member fn not supported")
-                    }
-                })
-                .collect::<Vec<_>>();
-            quote! {
-                #target_arch
-                #signature {
-                    #target_features #unsafe_signature #body
-                    unsafe {
-                        __unsafe_fn(#(#argument_names),*)
-                    }
-                }
-            }
+        let mut unsafe_signature = signature.clone();
+        unsafe_signature.unsafety = Some(token::Unsafe {
+            span: Span::call_site(),
         });
+        let body = &self.body;
+        tokens.extend(quote! { #target_arch #target_features #unsafe_signature #body });
     }
 }