aref.rs

  1// SPDX-License-Identifier: GPL-2.0
  2
  3//! Internal reference counting support.
  4//!
  5//! Many C types already have their own reference counting mechanism (e.g. by storing a
  6//! `refcount_t`). This module provides support for directly using their internal reference count
  7//! from Rust; instead of making users have to use an additional Rust-reference count in the form of
  8//! [`Arc`].
  9//!
 10//! The smart pointer [`ARef<T>`] acts similarly to [`Arc<T>`] in that it holds a refcount on the
 11//! underlying object, but this refcount is internal to the object. It essentially is a Rust
 12//! implementation of the `get_` and `put_` pattern used in C for reference counting.
 13//!
 14//! To make use of [`ARef<MyType>`], `MyType` needs to implement [`AlwaysRefCounted`]. It is a trait
 15//! for accessing the internal reference count of an object of the `MyType` type.
 16//!
 17//! [`Arc`]: crate::sync::Arc
 18//! [`Arc<T>`]: crate::sync::Arc
 19
 20use core::{marker::PhantomData, mem::ManuallyDrop, ops::Deref, ptr::NonNull};
 21
 22/// Types that are _always_ reference counted.
 23///
 24/// It allows such types to define their own custom ref increment and decrement functions.
 25/// Additionally, it allows users to convert from a shared reference `&T` to an owned reference
 26/// [`ARef<T>`].
 27///
 28/// This is usually implemented by wrappers to existing structures on the C side of the code. For
 29/// Rust code, the recommendation is to use [`Arc`](crate::sync::Arc) to create reference-counted
 30/// instances of a type.
 31///
 32/// # Safety
 33///
 34/// Implementers must ensure that increments to the reference count keep the object alive in memory
 35/// at least until matching decrements are performed.
 36///
 37/// Implementers must also ensure that all instances are reference-counted. (Otherwise they
 38/// won't be able to honour the requirement that [`AlwaysRefCounted::inc_ref`] keep the object
 39/// alive.)
 40pub unsafe trait AlwaysRefCounted {
 41    /// Increments the reference count on the object.
 42    fn inc_ref(&self);
 43
 44    /// Decrements the reference count on the object.
 45    ///
 46    /// Frees the object when the count reaches zero.
 47    ///
 48    /// # Safety
 49    ///
 50    /// Callers must ensure that there was a previous matching increment to the reference count,
 51    /// and that the object is no longer used after its reference count is decremented (as it may
 52    /// result in the object being freed), unless the caller owns another increment on the refcount
 53    /// (e.g., it calls [`AlwaysRefCounted::inc_ref`] twice, then calls
 54    /// [`AlwaysRefCounted::dec_ref`] once).
 55    unsafe fn dec_ref(obj: NonNull<Self>);
 56}
 57
 58/// An owned reference to an always-reference-counted object.
 59///
 60/// The object's reference count is automatically decremented when an instance of [`ARef`] is
 61/// dropped. It is also automatically incremented when a new instance is created via
 62/// [`ARef::clone`].
 63///
 64/// # Invariants
 65///
 66/// The pointer stored in `ptr` is non-null and valid for the lifetime of the [`ARef`] instance. In
 67/// particular, the [`ARef`] instance owns an increment on the underlying object's reference count.
 68pub struct ARef<T: AlwaysRefCounted> {
 69    ptr: NonNull<T>,
 70    _p: PhantomData<T>,
 71}
 72
 73// SAFETY: It is safe to send `ARef<T>` to another thread when the underlying `T` is `Sync` because
 74// it effectively means sharing `&T` (which is safe because `T` is `Sync`); additionally, it needs
 75// `T` to be `Send` because any thread that has an `ARef<T>` may ultimately access `T` using a
 76// mutable reference, for example, when the reference count reaches zero and `T` is dropped.
 77unsafe impl<T: AlwaysRefCounted + Sync + Send> Send for ARef<T> {}
 78
 79// SAFETY: It is safe to send `&ARef<T>` to another thread when the underlying `T` is `Sync`
 80// because it effectively means sharing `&T` (which is safe because `T` is `Sync`); additionally,
 81// it needs `T` to be `Send` because any thread that has a `&ARef<T>` may clone it and get an
 82// `ARef<T>` on that thread, so the thread may ultimately access `T` using a mutable reference, for
 83// example, when the reference count reaches zero and `T` is dropped.
 84unsafe impl<T: AlwaysRefCounted + Sync + Send> Sync for ARef<T> {}
 85
 86impl<T: AlwaysRefCounted> ARef<T> {
 87    /// Creates a new instance of [`ARef`].
 88    ///
 89    /// It takes over an increment of the reference count on the underlying object.
 90    ///
 91    /// # Safety
 92    ///
 93    /// Callers must ensure that the reference count was incremented at least once, and that they
 94    /// are properly relinquishing one increment. That is, if there is only one increment, callers
 95    /// must not use the underlying object anymore -- it is only safe to do so via the newly
 96    /// created [`ARef`].
 97    pub unsafe fn from_raw(ptr: NonNull<T>) -> Self {
 98        // INVARIANT: The safety requirements guarantee that the new instance now owns the
 99        // increment on the refcount.
100        Self {
101            ptr,
102            _p: PhantomData,
103        }
104    }
105
106    /// Consumes the `ARef`, returning a raw pointer.
107    ///
108    /// This function does not change the refcount. After calling this function, the caller is
109    /// responsible for the refcount previously managed by the `ARef`.
110    ///
111    /// # Examples
112    ///
113    /// ```
114    /// use core::ptr::NonNull;
115    /// use kernel::sync::aref::{ARef, AlwaysRefCounted};
116    ///
117    /// struct Empty {}
118    ///
119    /// # // SAFETY: TODO.
120    /// unsafe impl AlwaysRefCounted for Empty {
121    ///     fn inc_ref(&self) {}
122    ///     unsafe fn dec_ref(_obj: NonNull<Self>) {}
123    /// }
124    ///
125    /// let mut data = Empty {};
126    /// let ptr = NonNull::<Empty>::new(&mut data).unwrap();
127    /// # // SAFETY: TODO.
128    /// let data_ref: ARef<Empty> = unsafe { ARef::from_raw(ptr) };
129    /// let raw_ptr: NonNull<Empty> = ARef::into_raw(data_ref);
130    ///
131    /// assert_eq!(ptr, raw_ptr);
132    /// ```
133    pub fn into_raw(me: Self) -> NonNull<T> {
134        ManuallyDrop::new(me).ptr
135    }
136}
137
138impl<T: AlwaysRefCounted> Clone for ARef<T> {
139    fn clone(&self) -> Self {
140        self.inc_ref();
141        // SAFETY: We just incremented the refcount above.
142        unsafe { Self::from_raw(self.ptr) }
143    }
144}
145
146impl<T: AlwaysRefCounted> Deref for ARef<T> {
147    type Target = T;
148
149    fn deref(&self) -> &Self::Target {
150        // SAFETY: The type invariants guarantee that the object is valid.
151        unsafe { self.ptr.as_ref() }
152    }
153}
154
155impl<T: AlwaysRefCounted> From<&T> for ARef<T> {
156    fn from(b: &T) -> Self {
157        b.inc_ref();
158        // SAFETY: We just incremented the refcount above.
159        unsafe { Self::from_raw(NonNull::from(b)) }
160    }
161}
162
163impl<T: AlwaysRefCounted> Drop for ARef<T> {
164    fn drop(&mut self) {
165        // SAFETY: The type invariants guarantee that the `ARef` owns the reference we're about to
166        // decrement.
167        unsafe { T::dec_ref(self.ptr) };
168    }
169}