[PATCH 2/5] rust: sync: add READ_ONCE and WRITE_ONCE

[Alice Ryhl <aliceryhl@google.com>]

There are currently a few places in the kernel where we use volatile
reads when we really should be using `READ_ONCE`. To make it possible to
replace these with proper `READ_ONCE` calls, introduce a Rust version of
`READ_ONCE`.

I've written the code to use Rust's volatile ops directly when possible.
This results in a small amount of code duplication, but I think it makes
sense for READ_ONCE and WRITE_ONCE to be implemented in pure Rust when
possible. Otherwise they would unconditionally be a function call unless
you have a system where you can perform cross-language inlining.

I considered these functions in the bindings crate instead of kernel
crate. I actually think it would make a lot of sense. But it implies
some annoying complications on old compilers since the #![feature()]
invocations in kernel/lib.rs do not apply in the bindings crate.

For now, we do not support using READ_ONCE on compound types even if
they have the right size. This can be added later.

This fails checkpatch due to a misordered MAINTAINERS entry, but this is
a pre-existing problem.

Signed-off-by: Alice Ryhl <aliceryhl@google.com>
---
 MAINTAINERS                |   2 +
 rust/helpers/helpers.c     |   1 +
 rust/helpers/rwonce.c      |  34 ++++++++
 rust/kernel/sync.rs        |   2 +
 rust/kernel/sync/rwonce.rs | 188 +++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 227 insertions(+)

diff --git a/MAINTAINERS b/MAINTAINERS
index 12f49de7fe036c2439c00f9f4c67b2219d72a4c3..1d0cae158fe2cc7d99b6a64c11176b635e2d14e4 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -4117,9 +4117,11 @@ F:	arch/*/include/asm/atomic*.h
 F:	include/*/atomic*.h
 F:	include/linux/refcount.h
 F:	scripts/atomic/
+F:	rust/helpers/rwonce.c
 F:	rust/kernel/sync/atomic.rs
 F:	rust/kernel/sync/atomic/
 F:	rust/kernel/sync/refcount.rs
+F:	rust/kernel/sync/rwonce.rs
 
 ATTO EXPRESSSAS SAS/SATA RAID SCSI DRIVER
 M:	Bradley Grove <linuxdrivers@attotech.com>
diff --git a/rust/helpers/helpers.c b/rust/helpers/helpers.c
index 79c72762ad9c4b473971e6210c9577860d2e2b08..28b79ca7844fb744e5ad128238824921c055ec82 100644
--- a/rust/helpers/helpers.c
+++ b/rust/helpers/helpers.c
@@ -48,6 +48,7 @@
 #include "rcu.c"
 #include "refcount.c"
 #include "regulator.c"
+#include "rwonce.c"
 #include "scatterlist.c"
 #include "security.c"
 #include "signal.c"
diff --git a/rust/helpers/rwonce.c b/rust/helpers/rwonce.c
new file mode 100644
index 0000000000000000000000000000000000000000..55c621678cd632e728cb925b6a4a2e34e2fc4884
--- /dev/null
+++ b/rust/helpers/rwonce.c
@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: GPL-2.0
+
+/*
+ * Copyright (C) 2025 Google LLC.
+ */
+
+#ifdef CONFIG_ARCH_USE_CUSTOM_READ_ONCE
+
+__rust_helper u8 rust_helper_read_once_1(const u8 *ptr)
+{
+	return READ_ONCE(*ptr);
+}
+
+__rust_helper u16 rust_helper_read_once_2(const u16 *ptr)
+{
+	return READ_ONCE(*ptr);
+}
+
+__rust_helper u32 rust_helper_read_once_4(const u32 *ptr)
+{
+	return READ_ONCE(*ptr);
+}
+
+__rust_helper u64 rust_helper_read_once_8(const u64 *ptr)
+{
+	return READ_ONCE(*ptr);
+}
+
+__rust_helper void *rust_helper_read_once_ptr(void * const *ptr)
+{
+	return READ_ONCE(*ptr);
+}
+
+#endif
diff --git a/rust/kernel/sync.rs b/rust/kernel/sync.rs
index 5df87e2bd212e192b8a67644bd99f05b9d4afd75..a5bf7bdc3fa8a044786eafae39fe8844aeeef057 100644
--- a/rust/kernel/sync.rs
+++ b/rust/kernel/sync.rs
@@ -20,6 +20,7 @@
 pub mod poll;
 pub mod rcu;
 mod refcount;
