xref: /linux/rust/macros/vtable.rs (revision 23b0f90ba871f096474e1c27c3d14f455189d2d9)
1 // SPDX-License-Identifier: GPL-2.0
2 
3 use std::{
4     collections::HashSet,
5     iter::Extend, //
6 };
7 
8 use proc_macro2::{
9     Ident,
10     TokenStream, //
11 };
12 use quote::ToTokens;
13 use syn::{
14     parse_quote,
15     Error,
16     ImplItem,
17     Item,
18     ItemImpl,
19     ItemTrait,
20     Result,
21     TraitItem, //
22 };
23 
24 fn handle_trait(mut item: ItemTrait) -> Result<ItemTrait> {
25     let mut gen_items = Vec::new();
26 
27     gen_items.push(parse_quote! {
28          /// A marker to prevent implementors from forgetting to use [`#[vtable]`](vtable)
29          /// attribute when implementing this trait.
30          const USE_VTABLE_ATTR: ();
31     });
32 
33     for item in &item.items {
34         if let TraitItem::Fn(fn_item) = item {
35             let name = &fn_item.sig.ident;
36             let gen_const_name = Ident::new(
37                 &format!("HAS_{}", name.to_string().to_uppercase()),
38                 name.span(),
39             );
40 
41             // We don't know on the implementation-site whether a method is required or provided
42             // so we have to generate a const for all methods.
43             let cfg_attrs = crate::helpers::gather_cfg_attrs(&fn_item.attrs);
44             let comment =
45                 format!("Indicates if the `{name}` method is overridden by the implementor.");
46             gen_items.push(parse_quote! {
47                 #(#cfg_attrs)*
48                 #[doc = #comment]
49                 const #gen_const_name: bool = false;
50             });
51         }
52     }
53 
54     item.items.extend(gen_items);
55     Ok(item)
56 }
57 
58 fn handle_impl(mut item: ItemImpl) -> Result<ItemImpl> {
59     let mut gen_items = Vec::new();
60     let mut defined_consts = HashSet::new();
61 
62     // Iterate over all user-defined constants to gather any possible explicit overrides.
63     for item in &item.items {
64         if let ImplItem::Const(const_item) = item {
65             defined_consts.insert(const_item.ident.clone());
66         }
67     }
68 
69     gen_items.push(parse_quote! {
70         const USE_VTABLE_ATTR: () = ();
71     });
72 
73     for item in &item.items {
74         if let ImplItem::Fn(fn_item) = item {
75             let name = &fn_item.sig.ident;
76             let gen_const_name = Ident::new(
77                 &format!("HAS_{}", name.to_string().to_uppercase()),
78                 name.span(),
79             );
80             // Skip if it's declared already -- this allows user override.
81             if defined_consts.contains(&gen_const_name) {
82                 continue;
83             }
84             let cfg_attrs = crate::helpers::gather_cfg_attrs(&fn_item.attrs);
85             gen_items.push(parse_quote! {
86                 #(#cfg_attrs)*
87                 const #gen_const_name: bool = true;
88             });
89         }
90     }
91 
92     item.items.extend(gen_items);
93     Ok(item)
94 }
95 
96 pub(crate) fn vtable(input: Item) -> Result<TokenStream> {
97     match input {
98         Item::Trait(item) => Ok(handle_trait(item)?.into_token_stream()),
99         Item::Impl(item) => Ok(handle_impl(item)?.into_token_stream()),
100         _ => Err(Error::new_spanned(
101             input,
102             "`#[vtable]` attribute should only be applied to trait or impl block",
103         ))?,
104     }
105 }
106