xref: /linux/rust/kernel/maple_tree.rs (revision da939ef4c494246bc2102ecb628bbcc71d650410)
1 // SPDX-License-Identifier: GPL-2.0
2 
3 //! Maple trees.
4 //!
5 //! C header: [`include/linux/maple_tree.h`](srctree/include/linux/maple_tree.h)
6 //!
7 //! Reference: <https://docs.kernel.org/core-api/maple_tree.html>
8 
9 use core::{
10     marker::PhantomData,
11     ops::{Bound, RangeBounds},
12     ptr,
13 };
14 
15 use kernel::{
16     alloc::Flags,
17     error::to_result,
18     prelude::*,
19     types::{ForeignOwnable, Opaque},
20 };
21 
22 /// A maple tree optimized for storing non-overlapping ranges.
23 ///
24 /// # Invariants
25 ///
26 /// Each range in the maple tree owns an instance of `T`.
27 #[pin_data(PinnedDrop)]
28 #[repr(transparent)]
29 pub struct MapleTree<T: ForeignOwnable> {
30     #[pin]
31     tree: Opaque<bindings::maple_tree>,
32     _p: PhantomData<T>,
33 }
34 
35 #[inline]
36 fn to_maple_range(range: impl RangeBounds<usize>) -> Option<(usize, usize)> {
37     let first = match range.start_bound() {
38         Bound::Included(start) => *start,
39         Bound::Excluded(start) => start.checked_add(1)?,
40         Bound::Unbounded => 0,
41     };
42 
43     let last = match range.end_bound() {
44         Bound::Included(end) => *end,
45         Bound::Excluded(end) => end.checked_sub(1)?,
46         Bound::Unbounded => usize::MAX,
47     };
48 
49     if last < first {
50         return None;
51     }
52 
53     Some((first, last))
54 }
55 
56 impl<T: ForeignOwnable> MapleTree<T> {
57     /// Create a new maple tree.
58     ///
59     /// The tree will use the regular implementation with a higher branching factor, rather than
60     /// the allocation tree.
61     #[inline]
62     pub fn new() -> impl PinInit<Self> {
63         pin_init!(MapleTree {
64             // SAFETY: This initializes a maple tree into a pinned slot. The maple tree will be
65             // destroyed in Drop before the memory location becomes invalid.
66             tree <- Opaque::ffi_init(|slot| unsafe { bindings::mt_init_flags(slot, 0) }),
67             _p: PhantomData,
68         })
69     }
70 
71     /// Insert the value at the given index.
72     ///
73     /// # Errors
74     ///
75     /// If the maple tree already contains a range using the given index, then this call will
76     /// return an [`InsertErrorKind::Occupied`]. It may also fail if memory allocation fails.
77     ///
78     /// # Examples
79     ///
80     /// ```
81     /// use kernel::maple_tree::{InsertErrorKind, MapleTree};
82     ///
83     /// let tree = KBox::pin_init(MapleTree::<KBox<i32>>::new(), GFP_KERNEL)?;
84     ///
85     /// let ten = KBox::new(10, GFP_KERNEL)?;
86     /// let twenty = KBox::new(20, GFP_KERNEL)?;
87     /// let the_answer = KBox::new(42, GFP_KERNEL)?;
88     ///
89     /// // These calls will succeed.
90     /// tree.insert(100, ten, GFP_KERNEL)?;
91     /// tree.insert(101, twenty, GFP_KERNEL)?;
92     ///
93     /// // This will fail because the index is already in use.
94     /// assert_eq!(
95     ///     tree.insert(100, the_answer, GFP_KERNEL).unwrap_err().cause,
96     ///     InsertErrorKind::Occupied,
97     /// );
98     /// # Ok::<_, Error>(())
99     /// ```
100     #[inline]
101     pub fn insert(&self, index: usize, value: T, gfp: Flags) -> Result<(), InsertError<T>> {
102         self.insert_range(index..=index, value, gfp)
103     }
104 
105     /// Insert a value to the specified range, failing on overlap.
106     ///
107     /// This accepts the usual types of Rust ranges using the `..` and `..=` syntax for exclusive
108     /// and inclusive ranges respectively. The range must not be empty, and must not overlap with
109     /// any existing range.
110     ///
111     /// # Errors
112     ///
113     /// If the maple tree already contains an overlapping range, then this call will return an
114     /// [`InsertErrorKind::Occupied`]. It may also fail if memory allocation fails or if the
115     /// requested range is invalid (e.g. empty).
116     ///
117     /// # Examples
118     ///
119     /// ```
120     /// use kernel::maple_tree::{InsertErrorKind, MapleTree};
121     ///
122     /// let tree = KBox::pin_init(MapleTree::<KBox<i32>>::new(), GFP_KERNEL)?;
123     ///
124     /// let ten = KBox::new(10, GFP_KERNEL)?;
125     /// let twenty = KBox::new(20, GFP_KERNEL)?;
126     /// let the_answer = KBox::new(42, GFP_KERNEL)?;
127     /// let hundred = KBox::new(100, GFP_KERNEL)?;
128     ///
129     /// // Insert the value 10 at the indices 100 to 499.
130     /// tree.insert_range(100..500, ten, GFP_KERNEL)?;
131     ///
132     /// // Insert the value 20 at the indices 500 to 1000.
133     /// tree.insert_range(500..=1000, twenty, GFP_KERNEL)?;
134     ///
135     /// // This will fail due to overlap with the previous range on index 1000.
136     /// assert_eq!(
137     ///     tree.insert_range(1000..1200, the_answer, GFP_KERNEL).unwrap_err().cause,
138     ///     InsertErrorKind::Occupied,
139     /// );
140     ///
141     /// // When using .. to specify the range, you must be careful to ensure that the range is
142     /// // non-empty.
143     /// assert_eq!(
144     ///     tree.insert_range(72..72, hundred, GFP_KERNEL).unwrap_err().cause,
145     ///     InsertErrorKind::InvalidRequest,
146     /// );
147     /// # Ok::<_, Error>(())
148     /// ```
149     pub fn insert_range<R>(&self, range: R, value: T, gfp: Flags) -> Result<(), InsertError<T>>
150     where
151         R: RangeBounds<usize>,
152     {
153         let Some((first, last)) = to_maple_range(range) else {
154             return Err(InsertError {
155                 value,
156                 cause: InsertErrorKind::InvalidRequest,
157             });
158         };
159 
160         let ptr = T::into_foreign(value);
161 
162         // SAFETY: The tree is valid, and we are passing a pointer to an owned instance of `T`.
163         let res = to_result(unsafe {
164             bindings::mtree_insert_range(self.tree.get(), first, last, ptr, gfp.as_raw())
165         });
166 
167         if let Err(err) = res {
168             // SAFETY: As `mtree_insert_range` failed, it is safe to take back ownership.
169             let value = unsafe { T::from_foreign(ptr) };
170 
171             let cause = if err == ENOMEM {
172                 InsertErrorKind::AllocError(kernel::alloc::AllocError)
173             } else if err == EEXIST {
174                 InsertErrorKind::Occupied
175             } else {
176                 InsertErrorKind::InvalidRequest
177             };
178             Err(InsertError { value, cause })
179         } else {
180             Ok(())
181         }
182     }
183 
184     /// Erase the range containing the given index.
185     ///
186     /// # Examples
187     ///
188     /// ```
189     /// use kernel::maple_tree::MapleTree;
190     ///
191     /// let tree = KBox::pin_init(MapleTree::<KBox<i32>>::new(), GFP_KERNEL)?;
192     ///
193     /// let ten = KBox::new(10, GFP_KERNEL)?;
194     /// let twenty = KBox::new(20, GFP_KERNEL)?;
195     ///
196     /// tree.insert_range(100..500, ten, GFP_KERNEL)?;
197     /// tree.insert(67, twenty, GFP_KERNEL)?;
198     ///
199     /// assert_eq!(tree.erase(67).map(|v| *v), Some(20));
200     /// assert_eq!(tree.erase(275).map(|v| *v), Some(10));
201     ///
202     /// // The previous call erased the entire range, not just index 275.
203     /// assert!(tree.erase(127).is_none());
204     /// # Ok::<_, Error>(())
205     /// ```
206     #[inline]
207     pub fn erase(&self, index: usize) -> Option<T> {
208         // SAFETY: `self.tree` contains a valid maple tree.
209         let ret = unsafe { bindings::mtree_erase(self.tree.get(), index) };
210 
211         // SAFETY: If the pointer is not null, then we took ownership of a valid instance of `T`
212         // from the tree.
213         unsafe { T::try_from_foreign(ret) }
214     }
215 
216     /// Free all `T` instances in this tree.
217     ///
218     /// # Safety
219     ///
220     /// This frees Rust data referenced by the maple tree without removing it from the maple tree,
221     /// leaving it in an invalid state. The caller must ensure that this invalid state cannot be
222     /// observed by the end-user.
223     unsafe fn free_all_entries(self: Pin<&mut Self>) {
224         // SAFETY: The caller provides exclusive access to the entire maple tree, so we have
225         // exclusive access to the entire maple tree despite not holding the lock.
226         let mut ma_state = unsafe { MaState::new_raw(self.into_ref().get_ref(), 0, usize::MAX) };
227 
228         loop {
229             // This uses the raw accessor because we're destroying pointers without removing them
230             // from the maple tree, which is only valid because this is the destructor.
231             let ptr = ma_state.mas_find_raw(usize::MAX);
232             if ptr.is_null() {
233                 break;
234             }
235             // SAFETY: By the type invariants, this pointer references a valid value of type `T`.
236             // By the safety requirements, it is okay to free it without removing it from the maple
237             // tree.
238             drop(unsafe { T::from_foreign(ptr) });
239         }
240     }
241 }
242 
243 #[pinned_drop]
244 impl<T: ForeignOwnable> PinnedDrop for MapleTree<T> {
245     #[inline]
246     fn drop(mut self: Pin<&mut Self>) {
247         // We only iterate the tree if the Rust value has a destructor.
248         if core::mem::needs_drop::<T>() {
249             // SAFETY: Other than the below `mtree_destroy` call, the tree will not be accessed
250             // after this call.
251             unsafe { self.as_mut().free_all_entries() };
252         }
253 
254         // SAFETY: The tree is valid, and will not be accessed after this call.
255         unsafe { bindings::mtree_destroy(self.tree.get()) };
256     }
257 }
258 
259 /// A helper type used for navigating a [`MapleTree`].
260 ///
261 /// # Invariants
262 ///
263 /// For the duration of `'tree`:
264 ///
265 /// * The `ma_state` references a valid `MapleTree<T>`.
266 /// * The `ma_state` has read/write access to the tree.
267 pub struct MaState<'tree, T: ForeignOwnable> {
268     state: bindings::ma_state,
269     _phantom: PhantomData<&'tree mut MapleTree<T>>,
270 }
271 
272 impl<'tree, T: ForeignOwnable> MaState<'tree, T> {
273     /// Initialize a new `MaState` with the given tree.
274     ///
275     /// # Safety
276     ///
277     /// The caller must ensure that this `MaState` has read/write access to the maple tree.
278     #[inline]
279     unsafe fn new_raw(mt: &'tree MapleTree<T>, first: usize, end: usize) -> Self {
280         // INVARIANT:
281         // * Having a reference ensures that the `MapleTree<T>` is valid for `'tree`.
282         // * The caller ensures that we have read/write access.
283         Self {
284             state: bindings::ma_state {
285                 tree: mt.tree.get(),
286                 index: first,
287                 last: end,
288                 node: ptr::null_mut(),
289                 status: bindings::maple_status_ma_start,
290                 min: 0,
291                 max: usize::MAX,
292                 alloc: ptr::null_mut(),
293                 mas_flags: 0,
294                 store_type: bindings::store_type_wr_invalid,
295                 ..Default::default()
296             },
297             _phantom: PhantomData,
298         }
299     }
300 
301     #[inline]
302     fn as_raw(&mut self) -> *mut bindings::ma_state {
303         &raw mut self.state
304     }
305 
306     #[inline]
307     fn mas_find_raw(&mut self, max: usize) -> *mut c_void {
308         // SAFETY: By the type invariants, the `ma_state` is active and we have read/write access
309         // to the tree.
310         unsafe { bindings::mas_find(self.as_raw(), max) }
311     }
312 }
313 
314 /// Error type for failure to insert a new value.
315 pub struct InsertError<T> {
316     /// The value that could not be inserted.
317     pub value: T,
318     /// The reason for the failure to insert.
319     pub cause: InsertErrorKind,
320 }
321 
322 /// The reason for the failure to insert.
323 #[derive(PartialEq, Eq, Copy, Clone, Debug)]
324 pub enum InsertErrorKind {
325     /// There is already a value in the requested range.
326     Occupied,
327     /// Failure to allocate memory.
328     AllocError(kernel::alloc::AllocError),
329     /// The insertion request was invalid.
330     InvalidRequest,
331 }
332 
333 impl From<InsertErrorKind> for Error {
334     #[inline]
335     fn from(kind: InsertErrorKind) -> Error {
336         match kind {
337             InsertErrorKind::Occupied => EEXIST,
338             InsertErrorKind::AllocError(kernel::alloc::AllocError) => ENOMEM,
339             InsertErrorKind::InvalidRequest => EINVAL,
340         }
341     }
342 }
343 
344 impl<T> From<InsertError<T>> for Error {
345     #[inline]
346     fn from(insert_err: InsertError<T>) -> Error {
347         Error::from(insert_err.cause)
348     }
349 }
350