rust: support the ? operator in doctests

Now it is possible to have tests that use the `?` operator:

    /// ```
    /// # use kernel::{spawn_work_item, workqueue};
    /// spawn_work_item!(workqueue::system(), || pr_info!("x"))?;
    /// # Ok::<(), Error>(())
    /// ```

The logic is also simplified: instead of extracting the source
code and the originally-crate-level attributes separately,
wrap everything into an extra scope. This allows to reuse
the `rustdoc`-generated `unwrap()` and `Result` type.

This could break if `rustdoc` happens to emit an attribute that
can only be used at the crate-level, like a `feature(...)` one,
but the previous approach could also break in other ways.

There is a workaround for `Result` not being qualified in
the `rustdoc`-generated code -- it picks up our own `Result`,
and with the simplified logic here we cannot easily move our
`use` of the prelude inside the source code.

Suggested-by: Andreas Hindborg <andreas.hindborg@wdc.com>
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>

Author: ojeda

scripts/rustdoc_test_builder.py

@@ -13,45 +13,28 @@
 
 # `[^\s]*` removes the prefix (e.g. `_doctest_main_`) plus any
 # leading path (for `O=` builds).
-MAIN_RE = re.compile(
-    r"^"
-    r"fn main\(\) { "
-    r"#\[allow\(non_snake_case\)\] "
-    r"fn ([^\s]*rust_kernel_([a-zA-Z0-9_]+))\(\) {"
-    r"$"
-)
+TEST_NAME_RE = re.compile(r"fn [^\s]*rust_kernel_([a-zA-Z0-9_]+)\(\)")
 
 def main():
-    found_main = False
-    test_header = ""
-    test_body = ""
-    for line in sys.stdin.readlines():
-        main_match = MAIN_RE.match(line)
-        if main_match:
-            if found_main:
-                raise Exception("More than one `main` line found.")
-            found_main = True
-            function_name = main_match.group(1)
-            test_name = f"rust_kernel_doctest_{main_match.group(2)}"
-            continue
-
-        if found_main:
-            test_body += line
-        else:
-            test_header += line
-
-    if not found_main:
-        raise Exception("No `main` line found.")
-
-    call_line = f"}} {function_name}() }}"
-    if not test_body.endswith(call_line):
-        raise Exception("Unexpected end of test body.")
-    test_body = test_body[:-len(call_line)]
+    content = sys.stdin.read()
+    matches = TEST_NAME_RE.findall(content)
+    if len(matches) == 0:
+        raise Exception("No test name found.")
+    if len(matches) > 1:
+        raise Exception("More than one test name found.")
+
+    test_name = f"rust_kernel_doctest_{matches[0]}"
+
+    # Qualify `Result` to avoid the collision with our own `Result`
+    # coming from the prelude.
+    test_body = content.replace(
+        f'rust_kernel_{matches[0]}() -> Result<(), impl core::fmt::Debug> {{',
+        f'rust_kernel_{matches[0]}() -> core::result::Result<(), impl core::fmt::Debug> {{',
+    )
 
     with open(TESTS_DIR / f"{test_name}.json", "w") as fd:
         json.dump({
             "name": test_name,
-            "header": test_header,
             "body": test_body,
         }, fd, sort_keys=True, indent=4)
 

scripts/rustdoc_test_gen.py

@@ -29,6 +29,7 @@
     __KUNIT_TEST.store(__kunit_test, core::sync::atomic::Ordering::SeqCst);
 
     /// Overrides the usual [`assert!`] macro with one that calls KUnit instead.
+    #[allow(unused)]
     macro_rules! assert {{
         ($cond:expr $(,)?) => {{{{
             kernel::kunit_assert!(
@@ -39,6 +40,7 @@
     }}
 
     /// Overrides the usual [`assert_eq!`] macro with one that calls KUnit instead.
+    #[allow(unused)]
     macro_rules! assert_eq {{
         ($left:expr, $right:expr $(,)?) => {{{{
             kernel::kunit_assert_eq!(
@@ -50,9 +52,13 @@
     }}
 
     // Many tests need the prelude, so provide it by default.
+    #[allow(unused)]
     use kernel::prelude::*;
 
-    {test_body}
+    {{
+        {test_body}
+        main();
+    }}
 }}
 """
 RUST_TEMPLATE = """// SPDX-License-Identifier: GPL-2.0
@@ -93,8 +99,6 @@
 // an `AtomicPtr` to hold the context (though each test only writes once before
 // threads may be created).
 
-{rust_header}
-
 const __LOG_PREFIX: &[u8] = b"rust_kernel_doctests\\0";
 
 {rust_tests}
@@ -127,15 +131,12 @@
 """
 
 def main():
-    rust_header = set()
     rust_tests = ""
     c_test_declarations = ""
     c_test_cases = ""
     for filename in sorted(os.listdir(TESTS_DIR)):
         with open(TESTS_DIR / filename, "r") as fd:
             test = json.load(fd)
-            for line in test["header"].strip().split("\n"):
-                rust_header.add(line)
             rust_tests += RUST_TEMPLATE_TEST.format(
                 test_name = test["name"],
                 test_body = test["body"]
@@ -146,11 +147,9 @@ def main():
             c_test_cases += C_TEMPLATE_TEST_CASE.format(
                 test_name = test["name"]
             )
-    rust_header = sorted(rust_header)
 
     with open(RUST_FILE, "w") as fd:
         fd.write(RUST_TEMPLATE.format(
-            rust_header = "\n".join(rust_header).strip(),
             rust_tests = rust_tests.strip(),
         ))