Skip to main content

macros/
vtable.rs

1// SPDX-License-Identifier: GPL-2.0
2
3use std::{
4    collections::HashSet,
5    iter::Extend, //
6};
7
8use proc_macro2::{
9    Ident,
10    TokenStream, //
11};
12use quote::ToTokens;
13use syn::{
14    parse_quote,
15    Error,
16    ImplItem,
17    Item,
18    ItemImpl,
19    ItemTrait,
20    Result,
21    TraitItem, //
22};
23
24fn 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
58fn 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
96pub(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}