[PATCH RFC 1/2] rust: sync: introduce a way to create lock class from caller

[Gary Guo]

Rust provides a `#[track_caller]` mechanism that can be used to identify
callers from callees. This also works across multiple level of calls.
As lockdep only needs the address for static key classes, the address of
these `&'static Location` can be used as unique lock class keys.
This allows a more implicit way of creating lock classes without having to
use macros.

Create `LockClassKey::from_caller`, and use it for mutex and spinlocks.

Signed-off-by: Gary Guo <gary@xxxxxxxxxxx>
---
 rust/kernel/sync.rs               | 37 +++++++++++++++++++++++++++++++++----
 rust/kernel/sync/lock.rs          | 13 +++++++++++--
 rust/kernel/sync/lock/mutex.rs    |  3 +--
 rust/kernel/sync/lock/spinlock.rs |  3 +--
 4 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/rust/kernel/sync.rs b/rust/kernel/sync.rs
index 993dbf2caa0e..cf76fb37c460 100644
--- a/rust/kernel/sync.rs
+++ b/rust/kernel/sync.rs
@@ -5,8 +5,12 @@
 //! This module contains the kernel APIs related to synchronisation that have been ported or
 //! wrapped for usage by Rust code in the kernel.
 
-use crate::prelude::*;
-use crate::types::Opaque;
+use core::panic::Location;
+
+use crate::{
+    prelude::*,
+    types::Opaque, //
+};
 use pin_init;
 
 mod arc;
@@ -66,6 +70,31 @@ impl LockClassKey {
         }
     }
 
+    /// Obtain a statically allocated lock class key identified by the caller's location.
+    ///
+    /// Different caller locations will be guaranteed to have different lock class keys; however for
+    /// the same caller location, this may return different keys, e.g. if the caller is instantiated
+    /// in different object files.
+    #[inline]
+    #[track_caller]
+    pub const fn from_caller() -> Pin<&'static Self> {
+        static_assert!(size_of::<Location<'_>>() >= size_of::<LockClassKey>());
+
+        #[cfg(CONFIG_LOCKDEP)]
+        {
+            let caller = Location::caller();
+            // SAFETY: For static objects, lockdep does not touch the underlying memory, and only
+            // the address matters. `Location::caller()` is in rodata so it meets the requirement.
+            // The size check above makes sure that it does not overlap with other lock class keys.
+            unsafe { Pin::new_unchecked(&*core::ptr::from_ref(caller).cast::<Self>()) }
+        }
+
+        // A separate version if lockdep is disabled to avoid unnecessary `Location` being
+        // generated.
+        #[cfg(not(CONFIG_LOCKDEP))]
+        static_lock_class!()
+    }
+
     /// Initializes a dynamically allocated lock class key.
     ///
     /// In the common case of using a statically allocated lock class key, the
@@ -83,7 +112,7 @@ impl LockClassKey {
     /// let key_ptr = key.into_foreign();
     ///
     /// {
-    ///     stack_pin_init!(let num: SpinLock<u32> = SpinLock::new(
+    ///     stack_pin_init!(let num: SpinLock<u32> = SpinLock::new_with_lock_class(
     ///         0,
     ///         c"my_spinlock",
     ///         // SAFETY: `key_ptr` is returned by the above `into_foreign()`, whose
@@ -129,7 +158,7 @@ fn drop(self: Pin<&mut Self>) {
 /// use kernel::sync::{static_lock_class, Arc, SpinLock};
 ///
 /// fn new_locked_int() -> Result<Arc<SpinLock<u32>>> {
-///     Arc::pin_init(SpinLock::new(
+///     Arc::pin_init(SpinLock::new_with_lock_class(
 ///         42,
 ///         c"new_locked_int",
 ///         static_lock_class!(),
diff --git a/rust/kernel/sync/lock.rs b/rust/kernel/sync/lock.rs
index 10b6b5e9b024..447fec291cdf 100644
--- a/rust/kernel/sync/lock.rs
+++ b/rust/kernel/sync/lock.rs
@@ -127,8 +127,17 @@ unsafe impl<T: ?Sized + Send, B: Backend> Send for Lock<T, B> {}
 unsafe impl<T: ?Sized + Send, B: Backend> Sync for Lock<T, B> {}
 
 impl<T, B: Backend> Lock<T, B> {
-    /// Constructs a new lock initialiser.
-    pub fn new(
+    /// Constructs a new lock initialiser with a custom name.
+    #[inline]
+    #[track_caller]
+    pub fn new_with_name(t: impl PinInit<T>, name: &'static CStr) -> impl PinInit<Self> {
+        let key = LockClassKey::from_caller();
+        Self::new_with_lock_class(t, name, key)
+    }
+
+    /// Constructs a new lock initialiser with a custom name and lock class key.
+    #[inline]
+    pub fn new_with_lock_class(
         t: impl PinInit<T>,
         name: &'static CStr,
         key: Pin<&'static LockClassKey>,
diff --git a/rust/kernel/sync/lock/mutex.rs b/rust/kernel/sync/lock/mutex.rs
index cda0203efefb..3675ce244e08 100644
--- a/rust/kernel/sync/lock/mutex.rs
+++ b/rust/kernel/sync/lock/mutex.rs
@@ -11,8 +11,7 @@
 #[macro_export]
 macro_rules! new_mutex {
     ($inner:expr $(, $name:literal)? $(,)?) => {
-        $crate::sync::Mutex::new(
-            $inner, $crate::optional_name!($($name)?), $crate::static_lock_class!())
+        $crate::sync::Mutex::new_with_name($inner, $crate::optional_name!($($name)?))
     };
 }
 pub use new_mutex;
diff --git a/rust/kernel/sync/lock/spinlock.rs b/rust/kernel/sync/lock/spinlock.rs
index ef76fa07ca3a..091167ceda66 100644
--- a/rust/kernel/sync/lock/spinlock.rs
+++ b/rust/kernel/sync/lock/spinlock.rs
@@ -11,8 +11,7 @@
 #[macro_export]
 macro_rules! new_spinlock {
     ($inner:expr $(, $name:literal)? $(,)?) => {
-        $crate::sync::SpinLock::new(
-            $inner, $crate::optional_name!($($name)?), $crate::static_lock_class!())
+        $crate::sync::SpinLock::new_with_name($inner, $crate::optional_name!($($name)?))
     };
 }
 pub use new_spinlock;

-- 
2.54.0