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::{
13    quote,
14    ToTokens, //
15};
16use syn::{
17    parse_quote,
18    Error,
19    ImplItem,
20    Item,
21    ItemImpl,
22    ItemTrait,
23    Result,
24    TraitItem, //
25};
26
27fn handle_trait(mut item: ItemTrait) -> Result<ItemTrait> {
28    let mut gen_items = Vec::new();
29
30    gen_items.push(parse_quote! {
31         /// A marker to prevent implementors from forgetting to use [`#[vtable]`](vtable)
32         /// attribute when implementing this trait.
33         const USE_VTABLE_ATTR: ();
34    });
35
36    // Add `type OwnerModule: ModuleMetadata` as a required associated type if
37    // the trait does not already define it.
38    if !item
39        .items
40        .iter()
41        .any(|i| matches!(i, TraitItem::Type(t) if t.ident == "OwnerModule"))
42    {
43        gen_items.push(parse_quote! {
44            /// The module implementing this vtable trait.
45            ///
46            /// Automatically set to `crate::LocalModule` by the `#[vtable]`
47            /// impl macro.
48            type OwnerModule: ::kernel::ModuleMetadata;
49        });
50    }
51
52    for item in &mut item.items {
53        if let TraitItem::Fn(fn_item) = item {
54            let name = &fn_item.sig.ident;
55            let gen_const_name = Ident::new(
56                &format!("HAS_{}", name.to_string().to_uppercase()),
57                name.span(),
58            );
59
60            if fn_item
61                .attrs
62                .extract_if(.., |attr| attr.path().is_ident("optional"))
63                .count()
64                != 0
65            {
66                if let Some(default) = &fn_item.default {
67                    Err(Error::new_spanned(
68                        default,
69                        "`#[optional]` methods must not have a default implementation",
70                    ))?;
71                }
72
73                // Optional methods in Rust need a default implementation. Inject one that fails
74                // the build at compile time if not overridden.
75                fn_item.default = Some(parse_quote!({
76                    ::kernel::build_assert::build_error!(
77                        "This function must not be called, see the `#[vtable]` documentation."
78                    );
79                }));
80                // Ensure that the function is never code generated unless used.
81                fn_item.attrs.push(parse_quote!(#[inline]));
82            }
83
84            // We don't know on the implementation-site whether a method is required or provided
85            // so we have to generate a const for all methods.
86            // However, hide it for required methods as it will always be true.
87            let doc = if fn_item.default.is_some() {
88                let comment =
89                    format!("Indicates if the `{name}` method is overridden by the implementor.");
90                quote!(#[doc = #comment])
91            } else {
92                quote!(#[doc(hidden)])
93            };
94            let cfg_attrs = crate::helpers::gather_cfg_attrs(&fn_item.attrs);
95            gen_items.push(parse_quote! {
96                #(#cfg_attrs)*
97                #doc
98                const #gen_const_name: bool = false;
99            });
100        }
101    }
102
103    item.items.extend(gen_items);
104    Ok(item)
105}
106
107fn handle_impl(mut item: ItemImpl) -> Result<ItemImpl> {
108    let mut gen_items = Vec::new();
109    let mut defined_items = HashSet::new();
110
111    // Iterate over all user-defined items to gather any possible explicit overrides.
112    for item in &item.items {
113        match item {
114            ImplItem::Const(const_item) => {
115                defined_items.insert(const_item.ident.clone());
116            }
117            ImplItem::Type(type_item) => {
118                defined_items.insert(type_item.ident.clone());
119            }
120            _ => {}
121        }
122    }
123
124    gen_items.push(parse_quote! {
125        const USE_VTABLE_ATTR: () = ();
126    });
127
128    // Auto-insert `type OwnerModule = crate::LocalModule` if not explicitly defined.
129    // `crate::LocalModule` resolves to the real module type (via `module!`) or a
130    // dummy fallback in non-module contexts (e.g., doctests).
131    if !defined_items.contains(&parse_quote!(OwnerModule)) {
132        gen_items.push(parse_quote! {
133            type OwnerModule = crate::LocalModule;
134        });
135    }
136
137    for item in &item.items {
138        if let ImplItem::Fn(fn_item) = item {
139            let name = &fn_item.sig.ident;
140            let gen_const_name = Ident::new(
141                &format!("HAS_{}", name.to_string().to_uppercase()),
142                name.span(),
143            );
144            // Skip if it's declared already -- this allows user override.
145            if defined_items.contains(&gen_const_name) {
146                continue;
147            }
148            let cfg_attrs = crate::helpers::gather_cfg_attrs(&fn_item.attrs);
149            gen_items.push(parse_quote! {
150                #(#cfg_attrs)*
151                const #gen_const_name: bool = true;
152            });
153        }
154    }
155
156    item.items.extend(gen_items);
157    Ok(item)
158}
159
160pub(crate) fn vtable(input: Item) -> Result<TokenStream> {
161    match input {
162        Item::Trait(item) => Ok(handle_trait(item)?.into_token_stream()),
163        Item::Impl(item) => Ok(handle_impl(item)?.into_token_stream()),
164        _ => Err(Error::new_spanned(
165            input,
166            "`#[vtable]` attribute should only be applied to trait or impl block",
167        ))?,
168    }
169}