libhotpatch_macros/
lib.rs

1use proc_macro2::{Span, TokenStream};
2use quote::{ToTokens, quote};
3use syn::{
4    Abi, FnArg, ImplItemFn, LitByteStr, LitStr, Pat, PatWild, Token, parse_macro_input,
5    parse_quote, token::Extern,
6};
7
8use crate::hotpatch_fn::HotpatchFn;
9
10mod hotpatch_fn;
11
12#[proc_macro_attribute]
13pub fn hotpatch(
14    _: proc_macro::TokenStream,
15    input: proc_macro::TokenStream,
16) -> proc_macro::TokenStream {
17    let HotpatchFn { inner, outer } = parse_macro_input!(input as HotpatchFn);
18
19    let ImplItemFn {
20        mut attrs,
21        vis,
22        defaultness,
23        mut sig,
24        ..
25    } = outer;
26
27    attrs.push(parse_quote!(#[allow(improper_ctypes_definitions)]));
28
29    sig.abi.get_or_insert_with(|| Abi {
30        extern_token: Extern(Span::call_site()),
31        name: Some(LitStr::new("C-unwind", Span::call_site())),
32    });
33
34    let sig_str = quote!(sig).to_string();
35    let sig_lit = LitByteStr::new(sig_str.as_bytes(), Span::call_site());
36
37    let outer_fn = &sig.ident;
38    let inner_fn = &inner.sig.ident;
39
40    let args = inner.sig.inputs.iter().map(|input| match input {
41        FnArg::Receiver(_) => parse_quote!(self),
42        FnArg::Typed(typed) => fn_input_pat_to_ts(&typed.pat),
43    });
44
45    let wild = inner.sig.inputs.iter().map(|_| PatWild {
46        attrs: vec![],
47        underscore_token: Token![_](Span::call_site()),
48    });
49
50    quote! {
51        #(#attrs)*
52        #vis #defaultness #sig {
53            #inner
54            fn __libhotpatch_type_of() -> (u128, &'static str) {
55                let name = ::std::any::type_name_of_val(&#outer_fn);
56                let mut hasher = libhotpatch::Xxh3::new();
57                ::std::hash::Hash::hash(#sig_lit, &mut hasher);
58                ::std::hash::Hash::hash(name.as_bytes(), &mut hasher);
59                (hasher.digest128(), name)
60            }
61            #[libhotpatch::distributed_slice(libhotpatch::HOTPATCH_FN)]
62            #[linkme(crate = libhotpatch::linkme)]
63            static __LIBHOTPATCH_HOTPATCH_FN: (
64                ::std::sync::atomic::AtomicPtr<()>,
65                libhotpatch::LibraryHandle,
66                fn() -> (u128, &'static str),
67            ) = (
68                ::std::sync::atomic::AtomicPtr::new(#inner_fn as *mut ()),
69                libhotpatch::LibraryHandle::null(),
70                __libhotpatch_type_of,
71            );
72            libhotpatch::Watcher::get().map(libhotpatch::Watcher::poll);
73            let __libhotpatch_library_handle = __LIBHOTPATCH_HOTPATCH_FN.1.clone();
74            unsafe {
75                ::std::mem::transmute::<_, fn(#(#wild,)*) -> _>(
76                    __LIBHOTPATCH_HOTPATCH_FN.0.load(::std::sync::atomic::Ordering::Relaxed))(#(#args,)*
77                )
78            }
79        }
80    }
81    .into()
82}
83
84fn fn_input_pat_to_ts(pat: &Pat) -> TokenStream {
85    match pat {
86        Pat::Ident(pat_ident) => pat_ident.ident.clone().to_token_stream(),
87        Pat::Paren(pat_paren) => fn_input_pat_to_ts(&pat_paren.pat),
88        Pat::Reference(pat_ref) => fn_input_pat_to_ts(&pat_ref.pat),
89        Pat::Tuple(pat_tuple) => {
90            let elems = &pat_tuple.elems;
91            parse_quote!((#elems))
92        }
93        Pat::Struct(pat_struct) => {
94            let path = &pat_struct.path;
95            let members = pat_struct.fields.iter().map(|field_pat| &field_pat.member);
96
97            parse_quote!(#path { #(#members,)* })
98        }
99        Pat::TupleStruct(pat_tstruct) => {
100            let path = &pat_tstruct.path;
101            let elems = pat_tstruct.elems.iter().map(fn_input_pat_to_ts);
102
103            parse_quote!(#path(#(#elems,)*))
104        }
105        _ => panic!("unsupported type pattern in function input position"),
106    }
107}