kernel/alloc/
arrayvec.rs

1// SPDX-License-Identifier: GPL-2.0
2
3//! Implementation of [`ArrayVec`].
4
5use crate::{
6    alloc::kvec::{
7        impl_slice_eq,
8        PushError, //
9    },
10    const_assert,
11    error::{
12        code::EINVAL,
13        Error,
14        Result, //
15    },
16    fmt, //
17};
18
19use core::{
20    borrow::{
21        Borrow,
22        BorrowMut, //
23    },
24    mem::MaybeUninit,
25    ops::{
26        Deref,
27        DerefMut, //
28    },
29    ptr,
30    slice, //
31};
32
33use pin_init::{
34    init_from_closure,
35    Init, //
36};
37
38/// A fixed capacity vector that holds at most `N` elements.
39///
40/// # Invariants
41///
42/// - `len` is at most `N`.
43/// - The first `len` elements of `data` are initialized.
44///
45/// # Examples
46///
47/// ```
48/// use kernel::alloc::ArrayVec;
49///
50/// let mut v = ArrayVec::<u8, 4>::new();
51/// v.extend_from_slice(b"abc")?;
52/// assert_eq!(*v, *b"abc");
53///
54/// assert!(v.extend_from_slice(b"ab").is_err());
55///
56/// v.push(4u8)?;
57/// assert_eq!(*v, *b"abc\x04");
58/// assert!(v.push(5u8).is_err());
59///
60/// v.clear();
61/// assert!(v.is_empty());
62/// # Ok::<(), Error>(())
63/// ```
64pub struct ArrayVec<T, const N: usize> {
65    data: [MaybeUninit<T>; N],
66    len: usize,
67}
68
69/// The maximum size in bytes of an [`ArrayVec`] returned by value from [`ArrayVec::new`].
70const MAX_STACK_SIZE: usize = 512;
71
72impl<T, const N: usize> ArrayVec<T, N> {
73    /// Creates an empty [`ArrayVec`].
74    #[inline]
75    pub const fn new() -> Self {
76        const_assert!(
77            size_of::<Self>() <= MAX_STACK_SIZE,
78            "use `init_empty` instead of constructing a large ArrayVec on the stack"
79        );
80
81        // INVARIANT: An empty ArrayVec trivially has all its elements initialized.
82        Self {
83            data: [const { MaybeUninit::uninit() }; N],
84            len: 0,
85        }
86    }
87
88    /// Creates an initializer for an empty [`ArrayVec`].
89    ///
90    /// # Examples
91    ///
92    /// ```
93    /// use kernel::alloc::ArrayVec;
94    ///
95    /// let v = KBox::init(ArrayVec::<u8, 4096>::init_empty(), GFP_KERNEL)?;
96    /// assert!(v.is_empty());
97    /// # Ok::<(), Error>(())
98    /// ```
99    #[inline]
100    pub fn init_empty() -> impl Init<Self> {
101        Self::init_with(|_| Ok(()))
102    }
103
104    /// Creates an initializer for an [`ArrayVec`] populated by `f`.
105    ///
106    /// `f` gets an empty [`ArrayVec`] and can fill it in place.
107    ///
108    /// # Examples
109    ///
110    /// ```
111    /// use kernel::alloc::ArrayVec;
112    ///
113    /// let v = KBox::init(
114    ///     ArrayVec::<u8, 4096>::init_with(|v| v.extend_from_slice(b"abc")),
115    ///     GFP_KERNEL,
116    /// )?;
117    /// assert_eq!(**v, *b"abc");
118    /// # Ok::<(), Error>(())
119    /// ```
120    #[inline]
121    pub fn init_with<E>(f: impl FnOnce(&mut Self) -> Result<(), E>) -> impl Init<Self, E> {
122        let init_empty = |slot: *mut Self| {
123            // SAFETY: By the initializer contract `slot` is valid for writes. Once `len` is zero
124            // the slot holds a valid empty ArrayVec, since `data` requires no initialization.
125            // INVARIANT: An empty ArrayVec trivially has all its elements initialized.
126            unsafe { (&raw mut (*slot).len).write(0) };
127            Ok(())
128        };
129
130        // SAFETY: `init_empty` fully initializes the slot and never fails.
131        unsafe { init_from_closure(init_empty) }.chain(f)
132    }
133
134    /// Appends an element to the back of the [`ArrayVec`].
135    ///
136    /// Fails when the [`ArrayVec`] is full, handing the element back in [`PushError`].
137    #[inline]
138    pub fn push(&mut self, v: T) -> Result<(), PushError<T>> {
139        self.push_init(v)
140    }
141
142    /// Appends an element to the back of the [`ArrayVec`] by initializing it in place.
143    ///
144    /// Fails when the [`ArrayVec`] is full, handing the initializer back in [`PushError`].
145    #[inline]
146    pub fn push_init<I: Init<T>>(&mut self, init: I) -> Result<(), PushError<I>> {
147        self.try_push_init(init)
148            .map_err(|TryPushInitError::Full(init)| PushError(init))
149    }
150
151    /// Appends an element to the back of the [`ArrayVec`] by initializing it in place.
152    ///
153    /// Unlike [`ArrayVec::push_init`], the initializer may be fallible. If the [`ArrayVec`] is
154    /// full, the original initializer `init` is handed back in [`TryPushInitError::Full`]. If the
155    /// initializer itself fails, its error is returned in [`TryPushInitError::InitError`].
156    pub fn try_push_init<I, E>(&mut self, init: I) -> Result<(), TryPushInitError<I, E>>
157    where
158        I: Init<T, E>,
159    {
160        let Some(slot) = self.spare_capacity_mut().first_mut() else {
161            return Err(TryPushInitError::Full(init));
162        };
163
164        // SAFETY:
165        // - `slot.as_mut_ptr()` points to a spare slot, which is valid, uninitialized memory.
166        // - If `raw_try_init` fails, this returns without increasing `len` or touching `slot`.
167        // - `init` implements `Init<T, E>`, which removes the pinning requirement.
168        unsafe { pin_init::raw_try_init(slot.as_mut_ptr(), init) }
169            .map_err(TryPushInitError::InitError)?;
170
171        // INVARIANT: The element at index `len` was just initialized, and the new `len` does not
172        // exceed `N` because a spare slot existed.
173        self.len += 1;
174
175        Ok(())
176    }
177
178    /// Appends a clone of each element in `slice` to the back of the [`ArrayVec`].
179    ///
180    /// Fails with [`EINVAL`] if `slice` is longer than the remaining capacity.
181    pub fn extend_from_slice(&mut self, slice: &[T]) -> Result
182    where
183        T: Clone,
184    {
185        let Some(dst) = self.spare_capacity_mut().get_mut(..slice.len()) else {
186            return Err(EINVAL);
187        };
188
189        for (d, s) in dst.iter_mut().zip(slice) {
190            d.write(s.clone());
191        }
192        // INVARIANT: The next `slice.len()` elements after `len` were just initialized, and the
193        // new `len` does not exceed `N` because the spare capacity was enough.
194        self.len += slice.len();
195
196        Ok(())
197    }
198
199    /// Removes all elements.
200    #[inline]
201    pub fn clear(&mut self) {
202        let elems: *mut [T] = self.as_mut_slice();
203        // INVARIANT: An empty ArrayVec trivially has all its elements initialized.
204        self.len = 0;
205        // SAFETY: There are no references to the elements since we hold `&mut self`. The elements
206        // can't be dropped again because `len` is already 0.
207        unsafe { ptr::drop_in_place(elems) };
208    }
209
210    /// Returns the initialized elements as a slice.
211    #[inline]
212    pub fn as_slice(&self) -> &[T] {
213        let ptr = self.data.as_ptr().cast::<T>();
214        // SAFETY: `MaybeUninit<T>` has the same layout as `T`, and by the type invariants the first
215        // `len` elements of `data` are initialized.
216        unsafe { slice::from_raw_parts(ptr, self.len) }
217    }
218
219    /// Returns the initialized elements as a mutable slice.
220    #[inline]
221    pub fn as_mut_slice(&mut self) -> &mut [T] {
222        let ptr = self.data.as_mut_ptr().cast::<T>();
223        // SAFETY: `MaybeUninit<T>` has the same layout as `T`, and by the type invariants the first
224        // `len` elements of `data` are initialized.
225        unsafe { slice::from_raw_parts_mut(ptr, self.len) }
226    }
227
228    /// Returns a slice of `MaybeUninit<T>` for the remaining spare capacity of the [`ArrayVec`].
229    #[inline]
230    fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
231        // PANIC: `len` never exceeds `N` by the type invariants.
232        &mut self.data[self.len..]
233    }
234}
235
236/// Error type for [`ArrayVec::try_push_init`].
237pub enum TryPushInitError<I, E> {
238    /// The [`ArrayVec`] is full. Hand the initializer back.
239    Full(I),
240    /// The initializer failed.
241    InitError(E),
242}
243
244impl<I, E> fmt::Debug for TryPushInitError<I, E> {
245    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
246        match self {
247            TryPushInitError::Full(_) => write!(f, "Not enough capacity"),
248            TryPushInitError::InitError(_) => write!(f, "Initializer failed"),
249        }
250    }
251}
252
253impl<I, E> From<TryPushInitError<I, E>> for Error
254where
255    Error: From<E>,
256{
257    #[inline]
258    fn from(e: TryPushInitError<I, E>) -> Error {
259        match e {
260            TryPushInitError::Full(_) => EINVAL,
261            TryPushInitError::InitError(e) => Error::from(e),
262        }
263    }
264}
265
266impl<T, const N: usize> Default for ArrayVec<T, N> {
267    #[inline]
268    fn default() -> Self {
269        Self::new()
270    }
271}
272
273impl<T, const N: usize> Drop for ArrayVec<T, N> {
274    #[inline]
275    fn drop(&mut self) {
276        // SAFETY: The slice holds initialized elements that are never accessed again after this
277        // point.
278        unsafe { ptr::drop_in_place(self.as_mut_slice()) };
279    }
280}
281
282impl<T, const N: usize> Deref for ArrayVec<T, N> {
283    type Target = [T];
284
285    #[inline]
286    fn deref(&self) -> &Self::Target {
287        self.as_slice()
288    }
289}
290
291impl<T, const N: usize> DerefMut for ArrayVec<T, N> {
292    #[inline]
293    fn deref_mut(&mut self) -> &mut Self::Target {
294        self.as_mut_slice()
295    }
296}
297
298impl<T, const N: usize> Borrow<[T]> for ArrayVec<T, N> {
299    #[inline]
300    fn borrow(&self) -> &[T] {
301        self.as_slice()
302    }
303}
304
305impl<T, const N: usize> BorrowMut<[T]> for ArrayVec<T, N> {
306    #[inline]
307    fn borrow_mut(&mut self) -> &mut [T] {
308        self.as_mut_slice()
309    }
310}
311
312impl<T: Eq, const N: usize> Eq for ArrayVec<T, N> {}
313
314impl_slice_eq! {
315    [const N: usize, const M: usize] ArrayVec<T, N>, ArrayVec<U, M>,
316    [const N: usize] ArrayVec<T, N>, &[U],
317    [const N: usize] ArrayVec<T, N>, &mut [U],
318    [const N: usize] &[T], ArrayVec<U, N>,
319    [const N: usize] &mut [T], ArrayVec<U, N>,
320    [const N: usize] ArrayVec<T, N>, [U],
321    [const N: usize] [T], ArrayVec<U, N>,
322    [const N: usize, const M: usize] ArrayVec<T, N>, [U; M],
323    [const N: usize, const M: usize] ArrayVec<T, N>, &[U; M],
324}
325
326impl<'a, T, const N: usize> IntoIterator for &'a ArrayVec<T, N> {
327    type Item = &'a T;
328    type IntoIter = slice::Iter<'a, T>;
329
330    #[inline]
331    fn into_iter(self) -> Self::IntoIter {
332        self.iter()
333    }
334}
335
336impl<'a, T, const N: usize> IntoIterator for &'a mut ArrayVec<T, N> {
337    type Item = &'a mut T;
338    type IntoIter = slice::IterMut<'a, T>;
339
340    #[inline]
341    fn into_iter(self) -> Self::IntoIter {
342        self.iter_mut()
343    }
344}
345
346impl<T: fmt::Debug, const N: usize> fmt::Debug for ArrayVec<T, N> {
347    #[inline]
348    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
349        fmt::Debug::fmt(self.as_slice(), f)
350    }
351}
kernel/alloc/arrayvec.rs

kernel/alloc/
arrayvec.rs
