[PATCH v2 2/3] rust: const_eval: allow const trait method invocation in some contexts

[Gary Guo <gary@garyguo.net>]

Rust does not yet (as of 1.98) have stable const trait impl support; on
1.85 it does not have unstable support either. This leaves us having many
free functions that do the exact same as trait methods.

Create a `Const` wrapper type as the central place to put const trait
methods on as inherent methods, and a `const_call!` macro that dispatches
to these inherent methods as opposed to trait methods.

`const_call!` would check that the signature actually matches, and provides
some additional inference help. Once const trait impl is available, we can
also easily just convert the macro without having to touching all users at
once. Implement it as a proc macro, so it can accept the normal method
call syntax.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/kernel/const_eval.rs | 108 +++++++++++++++++++++++++++++++++++++++++++++-
 rust/kernel/prelude.rs    |   1 +
 rust/macros/const_eval.rs |  31 ++++++++++++-
 rust/macros/lib.rs        |   9 ++++
 4 files changed, 146 insertions(+), 3 deletions(-)

diff --git a/rust/kernel/const_eval.rs b/rust/kernel/const_eval.rs
index f1b79d82549d..2a6ea3f7b693 100644
--- a/rust/kernel/const_eval.rs
+++ b/rust/kernel/const_eval.rs
@@ -2,8 +2,114 @@
 
 //! Utilities for const evaluation.
 
+use core::{
+    ops::Deref, //
+};
+
 #[doc(inline)]
 pub use build_error::assert_in_const_eval;
 
 #[doc(inline)]
-pub use macros::const_eval_only;
+pub use macros::{
+    const_call,
+    const_eval_only, //
+};
+
+/// Constant wrapper type.
+///
+/// Rust does not yet (as of 1.98) have stable const trait impl support; on 1.85 it does not have
+/// unstable support either. Only inherent functions can be marked as const. There are a few const
+/// methods that we want to add on core types; and this mean that we cannot use extension trait on
+/// them.
+///
+/// This type serves as a middle layer. This type is local to the `kernel` crate, and thus we can
+/// define inherent methods on it. For core types, we will define it for `Const<Type>`. Other kernel
+/// crate types or even downstream types can also utilize it by defining inherent methods that
+/// *receive* `Const<Self>`.
+///
+/// Caller should use the `const_call!()` macro so it also checks that the type signature matches
+/// the trait.
+///
+/// # Examples
+///
+/// Say we want to define a extension method on u32. We can do
+/// ```no_run
+/// trait MyTrait {
+///     fn trait_method(self);
+/// }
+///
+/// impl MyTrait for u32 {
+///     fn trait_method(self) {
+///         /* impl */
+///     }
+/// }
+/// ```
+/// but we cannot mark it const.
+///
+/// Instead, we can do this
+/// ```ignore (doctest is outside kernel crate)
+/// trait MyTrait {
+///     fn trait_method(self);
+/// }
+///
+/// impl MyTrait for u32 {
+///     #[inline]
+///     fn trait_method(self) {
+///         // Forwarding impl
+///         Const(self).trait_method()
+///     }
+/// }
+///
+/// impl Const<u32> {
+///     pub const fn trait_method(self) {
+///         let Const(this) = self;
+///         /* impl */
+///     }
+/// }
+/// ```
+///
+/// For local or downstream types, implement it directly on the type with a different receiver:
+/// ```no_run
+/// # use kernel::const_eval::Const;
+/// trait MyTrait {
+///     fn trait_method(self);
+/// }
+///
+/// struct Foo;
+///
+/// impl MyTrait for Foo {
+///     #[inline]
+///     fn trait_method(self) {
+///         // Forwarding impl
+///         Const(self).trait_method()
+///     }
+/// }
+///
+/// impl Foo {
+///     pub const fn trait_method(self: Const<Self>) {
+///         let Const(this) = self;
+///         /* impl */
+///     }
+/// }
+/// ```
+///
+/// For caller of the method, one would simply replace `expr.method()` with `Const(expr).method()`;
+/// although the auto-ref coercion will be lost, so it a method expects `&self`, the caller would
+/// need to explicitly use the `Const(&expr).method()` syntax to call it.
+pub struct Const<T>(pub T);
+
+impl<T> Deref for Const<T> {
+    type Target = T;
+
+    #[inline]
+    fn deref(&self) -> &T {
+        &self.0
+    }
+}
+
+// Provide inference help only. Should never be code-generated.
+#[doc(hidden)]
+#[const_eval_only]
+pub const fn would_call<T, U, F: FnOnce(T) -> U>(_: T, _: F) -> U {
+    todo!()
+}
diff --git a/rust/kernel/prelude.rs b/rust/kernel/prelude.rs
index ca396f1f78a6..8f4c8cd2d8a2 100644
--- a/rust/kernel/prelude.rs
+++ b/rust/kernel/prelude.rs
@@ -91,6 +91,7 @@
         const_assert,
         static_assert, //
     },
+    const_eval::const_call,
     current,
     dev_alert,
     dev_crit,
diff --git a/rust/macros/const_eval.rs b/rust/macros/const_eval.rs
index 0664888d3b38..10a4a93f905a 100644
--- a/rust/macros/const_eval.rs
+++ b/rust/macros/const_eval.rs
@@ -1,9 +1,17 @@
 // SPDX-License-Identifier: GPL-2.0
 
-use proc_macro2::TokenStream;
-use quote::ToTokens;
+use proc_macro2::{
+    Span,
+    TokenStream, //
+};
+use quote::{
+    format_ident,
+    quote,
+    ToTokens, //
+};
 use syn::{
     parse_quote,
+    ExprMethodCall,
     ItemFn, //
 };
 
@@ -22,3 +30,22 @@ pub(crate) fn const_eval_only(mut input: ItemFn) -> TokenStream {
 
     input.into_token_stream()
 }
+
+pub(crate) fn const_call(mut input: ExprMethodCall) -> TokenStream {
+    let expr = input.receiver;
+
+    let expr_ident = format_ident!("expr", span = Span::mixed_site());
+    input.receiver = parse_quote!(#expr_ident);
+
+    let would_call = quote!(#input);
+    input.receiver = parse_quote!(::kernel::const_eval::Const(#expr_ident));
+
+    quote!({
+        let #expr_ident = #expr;
+        if false {
+            ::kernel::const_eval::would_call(#expr_ident, |#expr_ident| #would_call)
+        } else {
+            #input
+        }
+    })
+}
diff --git a/rust/macros/lib.rs b/rust/macros/lib.rs
index e47a8c35ccff..262539cccf38 100644
--- a/rust/macros/lib.rs
+++ b/rust/macros/lib.rs
@@ -339,6 +339,15 @@ pub fn concat_idents(input: TokenStream) -> TokenStream {
     concat_idents::concat_idents(parse_macro_input!(input)).into()
 }
 
+/// Call a trait method in const context.
+///
+/// This is a polyfill for Rust's const trait impl feature. Only work for specific methods that have
+/// dedicated const implementation.
+#[proc_macro]
+pub fn const_call(input: TokenStream) -> TokenStream {
+    const_eval::const_call(parse_macro_input!(input)).into()
+}
+
 /// Mark a function as usable from const evaluation only.
 ///
 /// Build will fail if the function is used for runtime code.

-- 
2.54.0