[PATCH 2/8] rust: device: add drvdata accessors

[Danilo Krummrich <dakr@kernel.org>]

Implement generic accessors for the private data of a driver bound to a
device.

Those accessors should be used by bus abstractions from their
corresponding core callbacks, such as probe(), remove(), etc.

Implementing them for device::Internal guarantees that driver's can't
interfere with the logic implemented by the bus abstraction.

Signed-off-by: Danilo Krummrich <dakr@kernel.org>
---
 rust/helpers/device.c | 10 ++++++++++
 rust/kernel/device.rs | 43 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/rust/helpers/device.c b/rust/helpers/device.c
index b2135c6686b0..9bf252649c75 100644
--- a/rust/helpers/device.c
+++ b/rust/helpers/device.c
@@ -8,3 +8,13 @@ int rust_helper_devm_add_action(struct device *dev,
 {
 	return devm_add_action(dev, action, data);
 }
+
+void *rust_helper_dev_get_drvdata(const struct device *dev)
+{
+	return dev_get_drvdata(dev);
+}
+
+void rust_helper_dev_set_drvdata(struct device *dev, void *data)
+{
+	dev_set_drvdata(dev, data);
+}
diff --git a/rust/kernel/device.rs b/rust/kernel/device.rs
index e9094d8322d5..146eba147d2f 100644
--- a/rust/kernel/device.rs
+++ b/rust/kernel/device.rs
@@ -6,7 +6,7 @@
 
 use crate::{
     bindings,
-    types::{ARef, Opaque},
+    types::{ARef, ForeignOwnable, Opaque},
 };
 use core::{fmt, marker::PhantomData, ptr};
 
@@ -62,6 +62,47 @@ pub unsafe fn get_device(ptr: *mut bindings::device) -> ARef<Self> {
     }
 }
 
+impl Device<Internal> {
+    /// Store a pointer to the bound driver's private data.
+    pub fn set_drvdata(&self, data: impl ForeignOwnable) {
+        // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
+        unsafe { bindings::dev_set_drvdata(self.as_raw(), data.into_foreign().cast()) }
+    }
+
+    /// Take ownership of the private data stored in this [`Device`].
+    ///
+    /// # Safety
+    ///
+    /// - Must only be called once after a preceding call to [`Device::set_drvdata`].
+    /// - The type `T` must match the type of the `ForeignOwnable` previously stored by
+    ///   [`Device::set_drvdata`].
+    pub unsafe fn drvdata_obtain<T: ForeignOwnable>(&self) -> T {
+        // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
+        let ptr = unsafe { bindings::dev_get_drvdata(self.as_raw()) };
+
+        // SAFETY: By the safety requirements of this function, `ptr` comes from a previous call to
+        // `into_foreign()`.
+        unsafe { T::from_foreign(ptr.cast()) }
+    }
+
+    /// Borrow the driver's private data bound to this [`Device`].
+    ///
+    /// # Safety
+    ///
+    /// - Must only be called after a preceding call to [`Device::set_drvdata`] and before
+    ///   [`Device::drvdata_obtain`].
+    /// - The type `T` must match the type of the `ForeignOwnable` previously stored by
+    ///   [`Device::set_drvdata`].
+    pub unsafe fn drvdata_borrow<T: ForeignOwnable>(&self) -> T::Borrowed<'_> {
+        // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
+        let ptr = unsafe { bindings::dev_get_drvdata(self.as_raw()) };
+
+        // SAFETY: By the safety requirements of this function, `ptr` comes from a previous call to
+        // `into_foreign()`.
+        unsafe { T::borrow(ptr.cast()) }
+    }
+}
+
 impl<Ctx: DeviceContext> Device<Ctx> {
     /// Obtain the raw `struct device *`.
     pub(crate) fn as_raw(&self) -> *mut bindings::device {
-- 
2.49.0