[RFC PATCH 1/7] rust: types: Introduce HasField trait and derive macro

[Boqun Feng <boqun.feng@gmail.com>]

In order to unify all the `Has*` infrastrutures, a generic `HasField`
trait is added along with a derive macro `#[derive(HasField)]` which
allows generate implementation of `HasField` automatically, e.g.

    #[derive(HasField)]
    struct Base {
        a: i32;
	b: i32;
	#[field]
	f1: Field<Base, 1>,
	#[field]
	f2: Field<Base, 2>,
    }

two implementations `impl HasField<Base, Field<Base, 1>>` and `impl
HasField<Base, Field<Base, 2>>` will be generated with `&raw mut` and
`container_of!()`.

This simplifies the usage of the current `Has*` traits, namely `HasWork`
and `HasHrTimer`, and eases the introduction of more `Field` type in the
future.

Signed-off-by: Boqun Feng <boqun.feng@gmail.com>
---
 rust/kernel/field.rs   | 73 ++++++++++++++++++++++++++++++++++++
 rust/kernel/lib.rs     |  1 +
 rust/kernel/prelude.rs |  4 +-
 rust/macros/field.rs   | 85 ++++++++++++++++++++++++++++++++++++++++++
 rust/macros/lib.rs     | 11 ++++++
 5 files changed, 173 insertions(+), 1 deletion(-)
 create mode 100644 rust/kernel/field.rs
 create mode 100644 rust/macros/field.rs

diff --git a/rust/kernel/field.rs b/rust/kernel/field.rs
new file mode 100644
index 000000000000..347387731d71
--- /dev/null
+++ b/rust/kernel/field.rs
@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Field types to describe a field inside a struct.
+
+/// A field.
+///
+/// The generic type `T` is usually the type that contains the field. For some field types, it
+/// needs to be generic over the type containing it, because it needs to be initialized with
+/// container-type-specific callbacks. For other types, simply implement [`Field<T>`] for all `T`
+/// to indicate there is no restriction.
+pub trait Field<T>: Sized {}
+
+/// A struct `T` that has a field `F`.
+///
+/// # Safety
+///
+/// The methods [`raw_get_field()`] and [`field_container_of()`] must return valid pointers and
+/// must be true inverses of each other; that is, they must satisfy the following invariants: -
+/// `field_container_of(raw_get_field(ptr)) == ptr` for any `ptr: *mut Self`. -
+/// `raw_get_field(field_container_of(ptr)) == ptr` for any `ptr: *mut Field<T>`.
+///
+/// Use [`macros::HasField`] to generate the impls automatically.
+///
+/// # Examples
+///
+/// ```
+/// # use core::marker::PhantomData;
+/// use kernel::{
+///     macros::HasField,
+///     field::{
+///         Field,
+///         HasField, //
+///     }, //
+/// };
+///
+/// struct Work<T, const ID: u64> {
+///     _x: isize,
+///     _inner: PhantomData<T>,
+/// }
+///
+/// // Declare that `Work` is a `Field`.
+/// impl<T, const ID: u64> Field<T> for Work<T, ID> {}
+///
+/// #[derive(HasField)]
+/// struct B {
+///     #[field]
+///     w: Work<B, 2>,
+///     a: i32,
+/// }
+///
+/// const _: () = {
+///     const fn assert_has_field<T: HasField<T, Work<T, 2>>>() { }
+///     assert_has_field::<B>();
+/// };
+/// ```
+///
+/// [`raw_get_field()`]: HasField::raw_get_field
+/// [`field_container_of()`]: HasField::field_container_of
+pub unsafe trait HasField<T, F: Field<T>> {
+    /// Returns a pointer to the [`Field<T>`] field.
+    ///
+    /// # Safety
+    ///
+    /// The provided pointer must point at a valid struct of type `Self`.
+    unsafe fn raw_get_field(ptr: *mut Self) -> *mut F;
+
+    /// Returns a pointer to the struct containing [`Field<T>`] field.
+    ///
+    /// # Safety
+    ///
+    /// The pointer must point at a [`Field<T>`] field in a struct of type `Self`.
+    unsafe fn field_container_of(ptr: *mut F) -> *mut Self;
+}
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index f812cf120042..36259aac1843 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -93,6 +93,7 @@
 pub mod drm;
 pub mod error;
 pub mod faux;
+pub mod field;
 #[cfg(CONFIG_RUST_FW_LOADER_ABSTRACTIONS)]
 pub mod firmware;
 pub mod fmt;
