[PATCH v4 2/2] rust: sync: add `Arc::into_unique_or_drop`

[Alice Ryhl]

Decrement the refcount of an `Arc`, but handle the case where it hits
zero by taking ownership of the now-unique `Arc`, instead of destroying
and deallocating it.

This is a dependency of the linked list that Rust Binder uses. The
linked list uses this method as part of its `ListArc` abstraction [1].

Boqun Feng has authored the examples.

Reviewed-by: Benno Lossin <benno.lossin@xxxxxxxxx>
Reviewed-by: Boqun Feng <boqun.feng@xxxxxxxxx>
Link: https://lore.kernel.org/r/20240402-linked-list-v1-1-b1c59ba7ae3b@xxxxxxxxxx [1]
Co-developed-by: Boqun Feng <boqun.feng@xxxxxxxxx>
Signed-off-by: Boqun Feng <boqun.feng@xxxxxxxxx>
Signed-off-by: Alice Ryhl <aliceryhl@xxxxxxxxxx>
---
 rust/kernel/sync/arc.rs | 62 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/rust/kernel/sync/arc.rs b/rust/kernel/sync/arc.rs
index 53addb8876c2..30777241d8cf 100644
--- a/rust/kernel/sync/arc.rs
+++ b/rust/kernel/sync/arc.rs
@@ -290,6 +290,68 @@ pub fn as_arc_borrow(&self) -> ArcBorrow<'_, T> {
     pub fn ptr_eq(this: &Self, other: &Self) -> bool {
         core::ptr::eq(this.ptr.as_ptr(), other.ptr.as_ptr())
     }
+
+    /// Converts this [`Arc`] into a [`UniqueArc`], or destroys it if it is not unique.
+    ///
+    /// When this destroys the `Arc`, it does so while properly avoiding races. This means that
+    /// this method will never call the destructor of the value.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::sync::{Arc, UniqueArc};
+    ///
+    /// let arc = Arc::try_new(42)?;
+    /// let unique_arc = arc.into_unique_or_drop();
+    ///
+    /// // The above conversion should succeed since refcount of `arc` is 1.
+    /// assert!(unique_arc.is_some());
+    ///
+    /// assert_eq!(*(unique_arc.unwrap()), 42);
+    ///
+    /// # Ok::<(), Error>(())
+    /// ```
+    ///
+    /// ```
+    /// use kernel::sync::{Arc, UniqueArc};
+    ///
+    /// let arc = Arc::try_new(42)?;
+    /// let another = arc.clone();
+    ///
+    /// let unique_arc = arc.into_unique_or_drop();
+    ///
+    /// // The above conversion should fail since refcount of `arc` is >1.
+    /// assert!(unique_arc.is_none());
+    ///
+    /// # Ok::<(), Error>(())
+    /// ```
+    pub fn into_unique_or_drop(self) -> Option<Pin<UniqueArc<T>>> {
+        // We will manually manage the refcount in this method, so we disable the destructor.
+        let me = ManuallyDrop::new(self);
+        // SAFETY: We own a refcount, so the pointer is still valid.
+        let refcount = unsafe { me.ptr.as_ref() }.refcount.get();
+
+        // If the refcount reaches a non-zero value, then we have destroyed this `Arc` and will
+        // return without further touching the `Arc`. If the refcount reaches zero, then there are
+        // no other arcs, and we can create a `UniqueArc`.
+        //
+        // SAFETY: We own a refcount, so the pointer is not dangling.
+        let is_zero = unsafe { bindings::refcount_dec_and_test(refcount) };
+        if is_zero {
+            // SAFETY: We have exclusive access to the arc, so we can perform unsynchronized
+            // accesses to the refcount.
+            unsafe { core::ptr::write(refcount, bindings::REFCOUNT_INIT(1)) };
+
+            // INVARIANT: We own the only refcount to this arc, so we may create a `UniqueArc`. We
+            // must pin the `UniqueArc` because the values was previously in an `Arc`, and they pin
+            // their values.
+            Some(Pin::from(UniqueArc {
+                inner: ManuallyDrop::into_inner(me),
+            }))
+        } else {
+            None
+        }
+    }
 }
 
 impl<T: 'static> ForeignOwnable for Arc<T> {

-- 
2.44.0.478.gd926399ef9-goog