Add OwnedMappedMutexGuard for OwnedMutexGuard

futures-util/src/lock/mod.rs

@@ -14,7 +14,8 @@ pub use self::bilock::{BiLock, BiLockAcquire, BiLockGuard, ReuniteError};
 #[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
 #[cfg(feature = "std")]
 pub use self::mutex::{
-    MappedMutexGuard, Mutex, MutexGuard, MutexLockFuture, OwnedMutexGuard, OwnedMutexLockFuture,
+    MappedMutexGuard, Mutex, MutexGuard, MutexLockFuture, OwnedMappedMutexGuard, OwnedMutexGuard,
+    OwnedMutexLockFuture,
 };
 
 #[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]

futures-util/src/lock/mutex.rs

@@ -1,3 +1,4 @@
+use core::ptr;
 use std::cell::UnsafeCell;
 use std::marker::PhantomData;
 use std::ops::{Deref, DerefMut};
@@ -283,6 +284,51 @@ pub struct OwnedMutexGuard<T: ?Sized> {
     mutex: Arc<Mutex<T>>,
 }
 
+/// A owned handle to a held `Mutex`.
+struct OwnedMutexGuardExtendedLock<T: ?Sized> {
+    mutex: Arc<Mutex<T>>,
+}
+
+impl<T: ?Sized> OwnedMutexGuard<T> {
+    /// `skip_drop` prevents the `OwnedMutexGuard` from being automatically dropped, allowing manual control over the drop process.
+    /// This method returns an `OwnedMutexGuardExtendedLock` which contains the original `Arc<Mutex<T>>`.
+    fn skip_drop(self) -> OwnedMutexGuardExtendedLock<T> {
+        // Prevents automatic drop by wrapping in `ManuallyDrop`
+        let man = mem::ManuallyDrop::new(self);
+        // Unsafely reads the `Arc<Mutex<T>>` from the `ManuallyDrop` wrapper
+        OwnedMutexGuardExtendedLock { mutex: unsafe { ptr::read(&man.mutex) } }
+    }
+
+    /// Returns a locked view over a portion of the locked data.
+    #[inline]
+    pub fn map<U, F>(mut this: Self, f: F) -> OwnedMappedMutexGuard<T, U>
+    where
+        U: ?Sized,
+        F: FnOnce(&mut T) -> &mut U,
+    {
+        let value = f(&mut *this) as *mut U;
+        let mutex = this.skip_drop();
+        OwnedMappedMutexGuard { value, mutex: mutex.mutex }
+    }
+
+    /// Returns a locked view over a portion of the locked data.
+    /// If the closure returns `None`, the original `OwnedMappedMutexGuard` is returned.
+    #[inline]
+    pub fn try_map<U, F>(mut this: Self, f: F) -> Result<OwnedMappedMutexGuard<T, U>, Self>
+    where
+        U: ?Sized,
+        F: FnOnce(&mut T) -> Option<&mut U>,
+    {
+        let value = match f(&mut *this) {
+            Some(data) => data as *mut U,
+            None => return Err(this),
+        };
+        let mutex = this.skip_drop();
+
+        Ok(OwnedMappedMutexGuard { value, mutex: mutex.mutex })
+    }
+}
+
 impl<T: ?Sized + fmt::Debug> fmt::Debug for OwnedMutexGuard<T> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("OwnedMutexGuard")
@@ -311,6 +357,60 @@ impl<T: ?Sized> DerefMut for OwnedMutexGuard<T> {
     }
 }
 
