[POC 6/6] rust: sync: lock: Add `Backend::BackendInContext`

[Boqun Feng <boqun.feng@gmail.com>]

`SpinLock`'s backend can be used for `SpinLockIrq`, if the interrupts
are disabled. And it actually provides performance gains since
interrupts are not needed to be disabled anymore. So add
`Backend::BackendInContext` to describe the case where one backend can
be used for another. Use the it to implement the `lock_with()` so that
`SpinLockIrq` can avoid disabling interrupts by using `SpinLock`'s
backend.

Signed-off-by: Boqun Feng <boqun.feng@gmail.com>
---
 rust/kernel/sync/lock.rs          | 25 +++++++++++++++++++++++--
 rust/kernel/sync/lock/mutex.rs    |  1 +
 rust/kernel/sync/lock/spinlock.rs |  9 +++++++++
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/rust/kernel/sync/lock.rs b/rust/kernel/sync/lock.rs
index 49b53433201c..4e3316feb497 100644
--- a/rust/kernel/sync/lock.rs
+++ b/rust/kernel/sync/lock.rs
@@ -24,10 +24,14 @@
 ///   is owned, that is, between calls to [`lock`] and [`unlock`].
 /// - Implementers must also ensure that [`relock`] uses the same locking method as the original
 ///   lock operation.
+/// - Implementers must ensure if [`BackendInContext`] is a [`Backend`], it's safe to acquire lock
+///   under the [`Context`], the [`State`] of two backends must be the same.
 ///
 /// [`lock`]: Backend::lock
 /// [`unlock`]: Backend::unlock
 /// [`relock`]: Backend::relock
+/// [`BackendInContext`]: Backend::BackendInContext
+/// [`Context`]: Backend::Context
 pub unsafe trait Backend {
     /// The state required by the lock.
     type State;
@@ -41,6 +45,9 @@ pub unsafe trait Backend {
     /// The context which can be provided to acquire the lock with a different backend.
     type Context<'a>;
 
+    /// The alternative backend we can use if a [`Context`] is provided.
+    type BackendInContext: Sized;
+
     /// Initialises the lock.
     ///
     /// # Safety
@@ -125,8 +132,22 @@ pub fn new(t: T, name: &'static CStr, key: &'static LockClassKey) -> impl PinIni
 impl<T: ?Sized, B: Backend> Lock<T, B> {
     /// Acquires the lock with the given context and gives the caller access to the data protected
     /// by it.
-    pub fn lock_with<'a>(&'a self, _context: B::Context<'a>) -> Guard<'a, T, B> {
-        todo!()
+    pub fn lock_with<'a>(&'a self, _context: B::Context<'a>) -> Guard<'a, T, B::BackendInContext>
+    where
+        B::BackendInContext: Backend,
+    {
+        // SAFETY: Per the safety guarantee of `Backend`, if `B::BackendIncontext` and `B` should
+        // have the same state, therefore the layout of the lock is the same so it's safe the
+        // convert one to another.
+        let lock = unsafe { &*(self as *const _ as *const Lock<T, B::BackendInContext>) };
+        // SAFETY: The constructor of the type calls `init`, so the existence of the object proves
+        // that `init` was called. Plus the safety guarantee of `Backend` guarantees that `B::state`
+        // is the same as `B::BackendInContext::state`, also it's safe to call another backend
+        // because there is `B::Context<'a>`.
+        let state = unsafe { B::BackendInContext::lock(lock.state.get()) };
+
+        // SAFETY: The lock was just acquired.
+        unsafe { Guard::new(lock, state) }
     }
 
     /// Acquires the lock and gives the caller access to the data protected by it.
diff --git a/rust/kernel/sync/lock/mutex.rs b/rust/kernel/sync/lock/mutex.rs
index 7c2c23994493..ddb7d06676f7 100644
--- a/rust/kernel/sync/lock/mutex.rs
+++ b/rust/kernel/sync/lock/mutex.rs
@@ -94,6 +94,7 @@ unsafe impl super::Backend for MutexBackend {
     type State = bindings::mutex;
     type GuardState = ();
     type Context<'a> = ();
+    type BackendInContext = ();
 
     unsafe fn init(
         ptr: *mut Self::State,
diff --git a/rust/kernel/sync/lock/spinlock.rs b/rust/kernel/sync/lock/spinlock.rs
index 8f9e1b27e474..3bec25b65c55 100644
--- a/rust/kernel/sync/lock/spinlock.rs
+++ b/rust/kernel/sync/lock/spinlock.rs
@@ -94,6 +94,7 @@ unsafe impl super::Backend for SpinLockBackend {
     type State = bindings::spinlock_t;
     type GuardState = ();
     type Context<'a> = ();
+    type BackendInContext = ();
 
     unsafe fn init(
         ptr: *mut Self::State,
@@ -146,6 +147,7 @@ macro_rules! new_spinlock_irq {
 ///
 /// ```
 /// use kernel::sync::{new_spinlock_irq, SpinLockIrq};
+/// use kernel::interrupt::InterruptDisabled;
 ///
 /// struct Inner {
 ///     a: u32,
@@ -168,6 +170,12 @@ macro_rules! new_spinlock_irq {
 ///     }
 /// }
 ///
+/// // Accessing an `Example` from a function that can only be called in no-irq contexts
+/// fn noirq_work(e: &Example, irq: &InterruptDisabled) {
+///     assert_eq!(e.c, 10);
+///     assert_eq!(e.d.lock_with(irq).a, 20);
+/// }
+///
 /// // Allocate a boxed `Example`
 /// let e = Box::pin_init(Example::new(), GFP_KERNEL)?;
 ///
@@ -186,6 +194,7 @@ unsafe impl super::Backend for SpinLockIrqBackend {
     type State = bindings::spinlock_t;
     type GuardState = ();
     type Context<'a> = &'a InterruptDisabled;
+    type BackendInContext = SpinLockBackend;
 
     unsafe fn init(
         ptr: *mut Self::State,
-- 
2.45.2