diff --git a/rust/kernel/prelude.rs b/rust/kernel/prelude.rs
index 2877e3f7b6d3..3668ef42046b 100644
--- a/rust/kernel/prelude.rs
+++ b/rust/kernel/prelude.rs
@@ -25,10 +25,12 @@
 pub use crate::alloc::{flags::*, Box, KBox, KVBox, KVVec, KVec, VBox, VVec, Vec};
 
 #[doc(no_inline)]
-pub use macros::{export, fmt, kunit_tests, module, vtable};
+pub use macros::{export, fmt, kunit_tests, module, vtable, HasField};
 
 pub use pin_init::{init, pin_data, pin_init, pinned_drop, InPlaceWrite, Init, PinInit, Zeroable};
 
+pub use super::field::{Field, HasField};
+
 pub use super::{build_assert, build_error};
 
 // `super::std_vendor` is hidden, which makes the macro inline for some reason.
diff --git a/rust/macros/field.rs b/rust/macros/field.rs
new file mode 100644
index 000000000000..3d32e5089f27
--- /dev/null
+++ b/rust/macros/field.rs
@@ -0,0 +1,85 @@
+// SPDX-License-Identifier: GPL-2.0
+
+use proc_macro2::{Span, TokenStream};
+use quote::quote;
+use syn::{
+    spanned::Spanned, Data, DataStruct, DeriveInput, Error, Fields, Generics, Ident, Result, Type,
+};
+
+fn impl_has_field(base: &Ident, field: &Ident, ty: &Type, generics: &Generics) -> TokenStream {
+    let (impl_generics, type_generics, where_clause) = generics.split_for_impl();
+
+    quote!(
+        // SAFETY: The implementation of `raw_get_field()` only compiles if the field has the
+        // right type.
+        unsafe impl #impl_generics
+        HasField<#base #type_generics, #ty>
+        for #base #type_generics
+        #where_clause {
+            #[inline(always)]
+            unsafe fn raw_get_field(ptr: *mut Self) -> *mut #ty {
+                // SAFETY: Per function safety requirement, the pointer is valid.
+                unsafe { &raw mut (*ptr).#field }
+            }
+
+            #[inline(always)]
+            unsafe fn field_container_of(ptr: *mut #ty) -> *mut Self {
+                // SAFETY: Per function safety requirement, the pointer is valid, and it points
+                // to the right field of the struct.
+                unsafe { kernel::container_of!(ptr, Self, #field) }
+            }
+        }
+    )
+}
+fn handle_struct(
+    ident: &Ident,
+    generics: &Generics,
+    st: &DataStruct,
+    span: Span,
+) -> Result<TokenStream> {
+    let mut impls = vec![];
+
+    if let Fields::Named(fields) = &st.fields {
+        for field in &fields.named {
+            let found = field
+                .attrs
+                .iter()
+                .find(|attr| attr.path().is_ident("field"));
+
+            if found.is_some() {
+                if let Some(name) = &field.ident {
+                    impls.push(impl_has_field(ident, name, &field.ty, generics));
+                }
+            }
+        }
+
+        Ok(quote!(
+            #(#impls)*
+        ))
+    } else {
+        Err(Error::new(
+            span,
+            "`#[derive(HasField)]` only supports structs with named fields",
+        ))
+    }
+}
+
+pub(crate) fn has_field(input: DeriveInput) -> Result<TokenStream> {
+    let span = input.span();
+    let data = &input.data;
+    let ident = &input.ident;
+    let generics = &input.generics;
+
+    if let Data::Struct(st) = data {
+        let impls = handle_struct(ident, generics, st, span)?;
+
+        Ok(quote!(
+            #impls
+        ))
+    } else {
+        Err(Error::new_spanned(
+            input,
+            "`#[derive(HasField)]` only supports structs",
+        ))
+    }
+}
diff --git a/rust/macros/lib.rs b/rust/macros/lib.rs
index 85b7938c08e5..4fccca0e11af 100644
--- a/rust/macros/lib.rs
+++ b/rust/macros/lib.rs
@@ -13,6 +13,7 @@
 
 mod concat_idents;
 mod export;
+mod field;
 mod fmt;
 mod helpers;
 mod kunit;
@@ -486,3 +487,13 @@ pub fn kunit_tests(attr: TokenStream, input: TokenStream) -> TokenStream {
         .unwrap_or_else(|e| e.into_compile_error())
         .into()
 }
+
+/// Derives the implementation for `HasField`.
+///
+/// See the documentation of `HasField` for more information.
+#[proc_macro_derive(HasField, attributes(field))]
+pub fn has_field(input: TokenStream) -> TokenStream {
+    field::has_field(parse_macro_input!(input))
+        .unwrap_or_else(|e| e.into_compile_error())
+        .into()
+}
-- 
2.50.1 (Apple Git-155)