kernel/sync/
aref.rs

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