rust: device: Add support for locking the device

Signed-off-by: Janne Grunau <j@jannau.net>

Author: jannau

rust/helpers/device.c

@@ -25,3 +25,13 @@ void rust_helper_dev_set_drvdata(struct device *dev, void *data)
 {
 	dev_set_drvdata(dev, data);
 }
+
+void rust_helper_device_lock(struct device *dev)
+{
+	device_lock(dev);
+}
+
+void rust_helper_device_unlock(struct device *dev)
+{
+	device_unlock(dev);
+}

rust/kernel/device.rs

@@ -6,7 +6,7 @@
 
 use crate::{
     bindings,
-    types::{ARef, ForeignOwnable, Opaque},
+    types::{ARef, ForeignOwnable, NotThreadSafe, Opaque},
 };
 use core::{fmt, marker::PhantomData, ptr};
 
@@ -398,6 +398,35 @@ impl<Ctx: DeviceContext> Device<Ctx> {
         // defined as a `#[repr(transparent)]` wrapper around `fwnode_handle`.
         Some(unsafe { &*fwnode_handle.cast() })
     }
+
+    /// Locks the [`Device`] for exclusive access.
+    pub fn lock(&self) -> Guard<'_, Ctx> {
+        // SAFETY: `self` is always valid by the type invariant.
+        unsafe { bindings::device_lock(self.as_raw()) };
+
+        Guard {
+            dev: self,
+            _not_send: NotThreadSafe,
+        }
+    }
+}
+
+/// A lock guard.
+///
+/// The lock is unlocked when the guard goes out of scope.
+#[must_use = "the lock unlocks immediately when the guard is unused"]
+pub struct Guard<'a, Ctx: DeviceContext = Normal> {
+    dev: &'a Device<Ctx>,
+    _not_send: NotThreadSafe,
+}
+
+impl<Ctx: DeviceContext> Drop for Guard<'_, Ctx> {
+    fn drop(&mut self) {
+        // SAFETY:
+        // - `self.xa.xa` is always valid by the type invariant.
+        // - The caller holds the lock, so it is safe to unlock it.
+        unsafe { bindings::device_unlock(self.dev.as_raw()) };
+    }
 }
 
 // SAFETY: `Device` is a transparent wrapper of a type that doesn't depend on `Device`'s generic