1 // SPDX-License-Identifier: GPL-2.0
2 
3 //! Internal reference counting support.
4 
5 use 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.)
25 pub unsafe trait AlwaysRefCounted {
26     /// Increments the reference count on the object.
inc_ref(&self)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).
dec_ref(obj: NonNull<Self>)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.
53 pub 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.
62 unsafe 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.
69 unsafe impl<T: AlwaysRefCounted + Sync + Send> Sync for ARef<T> {}
70 
71 impl<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`].
from_raw(ptr: NonNull<T>) -> Self82     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     /// ```
into_raw(me: Self) -> NonNull<T>118     pub fn into_raw(me: Self) -> NonNull<T> {
119         ManuallyDrop::new(me).ptr
120     }
121 }
122 
123 impl<T: AlwaysRefCounted> Clone for ARef<T> {
clone(&self) -> Self124     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 
131 impl<T: AlwaysRefCounted> Deref for ARef<T> {
132     type Target = T;
133 
deref(&self) -> &Self::Target134     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 
140 impl<T: AlwaysRefCounted> From<&T> for ARef<T> {
from(b: &T) -> Self141     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 
148 impl<T: AlwaysRefCounted> Drop for ARef<T> {
drop(&mut self)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 }
155