1// SPDX-License-Identifier: GPL-2.0
 2
 3use proc_macro::{Delimiter, Group, TokenStream, TokenTree};
 4use std::collections::HashSet;
 5use std::fmt::Write;
 6
 7pub(crate) fn vtable(_attr: TokenStream, ts: TokenStream) -> TokenStream {
 8    let mut tokens: Vec<_> = ts.into_iter().collect();
 9
10    // Scan for the `trait` or `impl` keyword.
11    let is_trait = tokens
12        .iter()
13        .find_map(|token| match token {
14            TokenTree::Ident(ident) => match ident.to_string().as_str() {
15                "trait" => Some(true),
16                "impl" => Some(false),
17                _ => None,
18            },
19            _ => None,
20        })
21        .expect("#[vtable] attribute should only be applied to trait or impl block");
22
23    // Retrieve the main body. The main body should be the last token tree.
24    let body = match tokens.pop() {
25        Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Brace => group,
26        _ => panic!("cannot locate main body of trait or impl block"),
27    };
28
29    let mut body_it = body.stream().into_iter();
30    let mut functions = Vec::new();
31    let mut consts = HashSet::new();
32    while let Some(token) = body_it.next() {
33        match token {
34            TokenTree::Ident(ident) if ident.to_string() == "fn" => {
35                let fn_name = match body_it.next() {
36                    Some(TokenTree::Ident(ident)) => ident.to_string(),
37                    // Possibly we've encountered a fn pointer type instead.
38                    _ => continue,
39                };
40                functions.push(fn_name);
41            }
42            TokenTree::Ident(ident) if ident.to_string() == "const" => {
43                let const_name = match body_it.next() {
44                    Some(TokenTree::Ident(ident)) => ident.to_string(),
45                    // Possibly we've encountered an inline const block instead.
46                    _ => continue,
47                };
48                consts.insert(const_name);
49            }
50            _ => (),
51        }
52    }
53
54    let mut const_items;
55    if is_trait {
56        const_items = "
57                /// A marker to prevent implementors from forgetting to use [`#[vtable]`](vtable)
58                /// attribute when implementing this trait.
59                const USE_VTABLE_ATTR: ();
60        "
61        .to_owned();
62
63        for f in functions {
64            let gen_const_name = format!("HAS_{}", f.to_uppercase());
65            // Skip if it's declared already -- this allows user override.
66            if consts.contains(&gen_const_name) {
67                continue;
68            }
69            // We don't know on the implementation-site whether a method is required or provided
70            // so we have to generate a const for all methods.
71            write!(
72                const_items,
73                "/// Indicates if the `{f}` method is overridden by the implementor.
74                const {gen_const_name}: bool = false;",
75            )
76            .unwrap();
77            consts.insert(gen_const_name);
78        }
79    } else {
80        const_items = "const USE_VTABLE_ATTR: () = ();".to_owned();
81
82        for f in functions {
83            let gen_const_name = format!("HAS_{}", f.to_uppercase());
84            if consts.contains(&gen_const_name) {
85                continue;
86            }
87            write!(const_items, "const {gen_const_name}: bool = true;").unwrap();
88        }
89    }
90
91    let new_body = vec![const_items.parse().unwrap(), body.stream()]
92        .into_iter()
93        .collect();
94    tokens.push(TokenTree::Group(Group::new(Delimiter::Brace, new_body)));
95    tokens.into_iter().collect()
96}