[PATCH v6 5/7] rust: ww_mutex: add context-free locking functions

["Onur Özkan" <work@onurozkan.dev>]

Adds new `WwMutex` functions (`lock`, `lock_interruptible`,
`lock_slow`, `lock_slow_interruptible` and `try_lock`) that
can be used without `WwAcquireCtx`. This provides simpler API
when deadlock avoidance grouping is not needed.

Signed-off-by: Onur Özkan <work@onurozkan.dev>
---
 rust/kernel/sync/lock/ww_mutex.rs | 162 ++++++++++++++++++++++--------
 1 file changed, 122 insertions(+), 40 deletions(-)

diff --git a/rust/kernel/sync/lock/ww_mutex.rs b/rust/kernel/sync/lock/ww_mutex.rs
index d289718d2c98..b415d6deae9b 100644
--- a/rust/kernel/sync/lock/ww_mutex.rs
+++ b/rust/kernel/sync/lock/ww_mutex.rs
@@ -138,6 +138,75 @@ pub fn new_wound_wait(name: &'static CStr) -> impl PinInit<Self> {
     }
 }

+/// Locking kinds used by [`lock_common`] to unify internal FFI locking logic.
+#[derive(Copy, Clone, Debug)]
+enum LockKind {
+    /// Blocks until lock is acquired.
+    Regular,
+    /// Blocks but can be interrupted by signals.
+    Interruptible,
+    /// Used in slow path after deadlock detection.
+    Slow,
+    /// Slow path but interruptible.
+    SlowInterruptible,
+    /// Does not block, returns immediately if busy.
+    Try,
+}
+
+/// Internal helper that unifies the different locking kinds.
+fn lock_common<'a, T: ?Sized>(
+    ww_mutex: &'a WwMutex<'a, T>,
+    ctx: Option<&WwAcquireCtx<'_>>,
+    kind: LockKind,
+) -> Result<WwMutexGuard<'a, T>> {
+    let ctx_ptr = ctx.map_or(core::ptr::null_mut(), |c| c.inner.get());
+
+    match kind {
+        LockKind::Regular => {
+            // SAFETY: `WwMutex` is always pinned. If `WwAcquireCtx` is `Some`, it is pinned,
+            // if `None`, it is set to `core::ptr::null_mut()`. Both cases are safe.
+            let ret = unsafe { bindings::ww_mutex_lock(ww_mutex.mutex.get(), ctx_ptr) };
+
+            to_result(ret)?;
+        }
+        LockKind::Interruptible => {
+            // SAFETY: `WwMutex` is always pinned. If `WwAcquireCtx` is `Some`, it is pinned,
+            // if `None`, it is set to `core::ptr::null_mut()`. Both cases are safe.
+            let ret =
+                unsafe { bindings::ww_mutex_lock_interruptible(ww_mutex.mutex.get(), ctx_ptr) };
+
+            to_result(ret)?;
+        }
+        LockKind::Slow => {
+            // SAFETY: `WwMutex` is always pinned. If `WwAcquireCtx` is `Some`, it is pinned,
+            // if `None`, it is set to `core::ptr::null_mut()`. Both cases are safe.
+            unsafe { bindings::ww_mutex_lock_slow(ww_mutex.mutex.get(), ctx_ptr) };
+        }
+        LockKind::SlowInterruptible => {
+            // SAFETY: `WwMutex` is always pinned. If `WwAcquireCtx` is `Some`, it is pinned,
+            // if `None`, it is set to `core::ptr::null_mut()`. Both cases are safe.
+            let ret = unsafe {
+                bindings::ww_mutex_lock_slow_interruptible(ww_mutex.mutex.get(), ctx_ptr)
+            };
+
+            to_result(ret)?;
+        }
+        LockKind::Try => {
+            // SAFETY: `WwMutex` is always pinned. If `WwAcquireCtx` is `Some`, it is pinned,
+            // if `None`, it is set to `core::ptr::null_mut()`. Both cases are safe.
+            let ret = unsafe { bindings::ww_mutex_trylock(ww_mutex.mutex.get(), ctx_ptr) };
+
+            if ret == 0 {
+                return Err(EBUSY);
+            } else {
+                to_result(ret)?;
+            }
+        }
+    };
+
+    Ok(WwMutexGuard::new(ww_mutex))
+}
+
 /// Groups multiple mutex acquisitions together for deadlock avoidance.
 ///
 /// Must be used when acquiring multiple mutexes of the same class.
@@ -196,14 +265,9 @@ pub fn done(&self) {
         unsafe { bindings::ww_acquire_done(self.inner.get()) };
     }

-    /// Locks the given mutex.
+    /// Locks the given mutex on this acquire context ([`WwAcquireCtx`]).
     pub fn lock<'a, T>(&'a self, ww_mutex: &'a WwMutex<'a, T>) -> Result<WwMutexGuard<'a, T>> {
-        // SAFETY: The mutex is pinned and valid.
-        let ret = unsafe { bindings::ww_mutex_lock(ww_mutex.mutex.get(), self.inner.get()) };
-
-        to_result(ret)?;
-
-        Ok(WwMutexGuard::new(ww_mutex))
+        lock_common(ww_mutex, Some(self), LockKind::Regular)
     }

     /// Similar to `lock`, but can be interrupted by signals.
@@ -211,24 +275,14 @@ pub fn lock_interruptible<'a, T>(
         &'a self,
         ww_mutex: &'a WwMutex<'a, T>,
     ) -> Result<WwMutexGuard<'a, T>> {