+/// An RAII guard returned by the `OwnedMutexGuard::map` and `OwnedMappedMutexGuard::map` methods.
+/// When this structure is dropped (falls out of scope), the lock will be unlocked.
+pub struct OwnedMappedMutexGuard<T: ?Sized, U: ?Sized = T> {
+    value: *mut U,
+    mutex: Arc<Mutex<T>>,
+}
+
+/// A owned handle to a held `Mutex`.
+struct OwnedMappedMutexGuardExtendedLock<T: ?Sized, U: ?Sized> {
+    _value: *mut U,
+    mutex: Arc<Mutex<T>>,
+}
+
+impl<T: ?Sized, U: ?Sized> OwnedMappedMutexGuard<T, U> {
+    /// `skip_drop` prevents the `OwnedMutexGuard` from being automatically dropped, allowing manual control over the drop process.
+    /// This method returns an `OwnedMappedMutexGuardExtendedLock` which contains the original `Arc<Mutex<T>>`.
+    fn skip_drop(self) -> OwnedMappedMutexGuardExtendedLock<T, U> {
+        // Prevents automatic drop by wrapping in `ManuallyDrop`
+        let selfi = mem::ManuallyDrop::new(self);
+        // Unsafely reads the `Arc<Mutex<T>>` from the `ManuallyDrop` wrapper
+        OwnedMappedMutexGuardExtendedLock {
+            mutex: unsafe { ptr::read(&selfi.mutex) },
+            _value: selfi.value,
+        }
+    }
+
+    /// Returns a locked view over a portion of the locked data.
+    #[inline]
+    pub fn map<C, F>(mut this: Self, f: F) -> OwnedMappedMutexGuard<T, C>
+    where
+        F: FnOnce(&mut U) -> &mut C,
+    {
+        let value = f(&mut *this) as *mut C;
+        let mutex = this.skip_drop();
+        OwnedMappedMutexGuard { value, mutex: mutex.mutex }
+    }
+
+    /// Returns a locked view over a portion of the locked data.
+    /// If the closure returns `None`, the original `OwnedMappedMutexGuard` is returned.
+    #[inline]
+    pub fn try_map<C, F>(mut this: Self, f: F) -> Result<OwnedMappedMutexGuard<T, C>, Self>
+    where
+        F: FnOnce(&mut U) -> Option<&mut C>,
+    {
+        let value = match f(&mut *this) {
+            Some(data) => data as *mut C,
+            None => return Err(this),
+        };
+        let mutex = this.skip_drop();
+
+        Ok(OwnedMappedMutexGuard { value, mutex: mutex.mutex })
+    }
+}
+
 /// A future which resolves when the target mutex has been successfully acquired.
 pub struct MutexLockFuture<'a, T: ?Sized> {
     // `None` indicates that the mutex was successfully acquired.
@@ -544,6 +644,52 @@ unsafe impl<T: ?Sized + Sync> Sync for OwnedMutexGuard<T> {}
 unsafe impl<T: ?Sized + Send, U: ?Sized + Send> Send for MappedMutexGuard<'_, T, U> {}
 unsafe impl<T: ?Sized + Sync, U: ?Sized + Sync> Sync for MappedMutexGuard<'_, T, U> {}
 
+impl<T: ?Sized, U: ?Sized> Drop for OwnedMappedMutexGuard<T, U> {
+    fn drop(&mut self) {
+        self.mutex.unlock()
+    }
+}
+
+unsafe impl<T, U> Sync for OwnedMappedMutexGuard<T, U>
+where
+    T: ?Sized + Send + Sync,
+    U: ?Sized + Send + Sync,
+{
+}
+unsafe impl<T, U> Send for OwnedMappedMutexGuard<T, U>
+where
+    T: ?Sized + Send,
+    U: ?Sized + Send,
+{
+}
+
+impl<T: ?Sized, U: ?Sized> Deref for OwnedMappedMutexGuard<T, U> {
+    type Target = U;
+    /// Deref implementation to allow access to the mapped value.
+    fn deref(&self) -> &Self::Target {
+        unsafe { &*self.value }
+    }
+}
+
+impl<T: ?Sized, U: ?Sized> DerefMut for OwnedMappedMutexGuard<T, U> {
+    /// DerefMut implementation to allow mutable access to the mapped value.
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        unsafe { &mut *self.value }
+    }
+}
+
+impl<T: ?Sized, U: ?Sized + fmt::Debug> fmt::Debug for OwnedMappedMutexGuard<T, U> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Debug::fmt(&**self, f)
+    }
+}
+
+impl<T: ?Sized, U: ?Sized + fmt::Display> fmt::Display for OwnedMappedMutexGuard<T, U> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Display::fmt(&**self, f)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;