kernel/sync/
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}