kernel/
xarray.rs

1// SPDX-License-Identifier: GPL-2.0
2
3//! XArray abstraction.
4//!
5//! C header: [`include/linux/xarray.h`](srctree/include/linux/xarray.h)
6
7use crate::{
8    alloc, bindings, build_assert,
9    error::{Error, Result},
10    ffi::c_void,
11    types::{ForeignOwnable, NotThreadSafe, Opaque},
12};
13use core::{iter, marker::PhantomData, pin::Pin, ptr::NonNull};
14use pin_init::{pin_data, pin_init, pinned_drop, PinInit};
15
16/// An array which efficiently maps sparse integer indices to owned objects.
17///
18/// This is similar to a [`crate::alloc::kvec::Vec<Option<T>>`], but more efficient when there are
19/// holes in the index space, and can be efficiently grown.
20///
21/// # Invariants
22///
23/// `self.xa` is always an initialized and valid [`bindings::xarray`] whose entries are either
24/// `XA_ZERO_ENTRY` or came from `T::into_foreign`.
25///
26/// # Examples
27///
28/// ```rust
29/// use kernel::alloc::KBox;
30/// use kernel::xarray::{AllocKind, XArray};
31///
32/// let xa = KBox::pin_init(XArray::new(AllocKind::Alloc1), GFP_KERNEL)?;
33///
34/// let dead = KBox::new(0xdead, GFP_KERNEL)?;
35/// let beef = KBox::new(0xbeef, GFP_KERNEL)?;
36///
37/// let mut guard = xa.lock();
38///
39/// assert_eq!(guard.get(0), None);
40///
41/// assert_eq!(guard.store(0, dead, GFP_KERNEL)?.as_deref(), None);
42/// assert_eq!(guard.get(0).copied(), Some(0xdead));
43///
44/// *guard.get_mut(0).unwrap() = 0xffff;
45/// assert_eq!(guard.get(0).copied(), Some(0xffff));
46///
47/// assert_eq!(guard.store(0, beef, GFP_KERNEL)?.as_deref().copied(), Some(0xffff));
48/// assert_eq!(guard.get(0).copied(), Some(0xbeef));
49///
50/// guard.remove(0);
51/// assert_eq!(guard.get(0), None);
52///
53/// # Ok::<(), Error>(())
54/// ```
55#[pin_data(PinnedDrop)]
56pub struct XArray<T: ForeignOwnable> {
57    #[pin]
58    xa: Opaque<bindings::xarray>,
59    _p: PhantomData<T>,
60}
61
62#[pinned_drop]
63impl<T: ForeignOwnable> PinnedDrop for XArray<T> {
64    fn drop(self: Pin<&mut Self>) {
65        self.iter().for_each(|ptr| {
66            let ptr = ptr.as_ptr();
67            // SAFETY: `ptr` came from `T::into_foreign`.
68            //
69            // INVARIANT: we own the only reference to the array which is being dropped so the
70            // broken invariant is not observable on function exit.
71            drop(unsafe { T::from_foreign(ptr) })
72        });
73
74        // SAFETY: `self.xa` is always valid by the type invariant.
75        unsafe { bindings::xa_destroy(self.xa.get()) };
76    }
77}
78
79/// Flags passed to [`XArray::new`] to configure the array's allocation tracking behavior.
80pub enum AllocKind {
81    /// Consider the first element to be at index 0.
82    Alloc,
83    /// Consider the first element to be at index 1.
84    Alloc1,
85}
86
87impl<T: ForeignOwnable> XArray<T> {
88    /// Creates a new initializer for this type.
89    pub fn new(kind: AllocKind) -> impl PinInit<Self> {
90        let flags = match kind {
91            AllocKind::Alloc => bindings::XA_FLAGS_ALLOC,
92            AllocKind::Alloc1 => bindings::XA_FLAGS_ALLOC1,
93        };
94        pin_init!(Self {
95            // SAFETY: `xa` is valid while the closure is called.
96            //
97            // INVARIANT: `xa` is initialized here to an empty, valid [`bindings::xarray`].
98            xa <- Opaque::ffi_init(|xa| unsafe {
99                bindings::xa_init_flags(xa, flags)
100            }),
101            _p: PhantomData,
102        })
103    }
104
105    fn iter(&self) -> impl Iterator<Item = NonNull<c_void>> + '_ {
106        let mut index = 0;
107
108        // SAFETY: `self.xa` is always valid by the type invariant.
109        iter::once(unsafe {
110            bindings::xa_find(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
111        })
112        .chain(iter::from_fn(move || {
113            // SAFETY: `self.xa` is always valid by the type invariant.
114            Some(unsafe {
115                bindings::xa_find_after(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
116            })
117        }))
118        .map_while(|ptr| NonNull::new(ptr.cast()))
119    }
120
121    /// Attempts to lock the [`XArray`] for exclusive access.
122    pub fn try_lock(&self) -> Option<Guard<'_, T>> {
123        // SAFETY: `self.xa` is always valid by the type invariant.
124        if (unsafe { bindings::xa_trylock(self.xa.get()) } != 0) {
125            Some(Guard {
126                xa: self,
127                _not_send: NotThreadSafe,
128            })
129        } else {
130            None
131        }
132    }
133
134    /// Locks the [`XArray`] for exclusive access.
135    pub fn lock(&self) -> Guard<'_, T> {
136        // SAFETY: `self.xa` is always valid by the type invariant.
137        unsafe { bindings::xa_lock(self.xa.get()) };
138
139        Guard {
140            xa: self,
141            _not_send: NotThreadSafe,
142        }
143    }
144}
145
146/// A lock guard.
147///
148/// The lock is unlocked when the guard goes out of scope.
149#[must_use = "the lock unlocks immediately when the guard is unused"]
150pub struct Guard<'a, T: ForeignOwnable> {
151    xa: &'a XArray<T>,
152    _not_send: NotThreadSafe,
153}
154
155impl<T: ForeignOwnable> Drop for Guard<'_, T> {
156    fn drop(&mut self) {
157        // SAFETY:
158        // - `self.xa.xa` is always valid by the type invariant.
159        // - The caller holds the lock, so it is safe to unlock it.
160        unsafe { bindings::xa_unlock(self.xa.xa.get()) };
161    }
162}
163
164/// The error returned by [`store`](Guard::store).
165///
166/// Contains the underlying error and the value that was not stored.
167pub struct StoreError<T> {
168    /// The error that occurred.
169    pub error: Error,
170    /// The value that was not stored.
171    pub value: T,
172}
173
174impl<T> From<StoreError<T>> for Error {
175    fn from(value: StoreError<T>) -> Self {
176        value.error
177    }
178}
179
180impl<'a, T: ForeignOwnable> Guard<'a, T> {
181    fn load<F, U>(&self, index: usize, f: F) -> Option<U>
182    where
183        F: FnOnce(NonNull<c_void>) -> U,
184    {
185        // SAFETY: `self.xa.xa` is always valid by the type invariant.
186        let ptr = unsafe { bindings::xa_load(self.xa.xa.get(), index) };
187        let ptr = NonNull::new(ptr.cast())?;
188        Some(f(ptr))
189    }
190
191    /// Provides a reference to the element at the given index.
192    pub fn get(&self, index: usize) -> Option<T::Borrowed<'_>> {
193        self.load(index, |ptr| {
194            // SAFETY: `ptr` came from `T::into_foreign`.
195            unsafe { T::borrow(ptr.as_ptr()) }
196        })
197    }
198
199    /// Provides a mutable reference to the element at the given index.
200    pub fn get_mut(&mut self, index: usize) -> Option<T::BorrowedMut<'_>> {
201        self.load(index, |ptr| {
202            // SAFETY: `ptr` came from `T::into_foreign`.
203            unsafe { T::borrow_mut(ptr.as_ptr()) }
204        })
205    }
206
207    /// Removes and returns the element at the given index.
208    pub fn remove(&mut self, index: usize) -> Option<T> {
209        // SAFETY:
210        // - `self.xa.xa` is always valid by the type invariant.
211        // - The caller holds the lock.
212        let ptr = unsafe { bindings::__xa_erase(self.xa.xa.get(), index) }.cast();
213        // SAFETY:
214        // - `ptr` is either NULL or came from `T::into_foreign`.
215        // - `&mut self` guarantees that the lifetimes of [`T::Borrowed`] and [`T::BorrowedMut`]
216        // borrowed from `self` have ended.
217        unsafe { T::try_from_foreign(ptr) }
218    }
219
220    /// Stores an element at the given index.
221    ///
222    /// May drop the lock if needed to allocate memory, and then reacquire it afterwards.
223    ///
224    /// On success, returns the element which was previously at the given index.
225    ///
226    /// On failure, returns the element which was attempted to be stored.
227    pub fn store(
228        &mut self,
229        index: usize,
230        value: T,
231        gfp: alloc::Flags,
232    ) -> Result<Option<T>, StoreError<T>> {
233        build_assert!(
234            T::FOREIGN_ALIGN >= 4,
235            "pointers stored in XArray must be 4-byte aligned"
236        );
237        let new = value.into_foreign();
238
239        let old = {
240            let new = new.cast();
241            // SAFETY:
242            // - `self.xa.xa` is always valid by the type invariant.
243            // - The caller holds the lock.
244            //
245            // INVARIANT: `new` came from `T::into_foreign`.
246            unsafe { bindings::__xa_store(self.xa.xa.get(), index, new, gfp.as_raw()) }
247        };
248
249        // SAFETY: `__xa_store` returns the old entry at this index on success or `xa_err` if an
250        // error happened.
251        let errno = unsafe { bindings::xa_err(old) };
252        if errno != 0 {
253            // SAFETY: `new` came from `T::into_foreign` and `__xa_store` does not take
254            // ownership of the value on error.
255            let value = unsafe { T::from_foreign(new) };
256            Err(StoreError {
257                value,
258                error: Error::from_errno(errno),
259            })
260        } else {
261            let old = old.cast();
262            // SAFETY: `ptr` is either NULL or came from `T::into_foreign`.
263            //
264            // NB: `XA_ZERO_ENTRY` is never returned by functions belonging to the Normal XArray
265            // API; such entries present as `NULL`.
266            Ok(unsafe { T::try_from_foreign(old) })
267        }
268    }
269}
270
271// SAFETY: `XArray<T>` has no shared mutable state so it is `Send` iff `T` is `Send`.
272unsafe impl<T: ForeignOwnable + Send> Send for XArray<T> {}
273
274// SAFETY: `XArray<T>` serialises the interior mutability it provides so it is `Sync` iff `T` is
275// `Send`.
276unsafe impl<T: ForeignOwnable + Send> Sync for XArray<T> {}