97 lines
3.5 KiB
Rust
97 lines
3.5 KiB
Rust
// SPDX-License-Identifier: GPL-2.0
|
|
|
|
use proc_macro::{Delimiter, Group, TokenStream, TokenTree};
|
|
use std::collections::HashSet;
|
|
use std::fmt::Write;
|
|
|
|
pub(crate) fn vtable(_attr: TokenStream, ts: TokenStream) -> TokenStream {
|
|
let mut tokens: Vec<_> = ts.into_iter().collect();
|
|
|
|
// Scan for the `trait` or `impl` keyword.
|
|
let is_trait = tokens
|
|
.iter()
|
|
.find_map(|token| match token {
|
|
TokenTree::Ident(ident) => match ident.to_string().as_str() {
|
|
"trait" => Some(true),
|
|
"impl" => Some(false),
|
|
_ => None,
|
|
},
|
|
_ => None,
|
|
})
|
|
.expect("#[vtable] attribute should only be applied to trait or impl block");
|
|
|
|
// Retrieve the main body. The main body should be the last token tree.
|
|
let body = match tokens.pop() {
|
|
Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Brace => group,
|
|
_ => panic!("cannot locate main body of trait or impl block"),
|
|
};
|
|
|
|
let mut body_it = body.stream().into_iter();
|
|
let mut functions = Vec::new();
|
|
let mut consts = HashSet::new();
|
|
while let Some(token) = body_it.next() {
|
|
match token {
|
|
TokenTree::Ident(ident) if ident.to_string() == "fn" => {
|
|
let fn_name = match body_it.next() {
|
|
Some(TokenTree::Ident(ident)) => ident.to_string(),
|
|
// Possibly we've encountered a fn pointer type instead.
|
|
_ => continue,
|
|
};
|
|
functions.push(fn_name);
|
|
}
|
|
TokenTree::Ident(ident) if ident.to_string() == "const" => {
|
|
let const_name = match body_it.next() {
|
|
Some(TokenTree::Ident(ident)) => ident.to_string(),
|
|
// Possibly we've encountered an inline const block instead.
|
|
_ => continue,
|
|
};
|
|
consts.insert(const_name);
|
|
}
|
|
_ => (),
|
|
}
|
|
}
|
|
|
|
let mut const_items;
|
|
if is_trait {
|
|
const_items = "
|
|
/// A marker to prevent implementors from forgetting to use [`#[vtable]`](vtable)
|
|
/// attribute when implementing this trait.
|
|
const USE_VTABLE_ATTR: ();
|
|
"
|
|
.to_owned();
|
|
|
|
for f in functions {
|
|
let gen_const_name = format!("HAS_{}", f.to_uppercase());
|
|
// Skip if it's declared already -- this allows user override.
|
|
if consts.contains(&gen_const_name) {
|
|
continue;
|
|
}
|
|
// We don't know on the implementation-site whether a method is required or provided
|
|
// so we have to generate a const for all methods.
|
|
write!(
|
|
const_items,
|
|
"/// Indicates if the `{f}` method is overridden by the implementor.
|
|
const {gen_const_name}: bool = false;",
|
|
)
|
|
.unwrap();
|
|
consts.insert(gen_const_name);
|
|
}
|
|
} else {
|
|
const_items = "const USE_VTABLE_ATTR: () = ();".to_owned();
|
|
|
|
for f in functions {
|
|
let gen_const_name = format!("HAS_{}", f.to_uppercase());
|
|
if consts.contains(&gen_const_name) {
|
|
continue;
|
|
}
|
|
write!(const_items, "const {gen_const_name}: bool = true;").unwrap();
|
|
}
|
|
}
|
|
|
|
let new_body = vec![const_items.parse().unwrap(), body.stream()]
|
|
.into_iter()
|
|
.collect();
|
|
tokens.push(TokenTree::Group(Group::new(Delimiter::Brace, new_body)));
|
|
tokens.into_iter().collect()
|
|
}
|