1 // SPDX-License-Identifier: GPL-2.0
2 
3 use proc_macro::{Delimiter, Group, TokenStream, TokenTree};
4 use std::collections::HashSet;
5 use std::fmt::Write;
6 
vtable(_attr: TokenStream, ts: TokenStream) -> TokenStream7 pub(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         }
78     } else {
79         const_items = "const USE_VTABLE_ATTR: () = ();".to_owned();
80 
81         for f in functions {
82             let gen_const_name = format!("HAS_{}", f.to_uppercase());
83             if consts.contains(&gen_const_name) {
84                 continue;
85             }
86             write!(const_items, "const {gen_const_name}: bool = true;").unwrap();
87         }
88     }
89 
90     let new_body = vec![const_items.parse().unwrap(), body.stream()]
91         .into_iter()
92         .collect();
93     tokens.push(TokenTree::Group(Group::new(Delimiter::Brace, new_body)));
94     tokens.into_iter().collect()
95 }
96