[PATCH 5/5] rust: print: add `continuable` as option

[Gary Guo]

Most users do not need the print to be continuable; have users that need
continuability of printk mark themselves as needing so explicitly.

Other users are assumed to actually want the log lines being committed
immediately. In printk, the trailing newline is treated more like a control
flag rather than actual new line. It is stripped way eventually. With
"continuable" not specified, naturally we should add "\n" automatically to
the print to convey this info and without requiring extra "\n" to be
specified. It is also natural to kernel developers with userspace Rust
background as neither `log` nor `tracing` crates needs trailing "\n". On
the other hand kernel C developers would add the trailing "\n" (or maybe
not, given many C code is missing trailing "\n" too), so it is not ideal to
require it being not present (or adding a redundant newline). Therefore,
make it optional; they will mean the same thing when the "continuable"
option is not present.

Suggested-by: Alice Ryhl <aliceryhl@xxxxxxxxxx>
Signed-off-by: Gary Guo <gary@xxxxxxxxxxx>
---
 lib/find_bit_benchmark_rust.rs  |  2 +-
 rust/kernel/print.rs            | 82 ++++++++++++++++++++++++++++++++++++++---
 samples/rust/rust_print_main.rs |  4 +-
 3 files changed, 80 insertions(+), 8 deletions(-)

