[PATCH v5 4/4] rust: adding OwnableRefCounted and SimpleOwnableRefCounted

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

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

Signed-off-by: Oliver Mangold <oliver.mangold@pm.me>
---
 rust/kernel/types.rs | 268 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 268 insertions(+)

diff --git a/rust/kernel/types.rs b/rust/kernel/types.rs
index e6f3308f931d90718d405443c3034a216388e0af..0acf95d322b6a213728916f0c7f4095aa3f0e0f0 100644
--- a/rust/kernel/types.rs
+++ b/rust/kernel/types.rs
@@ -552,6 +552,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
@@ -669,6 +675,268 @@ fn drop(&mut self) {
     }
 }
 
+/// A trait for objects that can be wrapped in either one of the reference types
+/// [`Owned`] and [`ARef`].
+///
+/// # Safety
+///
+/// - The same safety requirements as for [`Ownable`] and [`RefCounted`] apply.
+/// - the uniqueness invariant of [`Owned`] is upheld until dropped.
+/// - [`try_from_shared()`](OwnableRefCounted::into_shared) only returns an
+///   [`Owned`] if exactly one [`ARef`] exists.
+/// - [`into_shared()`](OwnableRefCounted::into_shared) set the reference count
+///   to the value which the returned [`ARef`] 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 a [`Owned`] to an [`ARef`].
+///
+/// # 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,
+///         )?;
+///         // SAFETY: we just allocated the `Foo`, thus it is valid
+///         Ok(unsafe { Owned::from_raw(NonNull::new(KBox::into_raw(result)).unwrap()) })
+///     }
+/// }
+///
+/// // SAFETY: we increment and decrement correctly 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 {
+///             // 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 conditions on implementing the trait
+        unsafe { ARef::from_raw(Owned::into_raw(this)) }
+    }
+}
+
+/// This trait allows to implement all of [`Ownable`], [`RefCounted`] and
+/// [`OwnableRefCounted`] together in a simplified way,
+/// only requiring to provide the methods [`inc_ref()`](SimpleOwnableRefCounted::inc_ref),
+/// [`dec_ref()`](SimpleOwnableRefCounted::dec_ref),
+/// and [`is_unique()`](SimpleOwnableRefCounted::is_unique).
+///
+/// For non-standard cases where conversion between [`Ownable`] and [`RefCounted`] needs
+/// or [`Ownable::release()`] and [`RefCounted::dec_ref()`] cannot be the same method,
+/// [`Ownable`], [`RefCounted`] and [`OwnableRefCounted`] should be implemented manually.
+///
+/// # Safety
+///
+/// - The same safety requirements as for [`Ownable`] and [`RefCounted`] apply.
+/// - [`is_unique`](SimpleOwnableRefCounted::is_unique) must only return `true`
+///   in case only one [`ARef`] exists and it is impossible for one to be obtained
+///   other than by cloning an existing [`ARef`] or converting an [`Owned`] to an [`ARef`].
+/// - It is safe to convert an unique [`ARef`] into an [`Owned`]
+///   simply by re-wrapping the underlying object without modifying the refcount.
+///
+/// # Examples
+///
+/// A minimal example implementation of [`SimpleOwnableRefCounted`]
+/// 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, SimpleOwnableRefCounted, Owned,
+/// };
+///
+/// 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,
+///         )?;
+///         // SAFETY: we just allocated the `Foo`, thus it is valid
+///         Ok(unsafe { Owned::from_raw(NonNull::new(KBox::into_raw(result)).unwrap()) })
+///     }
+/// }
+///
+/// // SAFETY: we implement the trait correctly by ensuring
+/// // - the `Foo` is only dropped then the refcount is zero
+/// // - `is_unique()` only returns `true` when the refcount is 1
+/// unsafe impl SimpleOwnableRefCounted 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 {
+///             // 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);
+///         }
+///     }
+///     fn is_unique(&self) -> bool {
+///         self.refcount.get() == 1
+///     }
+/// }
+///
+/// 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 SimpleOwnableRefCounted {
+    /// Checks if exactly one [`ARef`] to the object exists.
+    /// In case the object is [`Sync`] the check needs to be race-free.
+    fn is_unique(&self) -> bool;
+    /// Increments the reference count on the object.
+    fn inc_ref(&self);
+
+    /// Decrements the reference count on the object
+    /// when the [`SimpleOwnableRefCounted`] is dropped.
+    ///
+    /// Frees the object when the count reaches zero.
+    ///
+    /// # Safety
+    ///
+    /// The safety constraints for [`RefCounted::dec_ref`] and
+    /// [`Ownable::release`] both apply to this method, as it will be used
+    /// to implement both of these traits.
+    unsafe fn dec_ref(obj: NonNull<Self>);
+}
+
+// TODO: enable this when compiler supports it (>=1.85)
+// #[diagnostic::do_not_recommend]
+// SAFETY: safe by the requirements on implementation of [`SimpleOwnableRefCounted`]
+unsafe impl<T: SimpleOwnableRefCounted> OwnableRefCounted for T {
+    fn try_from_shared(this: ARef<Self>) -> Result<Owned<Self>, ARef<Self>> {
+        if T::is_unique(&*this) {
+            // SAFETY: safe by the requirements on implementation of [`SimpleOwnable`]
+            Ok(unsafe { Owned::from_raw(ARef::into_raw(this)) })
+        } else {
+            Err(this)
+        }
+    }
+}
+
+// TODO: enable this when compiler supports it (>=1.85)
+// #[diagnostic::do_not_recommend]
+// SAFETY: safe by the requirements on implementation of [`SimpleOwnableRefCounted`]
+unsafe impl<T: SimpleOwnableRefCounted> Ownable for T {
+    unsafe fn release(this: NonNull<Self>) {
+        // SAFETY: safe by the requirements on implementation
+        // of [`SimpleOwnableRefCounted::dec_ref()`]
+        unsafe { SimpleOwnableRefCounted::dec_ref(this) };
+    }
+}
+
+// TODO: enable this when compiler supports it (>=1.85)
+// #[diagnostic::do_not_recommend]
+// SAFETY: safe by the requirements on implementation of [`SimpleOwnableRefCounted`]
+unsafe impl<T: SimpleOwnableRefCounted> RefCounted for T {
+    fn inc_ref(&self) {
+        SimpleOwnableRefCounted::inc_ref(self);
+    }
+    unsafe fn dec_ref(this: NonNull<Self>) {
+        // SAFETY: safe by the requirements on implementation
+        // of [`SimpleOwnableRefCounted::dec_ref()`]
+        unsafe { SimpleOwnableRefCounted::dec_ref(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)
+    }
+}
+
 /// A sum type that always holds either a value of type `L` or `R`.
 ///
 /// # Examples

-- 
2.48.1