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}