3 files changed, 27 insertions, 55 deletions

diff --git a/scripts/remove-stale-files b/scripts/remove-stale-files
index 6e39fa8540df..f1f776caae00 100755
--- a/scripts/remove-stale-files
+++ b/scripts/remove-stale-files
@@ -26,3 +26,5 @@ rm -f scripts/selinux/genheaders/genheaders
 rm -f *.spec
 
 rm -f lib/test_fortify.log
+
+rm -f rust/doctests_kernel_generated_kunit.c
diff --git a/scripts/rustdoc_test_builder.rs b/scripts/rustdoc_test_builder.rs
index df864437cef7..0361e22853f2 100644
--- a/scripts/rustdoc_test_builder.rs
+++ b/scripts/rustdoc_test_builder.rs
@@ -72,6 +72,13 @@ fn main() {
     // Figure out a smaller test name based on the generated function name.
     let name = rustdoc_function_name.split_once("_rust_kernel_").unwrap().1;
 
+    // The `rustdoc` function name can include the absolute path when building with `O=` which is
+    // undesirable and creates overlong symbol names. Remap it to use a relative path.
+    // TODO: Remove when bumping the minimum version since with Rust >= 1.95 we already do the
+    // path remapping via the `rustdoc` `--remap-path-prefix` flag.
+    let trimmed_function_name = format!("_doctest_main_rust_kernel_{name}");
+    let body = body.replace(&rustdoc_function_name, &trimmed_function_name);
+
     let path = format!("rust/test/doctests/kernel/{name}");
 
     std::fs::write(path, body.as_bytes()).unwrap();
diff --git a/scripts/rustdoc_test_gen.rs b/scripts/rustdoc_test_gen.rs
index d087c0d9fcb3..754c5aafcce4 100644
--- a/scripts/rustdoc_test_gen.rs
+++ b/scripts/rustdoc_test_gen.rs
@@ -125,9 +125,7 @@ fn main() {
     // Sort paths.
     paths.sort();
 
-    let mut rust_tests = String::new();
-    let mut c_test_declarations = String::new();
-    let mut c_test_cases = String::new();
+    let mut tests = String::new();
     let mut body = String::new();
     let mut last_file = String::new();
     let mut number = 0;
@@ -172,28 +170,22 @@ fn main() {
 
         use std::fmt::Write;
         write!(
-            rust_tests,
+            tests,
             r#"/// Generated `{name}` KUnit test case from a Rust documentation test.
-#[no_mangle]
-pub extern "C" fn {kunit_name}(__kunit_test: *mut ::kernel::bindings::kunit) {{
-    /// Overrides the usual [`assert!`] macro with one that calls KUnit instead.
+#[test]
+fn {kunit_name}() {{
+    /// Overrides the usual [`file!`] macro with one that expands to the real path.
     #[allow(unused)]
-    macro_rules! assert {{
-        ($cond:expr $(,)?) => {{{{
-            ::kernel::kunit_assert!(
-                "{kunit_name}", c"{real_path}", __DOCTEST_ANCHOR - {line}, $cond
-            );
-        }}}}
+    macro_rules! file {{
+        () => {{ "{real_path}" }}
     }}
 
-    /// Overrides the usual [`assert_eq!`] macro with one that calls KUnit instead.
+    /// Overrides the usual [`line!`] macro with one that expands to the real line number.
     #[allow(unused)]
-    macro_rules! assert_eq {{
-        ($left:expr, $right:expr $(,)?) => {{{{
-            ::kernel::kunit_assert_eq!(
-                "{kunit_name}", c"{real_path}", __DOCTEST_ANCHOR - {line}, $left, $right
-            );
-        }}}}
+    macro_rules! line {{
+        // NOTE: This does not expand to a literal, but a constant expression.
+        // Therefore code that depends on `line!()` expanding to a literal needs special adjustment.
+        () => {{ const {{ ::core::line!() - __DOCTEST_ANCHOR + {line} }} }}
     }}
 
     // Many tests need the prelude, so provide it by default.
@@ -213,7 +205,7 @@ pub extern "C" fn {kunit_name}(__kunit_test: *mut ::kernel::bindings::kunit) {{
 
     /// The anchor where the test code body starts.
     #[allow(unused)]
-    static __DOCTEST_ANCHOR: i32 = ::core::line!() as i32 + {body_offset} + 2;
+    static __DOCTEST_ANCHOR: u32 = ::core::line!() + {body_offset} + 2;
     {{
         #![allow(unreachable_pub, clippy::disallowed_names)]
         {body}
@@ -224,14 +216,9 @@ pub extern "C" fn {kunit_name}(__kunit_test: *mut ::kernel::bindings::kunit) {{
 "#
         )
         .unwrap();
-
-        write!(c_test_declarations, "void {kunit_name}(struct kunit *);\n").unwrap();
-        write!(c_test_cases, "    KUNIT_CASE({kunit_name}),\n").unwrap();
     }
 
-    let rust_tests = rust_tests.trim();
-    let c_test_declarations = c_test_declarations.trim();
-    let c_test_cases = c_test_cases.trim();
+    let tests = tests.trim();
 
     write!(
         BufWriter::new(File::create("rust/doctests_kernel_generated.rs").unwrap()),
@@ -257,34 +244,10 @@ impl ModuleMetadata for LocalModule {{
     }};
 }}
 
-{rust_tests}
-"#
-    )
-    .unwrap();
-
-    write!(
-        BufWriter::new(File::create("rust/doctests_kernel_generated_kunit.c").unwrap()),
-        r#"/*
- * `kernel` crate documentation tests.
- */
-
-#include <kunit/test.h>
-
-{c_test_declarations}
-
-static struct kunit_case test_cases[] = {{
-    {c_test_cases}
-    {{ }}
-}};
-
-static struct kunit_suite test_suite = {{
-    .name = "rust_doctests_kernel",
-    .test_cases = test_cases,
-}};
-
-kunit_test_suite(test_suite);
-
-MODULE_LICENSE("GPL");
+#[kernel::macros::kunit_tests(rust_doctests_kernel)]
+mod tests {{
+{tests}
+}}
 "#
     )
     .unwrap();