[PATCH v11 4/4] rust: Add `OwnableRefCounted`

[Oliver Mangold <oliver.mangold@pm.me>]

Types implementing one of these traits can safely convert between an
`ARef<T>` and an `Owned<T>`.

This is useful for types which generally are accessed through an `ARef`
but have methods which can only safely be called when the reference is
unique, like e.g. `block::mq::Request::end_ok()`.

Signed-off-by: Oliver Mangold <oliver.mangold@pm.me>
Reviewed-by: Andreas Hindborg <a.hindborg@kernel.org>
---
 rust/kernel/types.rs         |  10 +++-
 rust/kernel/types/ownable.rs | 127 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 134 insertions(+), 3 deletions(-)

diff --git a/rust/kernel/types.rs b/rust/kernel/types.rs
index a22d5d39d89edfe357fdfe903ef42e5d5404237f..af0bda4ea50f54aed3c51327db0e43f960868501 100644
--- a/rust/kernel/types.rs
+++ b/rust/kernel/types.rs
@@ -12,7 +12,7 @@
 use pin_init::{PinInit, Zeroable};
 
 pub mod ownable;
-pub use ownable::{Ownable, OwnableMut, Owned};
+pub use ownable::{Ownable, OwnableMut, OwnableRefCounted, Owned};
 
 /// Used to transfer ownership to and from foreign (non-Rust) languages.
 ///
@@ -429,6 +429,8 @@ pub const fn raw_get(this: *const Self) -> *mut T {
 /// Note: Implementing this trait allows types to be wrapped in an [`ARef<Self>`]. It requires an
 /// internal reference count and provides only shared references. If unique references are required
 /// [`Ownable`] should be implemented which allows types to be wrapped in an [`Owned<Self>`].
+/// Implementing the trait [`OwnableRefCounted`] allows to convert between unique and shared
+/// references (i.e. [`Owned<Self>`] and [`ARef<Self>`]).
 ///
 /// # Safety
 ///
@@ -572,6 +574,12 @@ fn from(b: &T) -> Self {
     }
 }
 
