[PATCH v2 1/3] rust: const_eval: add `#[const_eval_only]` attribute

[Gary Guo]

We have a lot of helper const functions which are intended to be used
during const evaluation only and runtime calls should not be generated. Add
a macro to denote this explicitly. This is similar to C++'s consteval
keyword.

Convert device_id.rs as an example.

Signed-off-by: Gary Guo <gary@xxxxxxxxxxx>
---
 rust/build_error.rs       |  7 +++++++
 rust/kernel/const_eval.rs |  9 +++++++++
 rust/kernel/device_id.rs  |  4 ++++
 rust/kernel/lib.rs        |  1 +
 rust/macros/const_eval.rs | 24 ++++++++++++++++++++++++
 rust/macros/lib.rs        | 21 +++++++++++++++++++++
 6 files changed, 66 insertions(+)

diff --git a/rust/build_error.rs b/rust/build_error.rs
index fa24eeef9929..b7ef80596f1f 100644
--- a/rust/build_error.rs
+++ b/rust/build_error.rs
@@ -29,3 +29,10 @@
 pub const fn build_error(msg: &'static str) -> ! {
     panic!("{}", msg);
 }
+
+/// Assert that the code is in const evaluation.
+///
+/// Triggers a build error if called at runtime.
+#[inline(never)]
+#[export_name = "rust_const_eval_called_at_runtime"]
+pub const fn assert_in_const_eval() {}
diff --git a/rust/kernel/const_eval.rs b/rust/kernel/const_eval.rs
new file mode 100644
index 000000000000..f1b79d82549d
--- /dev/null
+++ b/rust/kernel/const_eval.rs
@@ -0,0 +1,9 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Utilities for const evaluation.
+
+#[doc(inline)]
+pub use build_error::assert_in_const_eval;
+
+#[doc(inline)]
+pub use macros::const_eval_only;
diff --git a/rust/kernel/device_id.rs b/rust/kernel/device_id.rs
index c81fca5b4986..dad9cadaeb1b 100644
--- a/rust/kernel/device_id.rs
+++ b/rust/kernel/device_id.rs
@@ -10,6 +10,8 @@
     mem::MaybeUninit, //
 };
 
+use crate::const_eval::const_eval_only;
+
 /// Marker trait to indicate a Rust device ID type represents a corresponding C device ID type.
 ///
 /// This is meant to be implemented by buses/subsystems so that they can use [`IdTable`] to
@@ -108,6 +110,7 @@ impl<T: RawDeviceId + RawDeviceIdIndex, U: 'static, const N: usize> IdArray<T, U
     /// Creates a new instance of the array.
     ///
     /// The contents are derived from the given identifiers and context information.
+    #[const_eval_only]
     pub const fn new(ids: [(T, &'static U); N]) -> Self {
         let mut raw_ids = [const { MaybeUninit::<T::RawType>::uninit() }; N];
 
@@ -144,6 +147,7 @@ impl<T: RawDeviceId, const N: usize> IdArray<T, (), N> {
     ///
     /// The contents are derived from the given identifiers and context information.
     /// If the device implements [`RawDeviceIdIndex`], consider using [`IdArray::new`] instead.
+    #[const_eval_only]
     pub const fn new_without_index(ids: [T; N]) -> Self {
         // SAFETY: `T` is layout-wise compatible with `T::RawType`, so is the array of them.
         let raw_ids: [MaybeUninit<T::RawType>; N] = unsafe { core::mem::transmute_copy(&ids) };
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 4d5c96ddc49c..d9ed25e96ff3 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -57,6 +57,7 @@
 pub mod clk;
 #[cfg(CONFIG_CONFIGFS_FS)]
 pub mod configfs;
+pub mod const_eval;
 pub mod cpu;
 #[cfg(CONFIG_CPU_FREQ)]
 pub mod cpufreq;
diff --git a/rust/macros/const_eval.rs b/rust/macros/const_eval.rs
new file mode 100644
index 000000000000..0664888d3b38
--- /dev/null
+++ b/rust/macros/const_eval.rs
@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: GPL-2.0
+
+use proc_macro2::TokenStream;
+use quote::ToTokens;
+use syn::{
+    parse_quote,
+    ItemFn, //
+};
+
+pub(crate) fn const_eval_only(mut input: ItemFn) -> TokenStream {
+    // Prevent code generation as the function is for const evaluation only.
+    input.attrs.push(parse_quote!(
+        #[inline(always)]
+    ));
+
+    input.block.stmts.insert(
+        0,
+        parse_quote!(
+            ::kernel::const_eval::assert_in_const_eval();
+        ),
+    );
+
+    input.into_token_stream()
+}
diff --git a/rust/macros/lib.rs b/rust/macros/lib.rs
index 24f96feaeb34..e47a8c35ccff 100644
--- a/rust/macros/lib.rs
+++ b/rust/macros/lib.rs
@@ -15,6 +15,7 @@
 #![cfg_attr(not(CONFIG_RUSTC_HAS_SPAN_FILE), feature(proc_macro_span))]
 
 mod concat_idents;
+mod const_eval;
 mod export;
 mod fmt;
 mod for_lt;
@@ -338,6 +339,26 @@ pub fn concat_idents(input: TokenStream) -> TokenStream {
     concat_idents::concat_idents(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.
+///
+/// # Examples
+///
+/// ```
+/// #[const_eval_only]
+/// const fn call_for_const_eval_only() {
+///     // This code will be executed only during const eval!
+/// }
+///
+/// const _: () = call_for_const_eval_only();
+/// ```
+#[proc_macro_attribute]
+pub fn const_eval_only(attr: TokenStream, input: TokenStream) -> TokenStream {
+    parse_macro_input!(attr as syn::parse::Nothing);
+    const_eval::const_eval_only(parse_macro_input!(input)).into()
+}
+
 /// Paste identifiers together.
 ///
 /// Within the `paste!` macro, identifiers inside `[<` and `>]` are concatenated together to form a

-- 
2.54.0