rust: add READ_ONCE and WRITE_ONCE

Provide Rust versions of READ_ONCE and WRITE_ONCE that have the same
semantics as the equivalent C functions.

The arm64+LTO implementation of READ_ONCE needs some more work to use
the same inline assembly as the C code.

Signed-off-by: Alice Ryhl <aliceryhl@google.com>

Author: Darksonn

rust/bindings/lib.rs

@@ -48,3 +48,6 @@ mod bindings_helper {
 }
 
 pub use bindings_raw::*;
+
+mod rwonce;
+pub use rwonce::{READ_ONCE, WRITE_ONCE};

rust/bindings/rwonce/arm64.rs

@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (C) 2024 Google LLC.
+
+#[cfg(not(CONFIG_LTO))]
+pub(super) use super::generic::__READ_ONCE;
+
+// When building with LTO, there is an increased risk of the compiler
+// converting an address dependency headed by a READ_ONCE() invocation
+// into a control dependency and consequently allowing for harmful
+// reordering by the CPU.
+//
+// Ensure that such transformations are harmless by overriding the generic
+// READ_ONCE() definition with one that provides RCpc acquire semantics
+// when building with LTO.
+#[cfg(CONFIG_LTO)]
+#[inline(always)]
+pub(super) unsafe fn __READ_ONCE<T>(ptr: *const T) -> T {
+    use core::sync::atomic::*;
+
+    macro_rules! load {
+        ($ptr:expr, $atomic:ident) => {{
+            let ptr = $ptr as *const $atomic;
+            // TODO: Use inline assembly instead.
+            let value = (&*ptr).load(Ordering::Acquire);
+            core::mem::transmute_copy::<_, T>(&value)
+        }};
+    }
+    
+    unsafe {
+        match core::mem::size_of::<T>() {
+            1 => load!(ptr, AtomicU8),
+            2 => load!(ptr, AtomicU16),
+            4 => load!(ptr, AtomicU32),
+            8 => load!(ptr, AtomicU64),
+            _ => core::ptr::read_volatile(ptr)
+        }
+    }
+}

rust/bindings/rwonce/generic.rs

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (C) 2024 Google LLC.
+
+#[inline(always)]
+pub(super) unsafe fn __READ_ONCE<T>(ptr: *const T) -> T {
+    unsafe { core::ptr::read_volatile(ptr) }
+}
+
+#[inline(always)]
+pub(super) unsafe fn __WRITE_ONCE<T>(ptr: *mut T, val: T) {
+    unsafe { core::ptr::write_volatile(ptr, val) }
+}

rust/bindings/rwonce/mod.rs

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (C) 2024 Google LLC.
+
+/// Generic implementations.
+///
+/// Use statements that mention the item by name take precedence over glob imports.
+mod generic;
+use self::generic::*;
+
+#[cfg(target_arch = "aarch64")]
+mod arm64;
+#[cfg(target_arch = "aarch64")]
+use self::arm64::__READ_ONCE;
+
+/// Read a value once.
+///
+/// # Safety
+///
+/// Same requirements as `READ_ONCE` from C.
+#[inline(always)]
+pub unsafe fn READ_ONCE<T: Copy>(ptr: *const T) -> T {
+    unsafe { __READ_ONCE(ptr) }
+}
+
+/// Write a value once.
+///
+/// # Safety
+///
+/// Same requirements as `WRITE_ONCE` from C.
+#[inline(always)]
+pub unsafe fn WRITE_ONCE<T: Copy>(ptr: *mut T, val: T) {
+    unsafe { __WRITE_ONCE(ptr, val) }
+}