1 files changed, 164 insertions, 0 deletions
diff --git a/scripts/rustdoc_test_gen.py b/scripts/rustdoc_test_gen.py
new file mode 100755
index 000000000000..ad9a94293ab5
--- /dev/null
+++ b/scripts/rustdoc_test_gen.py
@@ -0,0 +1,164 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: GPL-2.0
+"""rustdoc_test_gen - Generates KUnit tests from saved `rustdoc`-generated tests.
+"""
+
+import json
+import os
+import pathlib
+
+RUST_DIR = pathlib.Path("rust")
+TESTS_DIR = RUST_DIR / "test" / "doctests" / "kernel"
+
+RUST_FILE = RUST_DIR / "doctests_kernel_generated.rs"
+C_FILE = RUST_DIR / "doctests_kernel_generated_kunit.c"
+
+RUST_TEMPLATE_TEST = """
+/// Generated `{test_name}` KUnit test case from a Rust documentation test.
+#[no_mangle]
+pub fn {test_name}(__kunit_test: *mut kernel::bindings::kunit) {{
+    /// Provides mutual exclusion (see `# Implementation` notes).
+    static __KUNIT_TEST_MUTEX: kernel::sync::smutex::Mutex<()> =
+        kernel::sync::smutex::Mutex::new(());
+
+    /// Saved argument (see `# Implementation` notes).
+    static __KUNIT_TEST: core::sync::atomic::AtomicPtr<kernel::bindings::kunit> =
+        core::sync::atomic::AtomicPtr::new(core::ptr::null_mut());
+
+    let __kunit_test_mutex_guard = __KUNIT_TEST_MUTEX.lock();
+    __KUNIT_TEST.store(__kunit_test, core::sync::atomic::Ordering::SeqCst);
+
+    /// Overrides the usual [`assert!`] macro with one that calls KUnit instead.
+    macro_rules! assert {{
+        ($cond:expr $(,)?) => {{{{
+            kernel::kunit_assert!(
+                __KUNIT_TEST.load(core::sync::atomic::Ordering::SeqCst),
+                $cond
+            );
+        }}}}
+    }}
+
+    /// Overrides the usual [`assert_eq!`] macro with one that calls KUnit instead.
+    macro_rules! assert_eq {{
+        ($left:expr, $right:expr $(,)?) => {{{{
+            kernel::kunit_assert_eq!(
+                __KUNIT_TEST.load(core::sync::atomic::Ordering::SeqCst),
+                $left,
+                $right
+            );
+        }}}}
+    }}
+
+    // Many tests need the prelude, so provide it by default.
+    use kernel::prelude::*;
+
+    {test_body}
+}}
+"""
+RUST_TEMPLATE = """// SPDX-License-Identifier: GPL-2.0
+
+//! `kernel` crate documentation tests.
+
+// # Implementation
+//
+// KUnit gives us a context in the form of the `kunit_test` parameter that one
+// needs to pass back to other KUnit functions and macros.
+//
+// However, we want to keep this as an implementation detail because:
+//
+//   - Test code should not care about the implementation.
+//
+//   - Documentation looks worse if it needs to carry extra details unrelated
+//     to the piece being described.
+//
+//   - Test code should be able to define functions and call them, without
+//     having to carry the context (since functions cannot capture dynamic
+//     environment).
+//
+//   - Later on, we may want to be able to test non-kernel code (e.g. `core`,
+//     `alloc` or external crates) which likely use the standard library
+//     `assert*!` macros.
+//
+// For this reason, `static`s are used in the generated code to save the
+// argument which then gets read by the asserting macros. These macros then
+// call back into KUnit, instead of panicking.
+//
+// To avoid depending on whether KUnit allows to run tests concurrently and/or
+// reentrantly, we ensure mutual exclusion on our end. To ensure a single test
+// being killed does not trigger failure of every other test (timing out),
+// we provide different `static`s per test (which also allow for concurrent
+// execution, though KUnit runs them sequentially).
+//
+// Furthermore, since test code may create threads and assert from them, we use
+// 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}
+"""
+
+C_TEMPLATE_TEST_DECLARATION = "void {test_name}(struct kunit *);\n"
+C_TEMPLATE_TEST_CASE = "    KUNIT_CASE({test_name}),\n"
+C_TEMPLATE = """// SPDX-License-Identifier: GPL-2.0
+/*
+ * `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_kernel_doctests",
+    .test_cases = test_cases,
+}};
+
+kunit_test_suite(test_suite);
+
+MODULE_LICENSE("GPL");
+"""
+
+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"]
+            )
+            c_test_declarations += C_TEMPLATE_TEST_DECLARATION.format(
+                test_name = test["name"]
+            )
+            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(),
+        ))
+
+    with open(C_FILE, "w") as fd:
+        fd.write(C_TEMPLATE.format(
+            c_test_declarations=c_test_declarations.strip(),
+            c_test_cases=c_test_cases.strip(),
+        ))
+
+if __name__ == "__main__":
+    main()