+pub mod rwonce;
 mod set_once;
 
 pub use arc::{Arc, ArcBorrow, UniqueArc};
@@ -30,6 +31,7 @@
 pub use lock::spinlock::{new_spinlock, SpinLock, SpinLockGuard};
 pub use locked_by::LockedBy;
 pub use refcount::Refcount;
+pub use rwonce::{READ_ONCE, WRITE_ONCE};
 pub use set_once::SetOnce;
 
 /// Represents a lockdep class. It's a wrapper around C's `lock_class_key`.
diff --git a/rust/kernel/sync/rwonce.rs b/rust/kernel/sync/rwonce.rs
new file mode 100644
index 0000000000000000000000000000000000000000..a1660e43c9ef94011812d1816713cf031a73de1d
--- /dev/null
+++ b/rust/kernel/sync/rwonce.rs
@@ -0,0 +1,188 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (C) 2025 Google LLC.
+
+//! Rust version of the raw `READ_ONCE`/`WRITE_ONCE` functions.
+//!
+//! C header: [`include/asm-generic/rwonce.h`](srctree/include/asm-generic/rwonce.h)
+
+/// Read the pointer once.
+///
+/// # Safety
+///
+/// It must be safe to `READ_ONCE` the `ptr` with this type.
+#[inline(always)]
+#[must_use]
+#[track_caller]
+#[expect(non_snake_case)]
+pub unsafe fn READ_ONCE<T: RwOnceType>(ptr: *const T) -> T {
+    // SAFETY: It's safe to read `ptr` once with this type.
+    unsafe { T::read_once(ptr) }
+}
+
+/// Write the pointer once.
+///
+/// # Safety
+///
+/// It must be safe to `WRITE_ONCE` the `ptr` with this type.
+#[inline(always)]
+#[track_caller]
+#[expect(non_snake_case)]
+pub unsafe fn WRITE_ONCE<T: RwOnceType>(ptr: *mut T, val: T) {
+    // SAFETY: It's safe to write `ptr` once with this type.
+    unsafe { T::write_once(ptr, val) };
+}
+
+/// This module contains the generic implementations.
+#[expect(clippy::undocumented_unsafe_blocks)]
+#[expect(clippy::missing_safety_doc)]
+mod rwonce_generic_impl {
+    use core::ffi::c_void;
+    #[allow(unused_imports)]
+    use core::ptr::{read_volatile, write_volatile};
+
+    #[inline(always)]
+    #[track_caller]
+    #[cfg(not(CONFIG_ARCH_USE_CUSTOM_READ_ONCE))]
+    pub(super) unsafe fn read_once_1(ptr: *const u8) -> u8 {
+        unsafe { read_volatile::<u8>(ptr) }
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    #[cfg(not(CONFIG_ARCH_USE_CUSTOM_READ_ONCE))]
+    pub(super) unsafe fn read_once_2(ptr: *const u16) -> u16 {
+        unsafe { read_volatile::<u16>(ptr) }
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    #[cfg(not(CONFIG_ARCH_USE_CUSTOM_READ_ONCE))]
+    pub(super) unsafe fn read_once_4(ptr: *const u32) -> u32 {
+        unsafe { read_volatile::<u32>(ptr) }
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    #[cfg(not(CONFIG_ARCH_USE_CUSTOM_READ_ONCE))]
+    pub(super) unsafe fn read_once_8(ptr: *const u64) -> u64 {
+        unsafe { read_volatile::<u64>(ptr) }
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    #[cfg(not(CONFIG_ARCH_USE_CUSTOM_READ_ONCE))]
+    pub(super) unsafe fn read_once_ptr(ptr: *const *mut c_void) -> *mut c_void {
+        unsafe { read_volatile::<*mut c_void>(ptr) }
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    pub(super) unsafe fn write_once_1(ptr: *mut u8, val: u8) {
+        unsafe { write_volatile::<u8>(ptr, val) }
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    pub(super) unsafe fn write_once_2(ptr: *mut u16, val: u16) {
+        unsafe { write_volatile::<u16>(ptr, val) }
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    pub(super) unsafe fn write_once_4(ptr: *mut u32, val: u32) {
+        unsafe { write_volatile::<u32>(ptr, val) }
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    pub(super) unsafe fn write_once_8(ptr: *mut u64, val: u64) {
+        unsafe { write_volatile::<u64>(ptr, val) }
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    pub(super) unsafe fn write_once_ptr(ptr: *mut *mut c_void, val: *mut c_void) {
+        unsafe { write_volatile::<*mut c_void>(ptr, val) }
+    }
+}
+use rwonce_generic_impl::*;
+
+#[cfg(CONFIG_ARCH_USE_CUSTOM_READ_ONCE)]
+use bindings::{read_once_1, read_once_2, read_once_4, read_once_8, read_once_ptr};
+
+/// Rust trait for types that may be used with `READ_ONCE`/`WRITE_ONCE`.
+///
+/// This serves a similar purpose to the `compiletime_assert_rwonce_type` macro in the C header.
+pub trait RwOnceType {
+    /// The `READ_ONCE` for this type.
+    ///
+    /// # Safety
+    ///
+    /// It must be safe to `READ_ONCE` the `ptr` with this type.
+    unsafe fn read_once(ptr: *const Self) -> Self;
+
+    /// The `WRITE_ONCE` for this type.
+    ///
+    /// # Safety
+    ///
+    /// It must be safe to `WRITE_ONCE` the `ptr` with this type.
+    unsafe fn write_once(ptr: *mut Self, val: Self);
+}
+
+macro_rules! impl_rw_once_type {
+    ($($t:ty, $read:ident, $write:ident $(, <$u:ident>)?;)*) => {$(
+        #[allow(unknown_lints, reason = "unnecessary_transmutes is unknown prior to MSRV 1.88.0")]
+        #[allow(unnecessary_transmutes)]
+        #[allow(clippy::missing_transmute_annotations)]
+        #[allow(clippy::useless_transmute)]
+        impl$(<$u>)? RwOnceType for $t {
+            #[inline(always)]
+            #[track_caller]
+            unsafe fn read_once(ptr: *const Self) -> Self {
+                // SAFETY: The caller ensures we can `READ_ONCE`.
+                //
+                // Note that `transmute` fails to compile if the two types are of different sizes.
+                unsafe { core::mem::transmute($read(ptr.cast())) }
+            }
+
+            #[inline(always)]
+            #[track_caller]
+            unsafe fn write_once(ptr: *mut Self, val: Self) {
+                // SAFETY: The caller ensures we can `WRITE_ONCE`.
+                unsafe { $write(ptr.cast(), core::mem::transmute(val)) };
+            }
+        }
+    )*}
+}
+
+// These macros determine which types may be used with rwonce, and which helper function should be
+// used if so.
+//
+// Note that `core::mem::transmute` fails the build if the source and target type have different
+// sizes, so picking the wrong helper should lead to a build error.
+
+impl_rw_once_type! {
+    u8,   read_once_1, write_once_1;
+    i8,   read_once_1, write_once_1;
+    u16,  read_once_2, write_once_2;
+    i16,  read_once_2, write_once_2;
+    u32,  read_once_4, write_once_4;
+    i32,  read_once_4, write_once_4;
+    u64,  read_once_8, write_once_8;
+    i64,  read_once_8, write_once_8;
+    *mut T, read_once_ptr, write_once_ptr, <T>;
+    *const T, read_once_ptr, write_once_ptr, <T>;
+}
+
+#[cfg(target_pointer_width = "32")]
+impl_rw_once_type! {
+    usize, read_once_4, write_once_4;
+    isize, read_once_4, write_once_4;
+}
+
+#[cfg(target_pointer_width = "64")]
+impl_rw_once_type! {
+    usize, read_once_8, write_once_8;
+    isize, read_once_8, write_once_8;
+}

-- 
2.52.0.351.gbe84eed79e-goog