-        // SAFETY: The mutex is pinned and valid.
-        let ret = unsafe {
-            bindings::ww_mutex_lock_interruptible(ww_mutex.mutex.get(), self.inner.get())
-        };
-
-        to_result(ret)?;
-
-        Ok(WwMutexGuard::new(ww_mutex))
+        lock_common(ww_mutex, Some(self), LockKind::Interruptible)
     }

-    /// Locks the given mutex using the slow path.
+    /// Locks the given mutex on this acquire context ([`WwAcquireCtx`]) using the slow path.
     ///
     /// This function should be used when `lock` fails (typically due to a potential deadlock).
     pub fn lock_slow<'a, T>(&'a self, ww_mutex: &'a WwMutex<'a, T>) -> Result<WwMutexGuard<'a, T>> {
-        // SAFETY: The mutex is pinned and valid, and we're in the slow path.
-        unsafe { bindings::ww_mutex_lock_slow(ww_mutex.mutex.get(), self.inner.get()) };
-
-        Ok(WwMutexGuard::new(ww_mutex))
+        lock_common(ww_mutex, Some(self), LockKind::Slow)
     }

     /// Similar to `lock_slow`, but can be interrupted by signals.
@@ -236,30 +290,14 @@ pub fn lock_slow_interruptible<'a, T>(
         &'a self,
         ww_mutex: &'a WwMutex<'a, T>,
     ) -> Result<WwMutexGuard<'a, T>> {
-        // SAFETY: The mutex is pinned and valid, and we are in the slow path.
-        let ret = unsafe {
-            bindings::ww_mutex_lock_slow_interruptible(ww_mutex.mutex.get(), self.inner.get())
-        };
-
-        to_result(ret)?;
-
-        Ok(WwMutexGuard::new(ww_mutex))
+        lock_common(ww_mutex, Some(self), LockKind::SlowInterruptible)
     }

-    /// Tries to lock the mutex without blocking.
+    /// Tries to lock the mutex on this acquire context ([`WwAcquireCtx`]) without blocking.
     ///
     /// Unlike `lock`, no deadlock handling is performed.
     pub fn try_lock<'a, T>(&'a self, ww_mutex: &'a WwMutex<'a, T>) -> Result<WwMutexGuard<'a, T>> {
-        // SAFETY: The mutex is pinned and valid.
-        let ret = unsafe { bindings::ww_mutex_trylock(ww_mutex.mutex.get(), self.inner.get()) };
-
-        if ret == 0 {
-            return Err(EBUSY);
-        } else {
-            to_result(ret)?;
-        }
-
-        Ok(WwMutexGuard::new(ww_mutex))
+        lock_common(ww_mutex, Some(self), LockKind::Try)
     }
 }

@@ -355,7 +393,7 @@ pub fn new(t: T, ww_class: &'ww_class WwClass) -> impl PinInit<Self> {
     }
 }

-impl<T: ?Sized> WwMutex<'_, T> {
+impl<'ww_class, T: ?Sized> WwMutex<'ww_class, T> {
     /// Returns a raw pointer to the inner mutex.
     fn as_ptr(&self) -> *mut bindings::ww_mutex {
         self.mutex.get()
@@ -370,6 +408,35 @@ fn is_locked(&self) -> bool {
         // SAFETY: The mutex is pinned and valid.
         unsafe { bindings::ww_mutex_is_locked(self.mutex.get()) }
     }
+
+    /// Locks the given mutex without acquire context ([`WwAcquireCtx`]).
+    pub fn lock<'a>(&'a self) -> Result<WwMutexGuard<'a, T>> {
+        lock_common(self, None, LockKind::Regular)
+    }
+
+    /// Similar to `lock`, but can be interrupted by signals.
+    pub fn lock_interruptible<'a>(&'a self) -> Result<WwMutexGuard<'a, T>> {
+        lock_common(self, None, LockKind::Interruptible)
+    }
+
+    /// Locks the given mutex without acquire context ([`WwAcquireCtx`]) using the slow path.
+    ///
+    /// This function should be used when `lock` fails (typically due to a potential deadlock).
+    pub fn lock_slow<'a>(&'a self) -> Result<WwMutexGuard<'a, T>> {
+        lock_common(self, None, LockKind::Slow)
+    }
+
+    /// Similar to `lock_slow`, but can be interrupted by signals.
+    pub fn lock_slow_interruptible<'a>(&'a self) -> Result<WwMutexGuard<'a, T>> {
+        lock_common(self, None, LockKind::SlowInterruptible)
+    }
+
+    /// Tries to lock the mutex without acquire context ([`WwAcquireCtx`]) and without blocking.
+    ///
+    /// Unlike `lock`, no deadlock handling is performed.
+    pub fn try_lock<'a>(&'a self) -> Result<WwMutexGuard<'a, T>> {
+        lock_common(self, None, LockKind::Try)
+    }
 }

 /// A guard that provides exclusive access to the data protected
@@ -547,4 +614,19 @@ fn test_with_global_classes() -> Result {

         Ok(())
     }
+
+    #[test]
+    fn test_mutex_without_ctx() -> Result {
+        let mutex = Arc::pin_init(WwMutex::new(100, &TEST_WOUND_WAIT_CLASS), GFP_KERNEL)?;
+        let guard = mutex.lock()?;
+
+        assert_eq!(*guard, 100);
+        assert!(mutex.is_locked());
+
+        drop(guard);
+
+        assert!(!mutex.is_locked());
+
+        Ok(())
+    }
 }
--
2.50.0