+impl<T: OwnableRefCounted> From<Owned<T>> for ARef<T> {
+    fn from(b: Owned<T>) -> Self {
+        T::into_shared(b)
+    }
+}
+
 impl<T: RefCounted> Drop for ARef<T> {
     fn drop(&mut self) {
         // SAFETY: The type invariants guarantee that the `ARef` owns the reference we're about to
diff --git a/rust/kernel/types/ownable.rs b/rust/kernel/types/ownable.rs
index 80cd990f6601767aa5a742a6c0997f4f67d06453..b5626dead6bb25ea76a0ae577db1b130308d98b1 100644
--- a/rust/kernel/types/ownable.rs
+++ b/rust/kernel/types/ownable.rs
@@ -2,6 +2,7 @@
 
 //! Owned reference types.
 
+use crate::types::{ARef, RefCounted};
 use core::{
     marker::PhantomData,
     mem::ManuallyDrop,
@@ -18,8 +19,9 @@
 ///
 /// Note: Implementing this trait allows types to be wrapped in an [`Owned<Self>`]. This does not
 /// provide reference counting but represents a unique, owned reference. If reference counting is
-/// required [`RefCounted`](crate::types::RefCounted) should be implemented which allows types to be
-/// wrapped in an [`ARef<Self>`](crate::types::ARef).
+/// required [`RefCounted`] should be implemented which allows types to be wrapped in an
+/// [`ARef<Self>`]. Implementing the trait [`OwnableRefCounted`] allows to convert between unique
+/// and shared references (i.e. [`Owned<Self>`] and [`ARef<Self>`]).
 ///
 /// # Safety
 ///
@@ -132,3 +134,124 @@ fn drop(&mut self) {
         unsafe { T::release(self.ptr) };
     }
 }
+
+/// A trait for objects that can be wrapped in either one of the reference types [`Owned`] and
+/// [`ARef`].
+///
+/// # Safety
+///
+/// Implementers must ensure that:
+///
+/// - [`try_from_shared()`](OwnableRefCounted::into_shared) only returns an [`Owned<Self>`] if
+///   exactly one [`ARef<Self>`] exists.
+/// - [`into_shared()`](OwnableRefCounted::into_shared) set the reference count to the value which
+///   the returned [`ARef<Self>`] expects for an object with a single reference in existence. This
+///   implies that if [`into_shared()`](OwnableRefCounted::into_shared) is left on the default
+///   implementation, which just rewraps the underlying object, the reference count needs not to be
+///   modified when converting an [`Owned<Self>`] to an [`ARef<Self>`].
+///
+/// # Examples
+///
+/// A minimal example implementation of [`OwnableRefCounted`], [`Ownable`] and its usage with
+/// [`ARef`] and [`Owned`] looks like this:
+///
+/// ```
+/// # #![expect(clippy::disallowed_names)]
+/// use core::cell::Cell;
+/// use core::ptr::NonNull;
+/// use kernel::alloc::{flags, kbox::KBox, AllocError};
+/// use kernel::types::{
+///     ARef, RefCounted, Owned, Ownable, OwnableRefCounted,
+/// };
+///
+/// struct Foo {
+///     refcount: Cell<usize>,
+/// }
+///
+/// impl Foo {
+///     fn new() -> Result<Owned<Self>, AllocError> {
+///         // Use a `KBox` to handle the actual allocation.
+///         let result = KBox::new(
+///             Foo {
+///                 refcount: Cell::new(1),
+///             },
+///             flags::GFP_KERNEL,
+///         )?;
+///         let result = NonNull::new(KBox::into_raw(result))
+///             .expect("Raw pointer to newly allocation KBox is null, this should never happen.");
+///         // SAFETY: We just allocated the `Foo`, thus it is valid.
+///         Ok(unsafe { Owned::from_raw(result) })
+///     }
+/// }
+///
+/// // SAFETY: We increment and decrement each time the respective function is called and only free
+/// // the `Foo` when the refcount reaches zero.
+/// unsafe impl RefCounted for Foo {
+///     fn inc_ref(&self) {
+///         self.refcount.replace(self.refcount.get() + 1);
+///     }
+///
+///     unsafe fn dec_ref(this: NonNull<Self>) {
+///         // SAFETY: The underlying object is always valid when the function is called.
+///         let refcount = unsafe { &this.as_ref().refcount };
+///         let new_refcount = refcount.get() - 1;
+///         if new_refcount == 0 {
+///             // The `Foo` will be dropped when `KBox` goes out of scope.
+///             // SAFETY: The `Box<Foo>` is still alive as the old refcount is 1.
+///             unsafe { KBox::from_raw(this.as_ptr()) };
+///         } else {
+///             refcount.replace(new_refcount);
+///         }
+///     }
+/// }
+///
+/// // SAFETY: We only convert into an `Owned` when the refcount is 1.
+/// unsafe impl OwnableRefCounted for Foo {
+///     fn try_from_shared(this: ARef<Self>) -> Result<Owned<Self>, ARef<Self>> {
+///         if this.refcount.get() == 1 {
+///             // SAFETY: The `Foo` is still alive as the refcount is 1.
+///             Ok(unsafe { Owned::from_raw(ARef::into_raw(this)) })
+///         } else {
+///             Err(this)
+///         }
+///     }
+/// }
+///
+/// // SAFETY: We are not `AlwaysRefCounted`.
+/// unsafe impl Ownable for Foo {
+///     unsafe fn release(this: NonNull<Self>) {
+///         // SAFETY: Using `dec_ref()` from `RefCounted` to release is okay, as the refcount is
+///         // always 1 for an `Owned<Foo>`.
+///         unsafe{ Foo::dec_ref(this) };
+///     }
+/// }
+///
+/// let foo = Foo::new().unwrap();
+/// let mut foo = ARef::from(foo);
+/// {
+///     let bar = foo.clone();
+///     assert!(Owned::try_from(bar).is_err());
+/// }
+/// assert!(Owned::try_from(foo).is_ok());
+/// ```
+pub unsafe trait OwnableRefCounted: RefCounted + Ownable + Sized {
+    /// Checks if the [`ARef`] is unique and convert it to an [`Owned`] it that is that case.
+    /// Otherwise it returns again an [`ARef`] to the same underlying object.
+    fn try_from_shared(this: ARef<Self>) -> Result<Owned<Self>, ARef<Self>>;
+
+    /// Converts the [`Owned`] into an [`ARef`].
+    fn into_shared(this: Owned<Self>) -> ARef<Self> {
+        // SAFETY: Safe by the requirements on implementing the trait.
+        unsafe { ARef::from_raw(Owned::into_raw(this)) }
+    }
+}
+
+impl<T: OwnableRefCounted> TryFrom<ARef<T>> for Owned<T> {
+    type Error = ARef<T>;
+    /// Tries to convert the [`ARef`] to an [`Owned`] by calling
+    /// [`try_from_shared()`](OwnableRefCounted::try_from_shared). In case the [`ARef`] is not
+    /// unique, it returns again an [`ARef`] to the same underlying object.
+    fn try_from(b: ARef<T>) -> Result<Owned<T>, Self::Error> {
+        T::try_from_shared(b)
+    }
+}

-- 
2.49.0