diff --git a/lib/find_bit_benchmark_rust.rs b/lib/find_bit_benchmark_rust.rs
index 6bdc51de2f30..7fd7d4b26d0e 100644
--- a/lib/find_bit_benchmark_rust.rs
+++ b/lib/find_bit_benchmark_rust.rs
@@ -62,7 +62,7 @@ fn test_next_zero_bit(bitmap: &BitmapVec) {
 }
 
 fn find_bit_test() {
-    pr_err!("Benchmark");
+    pr_err!(continuable, "Benchmark");
     pr_cont!("\nStart testing find_bit() Rust with random-filled bitmap");
 
     let mut bitmap = BitmapVec::new(BITMAP_LEN, GFP_KERNEL).expect("alloc bitmap failed");
diff --git a/rust/kernel/print.rs b/rust/kernel/print.rs
index 2c82f5e82700..0382a0bf71ec 100644
--- a/rust/kernel/print.rs
+++ b/rust/kernel/print.rs
@@ -42,6 +42,7 @@ pub mod level {
     use core::ffi::CStr;
 
     pub trait Level {
+        const FORMAT_STRING_CONTINUABLE: &'static CStr;
         const FORMAT_STRING: &'static CStr;
     }
 
@@ -51,7 +52,7 @@ pub trait Level {
     /// given `prefix`, which are the kernel's `KERN_*` constants.
     ///
     /// [`_printk`]: srctree/include/linux/printk.h
-    const fn generate(prefix: &[u8; 3]) -> [u8; 10] {
+    const fn generate_continuable(prefix: &[u8; 3]) -> [u8; 10] {
         // Ensure the `KERN_*` macros are what we expect.
         assert!(prefix.len() == 3);
         assert!(prefix[0] == b'\x01');
@@ -63,12 +64,31 @@ pub trait Level {
         fmt
     }
 
+    const fn generate(prefix: &[u8; 3]) -> [u8; 11] {
+        // Ensure the `KERN_*` macros are what we expect.
+        assert!(prefix.len() == 3);
+        assert!(prefix[0] == b'\x01');
+        assert!(prefix[1] >= b'0' && prefix[1] <= b'7');
+
+        let mut fmt = *b"\0\0%s: %pA\n\0";
+        fmt[0] = prefix[0];
+        fmt[1] = prefix[1];
+        fmt
+    }
+
     // #[rustfmt::skip] // Rustfmt formats the macro awkwardly.
     macro_rules! define_level {
         ($lvl:ident) => {
             pub struct $lvl;
 
             impl Level for $lvl {
+                const FORMAT_STRING_CONTINUABLE: &'static CStr = match CStr::from_bytes_with_nul(
+                    &generate_continuable(macros::paste!(bindings::[<KERN_ $lvl>])),
+                ) {
+                    Ok(v) => v,
+                    Err(_) => unreachable!(),
+                };
+
                 const FORMAT_STRING: &'static CStr = match CStr::from_bytes_with_nul(&generate(
                     macros::paste!(bindings::[<KERN_ $lvl>]),
                 )) {
@@ -89,6 +109,26 @@ impl Level for $lvl {
     define_level!(DEBUG);
 }
 
+/// Prints a message via the kernel's [`_printk`] without newline.
+///
+/// Public but hidden since it should only be used from public macros.
+///
+/// [`_printk`]: srctree/include/linux/_printk.h
+#[doc(hidden)]
+#[cfg_attr(not(CONFIG_PRINTK), allow(unused_variables))]
+pub fn call_printk_continuable<Lvl: level::Level>(module_name: &CStr, args: fmt::Arguments<'_>) {
+    // `_printk` does not seem to fail in any path.
+    #[cfg(CONFIG_PRINTK)]
+    // SAFETY: TODO.
+    unsafe {
+        bindings::_printk(
+            Lvl::FORMAT_STRING_CONTINUABLE.as_char_ptr(),
+            module_name.as_char_ptr(),
+            core::ptr::from_ref(&args).cast::<c_void>(),
+        );
+    }
+}
+
 /// Prints a message via the kernel's [`_printk`].
 ///
 /// Public but hidden since it should only be used from public macros.
@@ -147,16 +187,48 @@ pub fn call_printk_cont(args: fmt::Arguments<'_>) {
 #[macro_export]
 #[expect(clippy::crate_in_macro_def)]
 macro_rules! print_macro (
-    ($level:ident, $($arg:tt)+) => (
-        match $crate::prelude::fmt!($($arg)+) {
+    ($level:ident, continuable, $fmt:literal $($arg:tt)*) => (
+        match $crate::prelude::fmt!($fmt $($arg)*) {
             args => {
-                $crate::print::call_printk::<$crate::print::level::$level>(
+                // Perform a check on the format string to ensure last byte is not "\n",
+                // otherwise, the continuable option is pointless.
+                const {
+                    let bytes = str::as_bytes($fmt);
+                    ::core::assert!(
+                        bytes[bytes.len() - 1] != b'\n',
+                        r#"continuable format should not terminate with "\n""#
+                    );
+                }
+
+                $crate::print::call_printk_continuable::<$crate::print::level::$level>(
                     crate::__LOG_PREFIX,
                     args,
                 );
             }
         }
     );
+
+    ($level:ident, $fmt:literal $($arg:tt)*) => (
+        match $crate::prelude::fmt!($fmt $($arg)*) {
+            args => {
+                if const {
+                    let bytes = str::as_bytes($fmt);
+                    bytes[bytes.len() - 1] == b'\n'
+                } {
+                    // Format string already have newline, don't add extra ones.
+                    $crate::print::call_printk_continuable::<$crate::print::level::$level>(
+                        crate::__LOG_PREFIX,
+                        args,
+                    );
+                } else {
+                    $crate::print::call_printk::<$crate::print::level::$level>(
+                        crate::__LOG_PREFIX,
+                        args,
+                    );
+                }
+            }
+        }
+    );
 );
 
 /// Stub for doctests
@@ -400,7 +472,7 @@ macro_rules! pr_debug (
 ///
 /// ```
 /// # use kernel::pr_cont;
-/// pr_info!("hello");
+/// pr_info!(continuable, "hello");
 /// pr_cont!(" {}\n", "there");
 /// ```
 #[macro_export]
diff --git a/samples/rust/rust_print_main.rs b/samples/rust/rust_print_main.rs
index 682207c81fc2..b42a9e6c9d72 100644
--- a/samples/rust/rust_print_main.rs
+++ b/samples/rust/rust_print_main.rs
@@ -70,7 +70,7 @@ fn init(_module: &'static ThisModule) -> Result<Self> {
         pr_notice!("Notice message (level 5) without args\n");
         pr_info!("Info message (level 6) without args\n");
 
-        pr_info!("A line that");
+        pr_info!(continuable, "A line that");
         pr_cont!(" is continued");
         pr_cont!(" without args\n");
 
@@ -82,7 +82,7 @@ fn init(_module: &'static ThisModule) -> Result<Self> {
         pr_notice!("{} message (level {}) with args\n", "Notice", 5);
         pr_info!("{} message (level {}) with args\n", "Info", 6);
 
-        pr_info!("A {} that", "line");
+        pr_info!(continuable, "A {} that", "line");
         pr_cont!(" is {}", "continued");
         pr_cont!(" with {}\n", "args");
 

-- 
2.54.0