rust: implement a #[vtable] macro

Use a single `#[vtable]` macro to replace the current boilerplating
of `ToUse`, `USE_NONE`, `declare_file_operations` required for
declaring and implementing traits that maps to Linux's pure vtables and
contains optional methods.

Signed-off-by: Gary Guo <[email protected]>

Author: nbdd0121

rust/kernel/prelude.rs

@@ -15,7 +15,7 @@ pub use core::pin::Pin;
 
 pub use alloc::{boxed::Box, string::String, vec::Vec};
 
-pub use macros::module;
+pub use macros::{module, vtable};
 
 pub use super::build_assert;
 

rust/macros/lib.rs

@@ -4,6 +4,7 @@
 
 mod helpers;
 mod module;
+mod vtable;
 
 use proc_macro::TokenStream;
 
@@ -92,3 +93,54 @@ use proc_macro::TokenStream;
 pub fn module(ts: TokenStream) -> TokenStream {
     module::module(ts)
 }
+
+/// Declares or implements a vtable trait.
+///
+/// Linux's use of pure vtables is very close to Rust traits, but they differ
+/// in how unimplemented functions are represented. In Rust, traits can provide
+/// default implementation for all non-required methods (and the default
+/// implementation could just return `Error::EINVAL`); Linux typically use C
+/// `NULL` pointers to represent these functions.
+///
+/// This attribute is intended to close the gap. Traits can be declared and
+/// implemented with the `#[vtable]` attribute, and a `HAS_*` associated constant
+/// will be generated for each method in the trait, indicating if the implementor
+/// has overriden a method.
+///
+/// This attribute is not needed if all methods are required.
+///
+/// # Examples
+///
+/// ```ignore
+/// use kernel::prelude::*;
+///
+/// // Declares a `#[vtable]` trait
+/// #[vtable]
+/// pub trait Operations: Send + Sync + Sized {
+///     fn foo(&self) -> Result<()> {
+///         Err(EINVAL)
+///     }
+///
+///     fn bar(&self) -> Result<()> {
+///         Err(EINVAL)
+///     }
+/// }
+///
+/// struct Foo;
+///
+/// // Implements the `#[vtable]` trait
+/// #[vtable]
+/// impl Operations for Foo {
+///     fn foo(&self) -> Result<()> {
+/// #        Err(EINVAL)
+///         /* ... */
+///     }
+/// }
+///
+/// assert_eq!(<Foo as Operations>::HAS_FOO, true);
+/// assert_eq!(<Foo as Operations>::HAS_BAR, false);
+/// ```
+#[proc_macro_attribute]
+pub fn vtable(attr: TokenStream, ts: TokenStream) -> TokenStream {
+    vtable::vtable(attr, ts)
+}

rust/macros/vtable.rs

@@ -0,0 +1,90 @@
+// 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 overriden by the implementor.
+                const {gen_const_name}: bool = false;",
+            )
+            .unwrap();
+        }
+    } 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()
+}