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    // Add `type OwnerModule: ModuleMetadata` as a required associated type if
34    // the trait does not already define it.
35    if !item
36        .items
37        .iter()
38        .any(|i| matches!(i, TraitItem::Type(t) if t.ident == "OwnerModule"))
39    {
40        gen_items.push(parse_quote! {
41            /// The module implementing this vtable trait.
42            ///
43            /// Automatically set to `crate::LocalModule` by the `#[vtable]`
44            /// impl macro.
45            type OwnerModule: ::kernel::ModuleMetadata;
46        });
47    }
48
49    for item in &item.items {
50        if let TraitItem::Fn(fn_item) = item {
51            let name = &fn_item.sig.ident;
52            let gen_const_name = Ident::new(
53                &format!("HAS_{}", name.to_string().to_uppercase()),
54                name.span(),
55            );
56
57            // We don't know on the implementation-site whether a method is required or provided
58            // so we have to generate a const for all methods.
59            let cfg_attrs = crate::helpers::gather_cfg_attrs(&fn_item.attrs);
60            let comment =
61                format!("Indicates if the `{name}` method is overridden by the implementor.");
62            gen_items.push(parse_quote! {
63                #(#cfg_attrs)*
64                #[doc = #comment]
65                const #gen_const_name: bool = false;
66            });
67        }
68    }
69
70    item.items.extend(gen_items);
71    Ok(item)
72}
73
74fn handle_impl(mut item: ItemImpl) -> Result<ItemImpl> {
75    let mut gen_items = Vec::new();
76    let mut defined_items = HashSet::new();
77
78    // Iterate over all user-defined items to gather any possible explicit overrides.
79    for item in &item.items {
80        match item {
81            ImplItem::Const(const_item) => {
82                defined_items.insert(const_item.ident.clone());
83            }
84            ImplItem::Type(type_item) => {
85                defined_items.insert(type_item.ident.clone());
86            }
87            _ => {}
88        }
89    }
90
91    gen_items.push(parse_quote! {
92        const USE_VTABLE_ATTR: () = ();
93    });
94
95    // Auto-insert `type OwnerModule = crate::LocalModule` if not explicitly defined.
96    // `crate::LocalModule` resolves to the real module type (via `module!`) or a
97    // dummy fallback in non-module contexts (e.g., doctests).
98    if !defined_items.contains(&parse_quote!(OwnerModule)) {
99        gen_items.push(parse_quote! {
100            type OwnerModule = crate::LocalModule;
101        });
102    }
103
104    for item in &item.items {
105        if let ImplItem::Fn(fn_item) = item {
106            let name = &fn_item.sig.ident;
107            let gen_const_name = Ident::new(
108                &format!("HAS_{}", name.to_string().to_uppercase()),
109                name.span(),
110            );
111            // Skip if it's declared already -- this allows user override.
112            if defined_items.contains(&gen_const_name) {
113                continue;
114            }
115            let cfg_attrs = crate::helpers::gather_cfg_attrs(&fn_item.attrs);
116            gen_items.push(parse_quote! {
117                #(#cfg_attrs)*
118                const #gen_const_name: bool = true;
119            });
120        }
121    }
122
123    item.items.extend(gen_items);
124    Ok(item)
125}
126
127pub(crate) fn vtable(input: Item) -> Result<TokenStream> {
128    match input {
129        Item::Trait(item) => Ok(handle_trait(item)?.into_token_stream()),
130        Item::Impl(item) => Ok(handle_impl(item)?.into_token_stream()),
131        _ => Err(Error::new_spanned(
132            input,
133            "`#[vtable]` attribute should only be applied to trait or impl block",
134        ))?,
135    }
136}