[PATCH 3/3] rust: sync: Add IrqSpinLock

[Lyude Paul]

A variant of SpinLock that is expected to be used in noirq contexts, and
thus requires that the user provide an kernel::irq::IrqDisabled to prove
they are in such a context upon lock acquisition. This is the rust
equivalent of spin_lock_irqsave()/spin_lock_irqrestore().

Signed-off-by: Lyude Paul <lyude@xxxxxxxxxx>
---
 rust/kernel/sync.rs               |   2 +-
 rust/kernel/sync/lock/spinlock.rs | 101 ++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 1 deletion(-)

diff --git a/rust/kernel/sync.rs b/rust/kernel/sync.rs
index 14a79ebbb42d5..d3e2e9cb0e120 100644
--- a/rust/kernel/sync.rs
+++ b/rust/kernel/sync.rs
@@ -15,7 +15,7 @@
 pub use arc::{Arc, ArcBorrow, UniqueArc};
 pub use condvar::{new_condvar, CondVar, CondVarTimeoutResult};
 pub use lock::mutex::{new_mutex, Mutex};
-pub use lock::spinlock::{new_spinlock, SpinLock};
+pub use lock::spinlock::{new_irq_spinlock, new_spinlock, IrqSpinLock, SpinLock};
 pub use lock::LockContainer;
 pub use locked_by::LockedBy;
 
diff --git a/rust/kernel/sync/lock/spinlock.rs b/rust/kernel/sync/lock/spinlock.rs
index ea5c5bc1ce12e..f63c75c19344f 100644
--- a/rust/kernel/sync/lock/spinlock.rs
+++ b/rust/kernel/sync/lock/spinlock.rs
@@ -3,6 +3,8 @@
 //! A kernel spinlock.
 //!
 //! This module allows Rust code to use the kernel's `spinlock_t`.
+use core::marker::*;
+use kernel::{irq::*, prelude::*};
 
 /// Creates a [`SpinLock`] initialiser with the given name and a newly-created lock class.
 ///
@@ -115,3 +117,102 @@ unsafe fn unlock(ptr: *mut Self::State, _guard_state: &Self::GuardState) {
         unsafe { bindings::spin_unlock(ptr) }
     }
 }
+
+/// Creates a [`IrqSpinLock`] initialiser with the given name and a newly-created lock class.
+///
+/// It uses the name if one is given, otherwise it generates one based on the file name and line
+/// number.
+#[macro_export]
+macro_rules! new_irq_spinlock {
+    ($inner:expr $(, $name:literal)? $(,)?) => {
+        $crate::sync::IrqSpinLock::new(
+            $inner, $crate::optional_name!($($name)?), $crate::static_lock_class!())
+    };
+}
+pub use new_irq_spinlock;
+
+/// A spinlock that may be acquired when interrupts are disabled.
+///
+/// A version of [`SpinLock`] that can only be used in contexts where interrupts for the local CPU
+/// are disabled. It requires that the user acquiring the lock provide proof that interrupts are
+/// disabled through [`IrqDisabled`]
+///
+/// For more info, see [`SpinLock`].
+///
+/// # Examples
+///
+/// The following example shows how to declare, allocate initialise and access a struct (`Example`)
+/// that contains an inner struct (`Inner`) that is protected by a spinlock.
+///
+/// ```
+/// use kernel::{
+///     sync::{new_irq_spinlock, IrqSpinLock},
+///     irq::{with_irqs_disabled, IrqDisabled}
+/// };
+///
+/// struct Inner {
+///     a: u32,
+///     b: u32,
+/// }
+///
+/// #[pin_data]
+/// struct Example {
+///     c: u32,
+///     #[pin]
+///     d: IrqSpinLock<Inner>,
+/// }
+///
+/// impl Example {
+///     fn new() -> impl PinInit<Self> {
+///         pin_init!(Self {
+///             c: 10,
+///             d <- new_irq_spinlock!(Inner { a: 20, b: 30 }),
+///         })
+///     }
+/// }
+///
+/// // Accessing an `Example` from a function that can only be called in no-irq contexts
+/// fn noirq_work(e: &Example, irq: &IrqDisabled<'_>) {
+///     assert_eq!(e.c, 10);
+///     assert_eq!(e.d.lock(&irq).a, 20);
+/// }
+///
+/// // Allocate a boxed `Example`
+/// let e = Box::pin_init(Example::new(), GFP_KERNEL)?;
+///
+/// // Accessing an `Example` from a context where IRQs may not be disabled already.
+/// let b = with_irqs_disabled(|irq| {
+///     noirq_work(&e, &irq);
+///     e.d.lock(&irq).b
+/// );
+/// assert_eq!(b, 30);
+/// # Ok::<(), Error>(())
+/// ```
+#[pin_data]
+pub struct IrqSpinLock<T> {
+    #[pin]
+    inner: SpinLock<T>,
+    #[pin]
+    _p: PhantomPinned,
+}
+
+impl<T> IrqSpinLock<T> {
+    /// Constructs a new IRQ spinlock initialiser
+    pub fn new(t: T, name: &'static CStr, key: &'static super::LockClassKey) -> impl PinInit<Self> {
+        pin_init!(Self {
+            inner <- SpinLock::new(t, name, key),
+            _p: PhantomPinned,
+        })
+    }
+
+    /// Acquires the lock and gives the caller access to the data protected by it
+    pub fn lock<'a>(&'a self, _irq: &'a IrqDisabled<'a>) -> super::Guard<'a, T, SpinLockBackend> {
+        self.inner.lock()
+    }
+}
+
+impl<T> super::LockContainer<T, SpinLockBackend> for IrqSpinLock<T> {
+    unsafe fn get_lock_ref(&self) -> &super::Lock<T, SpinLockBackend> {
+        &self.inner
+    }
+}
-- 
2.45.2