kvec.rs 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919
  1. // SPDX-License-Identifier: GPL-2.0
  2. //! Implementation of [`Vec`].
  3. // May not be needed in Rust 1.87.0 (pending beta backport).
  4. #![allow(clippy::ptr_eq)]
  5. use super::{
  6. allocator::{KVmalloc, Kmalloc, Vmalloc},
  7. layout::ArrayLayout,
  8. AllocError, Allocator, Box, Flags,
  9. };
  10. use core::{
  11. fmt,
  12. marker::PhantomData,
  13. mem::{ManuallyDrop, MaybeUninit},
  14. ops::Deref,
  15. ops::DerefMut,
  16. ops::Index,
  17. ops::IndexMut,
  18. ptr,
  19. ptr::NonNull,
  20. slice,
  21. slice::SliceIndex,
  22. };
  23. /// Create a [`KVec`] containing the arguments.
  24. ///
  25. /// New memory is allocated with `GFP_KERNEL`.
  26. ///
  27. /// # Examples
  28. ///
  29. /// ```
  30. /// let mut v = kernel::kvec![];
  31. /// v.push(1, GFP_KERNEL)?;
  32. /// assert_eq!(v, [1]);
  33. ///
  34. /// let mut v = kernel::kvec![1; 3]?;
  35. /// v.push(4, GFP_KERNEL)?;
  36. /// assert_eq!(v, [1, 1, 1, 4]);
  37. ///
  38. /// let mut v = kernel::kvec![1, 2, 3]?;
  39. /// v.push(4, GFP_KERNEL)?;
  40. /// assert_eq!(v, [1, 2, 3, 4]);
  41. ///
  42. /// # Ok::<(), Error>(())
  43. /// ```
  44. #[macro_export]
  45. macro_rules! kvec {
  46. () => (
  47. $crate::alloc::KVec::new()
  48. );
  49. ($elem:expr; $n:expr) => (
  50. $crate::alloc::KVec::from_elem($elem, $n, GFP_KERNEL)
  51. );
  52. ($($x:expr),+ $(,)?) => (
  53. match $crate::alloc::KBox::new_uninit(GFP_KERNEL) {
  54. Ok(b) => Ok($crate::alloc::KVec::from($crate::alloc::KBox::write(b, [$($x),+]))),
  55. Err(e) => Err(e),
  56. }
  57. );
  58. }
  59. /// The kernel's [`Vec`] type.
  60. ///
  61. /// A contiguous growable array type with contents allocated with the kernel's allocators (e.g.
  62. /// [`Kmalloc`], [`Vmalloc`] or [`KVmalloc`]), written `Vec<T, A>`.
  63. ///
  64. /// For non-zero-sized values, a [`Vec`] will use the given allocator `A` for its allocation. For
  65. /// the most common allocators the type aliases [`KVec`], [`VVec`] and [`KVVec`] exist.
  66. ///
  67. /// For zero-sized types the [`Vec`]'s pointer must be `dangling_mut::<T>`; no memory is allocated.
  68. ///
  69. /// Generally, [`Vec`] consists of a pointer that represents the vector's backing buffer, the
  70. /// capacity of the vector (the number of elements that currently fit into the vector), its length
  71. /// (the number of elements that are currently stored in the vector) and the `Allocator` type used
  72. /// to allocate (and free) the backing buffer.
  73. ///
  74. /// A [`Vec`] can be deconstructed into and (re-)constructed from its previously named raw parts
  75. /// and manually modified.
  76. ///
  77. /// [`Vec`]'s backing buffer gets, if required, automatically increased (re-allocated) when elements
  78. /// are added to the vector.
  79. ///
  80. /// # Invariants
  81. ///
  82. /// - `self.ptr` is always properly aligned and either points to memory allocated with `A` or, for
  83. /// zero-sized types, is a dangling, well aligned pointer.
  84. ///
  85. /// - `self.len` always represents the exact number of elements stored in the vector.
  86. ///
  87. /// - `self.layout` represents the absolute number of elements that can be stored within the vector
  88. /// without re-allocation. For ZSTs `self.layout`'s capacity is zero. However, it is legal for the
  89. /// backing buffer to be larger than `layout`.
  90. ///
  91. /// - The `Allocator` type `A` of the vector is the exact same `Allocator` type the backing buffer
  92. /// was allocated with (and must be freed with).
  93. pub struct Vec<T, A: Allocator> {
  94. ptr: NonNull<T>,
  95. /// Represents the actual buffer size as `cap` times `size_of::<T>` bytes.
  96. ///
  97. /// Note: This isn't quite the same as `Self::capacity`, which in contrast returns the number of
  98. /// elements we can still store without reallocating.
  99. layout: ArrayLayout<T>,
  100. len: usize,
  101. _p: PhantomData<A>,
  102. }
  103. /// Type alias for [`Vec`] with a [`Kmalloc`] allocator.
  104. ///
  105. /// # Examples
  106. ///
  107. /// ```
  108. /// let mut v = KVec::new();
  109. /// v.push(1, GFP_KERNEL)?;
  110. /// assert_eq!(&v, &[1]);
  111. ///
  112. /// # Ok::<(), Error>(())
  113. /// ```
  114. pub type KVec<T> = Vec<T, Kmalloc>;
  115. /// Type alias for [`Vec`] with a [`Vmalloc`] allocator.
  116. ///
  117. /// # Examples
  118. ///
  119. /// ```
  120. /// let mut v = VVec::new();
  121. /// v.push(1, GFP_KERNEL)?;
  122. /// assert_eq!(&v, &[1]);
  123. ///
  124. /// # Ok::<(), Error>(())
  125. /// ```
  126. pub type VVec<T> = Vec<T, Vmalloc>;
  127. /// Type alias for [`Vec`] with a [`KVmalloc`] allocator.
  128. ///
  129. /// # Examples
  130. ///
  131. /// ```
  132. /// let mut v = KVVec::new();
  133. /// v.push(1, GFP_KERNEL)?;
  134. /// assert_eq!(&v, &[1]);
  135. ///
  136. /// # Ok::<(), Error>(())
  137. /// ```
  138. pub type KVVec<T> = Vec<T, KVmalloc>;
  139. // SAFETY: `Vec` is `Send` if `T` is `Send` because `Vec` owns its elements.
  140. unsafe impl<T, A> Send for Vec<T, A>
  141. where
  142. T: Send,
  143. A: Allocator,
  144. {
  145. }
  146. // SAFETY: `Vec` is `Sync` if `T` is `Sync` because `Vec` owns its elements.
  147. unsafe impl<T, A> Sync for Vec<T, A>
  148. where
  149. T: Sync,
  150. A: Allocator,
  151. {
  152. }
  153. impl<T, A> Vec<T, A>
  154. where
  155. A: Allocator,
  156. {
  157. #[inline]
  158. const fn is_zst() -> bool {
  159. core::mem::size_of::<T>() == 0
  160. }
  161. /// Returns the number of elements that can be stored within the vector without allocating
  162. /// additional memory.
  163. pub fn capacity(&self) -> usize {
  164. if const { Self::is_zst() } {
  165. usize::MAX
  166. } else {
  167. self.layout.len()
  168. }
  169. }
  170. /// Returns the number of elements stored within the vector.
  171. #[inline]
  172. pub fn len(&self) -> usize {
  173. self.len
  174. }
  175. /// Forcefully sets `self.len` to `new_len`.
  176. ///
  177. /// # Safety
  178. ///
  179. /// - `new_len` must be less than or equal to [`Self::capacity`].
  180. /// - If `new_len` is greater than `self.len`, all elements within the interval
  181. /// [`self.len`,`new_len`) must be initialized.
  182. #[inline]
  183. pub unsafe fn set_len(&mut self, new_len: usize) {
  184. debug_assert!(new_len <= self.capacity());
  185. // INVARIANT: By the safety requirements of this method `new_len` represents the exact
  186. // number of elements stored within `self`.
  187. self.len = new_len;
  188. }
  189. /// Returns a slice of the entire vector.
  190. #[inline]
  191. pub fn as_slice(&self) -> &[T] {
  192. self
  193. }
  194. /// Returns a mutable slice of the entire vector.
  195. #[inline]
  196. pub fn as_mut_slice(&mut self) -> &mut [T] {
  197. self
  198. }
  199. /// Returns a mutable raw pointer to the vector's backing buffer, or, if `T` is a ZST, a
  200. /// dangling raw pointer.
  201. #[inline]
  202. pub fn as_mut_ptr(&mut self) -> *mut T {
  203. self.ptr.as_ptr()
  204. }
  205. /// Returns a raw pointer to the vector's backing buffer, or, if `T` is a ZST, a dangling raw
  206. /// pointer.
  207. #[inline]
  208. pub fn as_ptr(&self) -> *const T {
  209. self.ptr.as_ptr()
  210. }
  211. /// Returns `true` if the vector contains no elements, `false` otherwise.
  212. ///
  213. /// # Examples
  214. ///
  215. /// ```
  216. /// let mut v = KVec::new();
  217. /// assert!(v.is_empty());
  218. ///
  219. /// v.push(1, GFP_KERNEL);
  220. /// assert!(!v.is_empty());
  221. /// ```
  222. #[inline]
  223. pub fn is_empty(&self) -> bool {
  224. self.len() == 0
  225. }
  226. /// Creates a new, empty `Vec<T, A>`.
  227. ///
  228. /// This method does not allocate by itself.
  229. #[inline]
  230. pub const fn new() -> Self {
  231. // INVARIANT: Since this is a new, empty `Vec` with no backing memory yet,
  232. // - `ptr` is a properly aligned dangling pointer for type `T`,
  233. // - `layout` is an empty `ArrayLayout` (zero capacity)
  234. // - `len` is zero, since no elements can be or have been stored,
  235. // - `A` is always valid.
  236. Self {
  237. ptr: NonNull::dangling(),
  238. layout: ArrayLayout::empty(),
  239. len: 0,
  240. _p: PhantomData::<A>,
  241. }
  242. }
  243. /// Returns a slice of `MaybeUninit<T>` for the remaining spare capacity of the vector.
  244. pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
  245. // SAFETY:
  246. // - `self.len` is smaller than `self.capacity` and hence, the resulting pointer is
  247. // guaranteed to be part of the same allocated object.
  248. // - `self.len` can not overflow `isize`.
  249. let ptr = unsafe { self.as_mut_ptr().add(self.len) } as *mut MaybeUninit<T>;
  250. // SAFETY: The memory between `self.len` and `self.capacity` is guaranteed to be allocated
  251. // and valid, but uninitialized.
  252. unsafe { slice::from_raw_parts_mut(ptr, self.capacity() - self.len) }
  253. }
  254. /// Appends an element to the back of the [`Vec`] instance.
  255. ///
  256. /// # Examples
  257. ///
  258. /// ```
  259. /// let mut v = KVec::new();
  260. /// v.push(1, GFP_KERNEL)?;
  261. /// assert_eq!(&v, &[1]);
  262. ///
  263. /// v.push(2, GFP_KERNEL)?;
  264. /// assert_eq!(&v, &[1, 2]);
  265. /// # Ok::<(), Error>(())
  266. /// ```
  267. pub fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
  268. self.reserve(1, flags)?;
  269. // SAFETY:
  270. // - `self.len` is smaller than `self.capacity` and hence, the resulting pointer is
  271. // guaranteed to be part of the same allocated object.
  272. // - `self.len` can not overflow `isize`.
  273. let ptr = unsafe { self.as_mut_ptr().add(self.len) };
  274. // SAFETY:
  275. // - `ptr` is properly aligned and valid for writes.
  276. unsafe { core::ptr::write(ptr, v) };
  277. // SAFETY: We just initialised the first spare entry, so it is safe to increase the length
  278. // by 1. We also know that the new length is <= capacity because of the previous call to
  279. // `reserve` above.
  280. unsafe { self.set_len(self.len() + 1) };
  281. Ok(())
  282. }
  283. /// Creates a new [`Vec`] instance with at least the given capacity.
  284. ///
  285. /// # Examples
  286. ///
  287. /// ```
  288. /// let v = KVec::<u32>::with_capacity(20, GFP_KERNEL)?;
  289. ///
  290. /// assert!(v.capacity() >= 20);
  291. /// # Ok::<(), Error>(())
  292. /// ```
  293. pub fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError> {
  294. let mut v = Vec::new();
  295. v.reserve(capacity, flags)?;
  296. Ok(v)
  297. }
  298. /// Creates a `Vec<T, A>` from a pointer, a length and a capacity using the allocator `A`.
  299. ///
  300. /// # Examples
  301. ///
  302. /// ```
  303. /// let mut v = kernel::kvec![1, 2, 3]?;
  304. /// v.reserve(1, GFP_KERNEL)?;
  305. ///
  306. /// let (mut ptr, mut len, cap) = v.into_raw_parts();
  307. ///
  308. /// // SAFETY: We've just reserved memory for another element.
  309. /// unsafe { ptr.add(len).write(4) };
  310. /// len += 1;
  311. ///
  312. /// // SAFETY: We only wrote an additional element at the end of the `KVec`'s buffer and
  313. /// // correspondingly increased the length of the `KVec` by one. Otherwise, we construct it
  314. /// // from the exact same raw parts.
  315. /// let v = unsafe { KVec::from_raw_parts(ptr, len, cap) };
  316. ///
  317. /// assert_eq!(v, [1, 2, 3, 4]);
  318. ///
  319. /// # Ok::<(), Error>(())
  320. /// ```
  321. ///
  322. /// # Safety
  323. ///
  324. /// If `T` is a ZST:
  325. ///
  326. /// - `ptr` must be a dangling, well aligned pointer.
  327. ///
  328. /// Otherwise:
  329. ///
  330. /// - `ptr` must have been allocated with the allocator `A`.
  331. /// - `ptr` must satisfy or exceed the alignment requirements of `T`.
  332. /// - `ptr` must point to memory with a size of at least `size_of::<T>() * capacity` bytes.
  333. /// - The allocated size in bytes must not be larger than `isize::MAX`.
  334. /// - `length` must be less than or equal to `capacity`.
  335. /// - The first `length` elements must be initialized values of type `T`.
  336. ///
  337. /// It is also valid to create an empty `Vec` passing a dangling pointer for `ptr` and zero for
  338. /// `cap` and `len`.
  339. pub unsafe fn from_raw_parts(ptr: *mut T, length: usize, capacity: usize) -> Self {
  340. let layout = if Self::is_zst() {
  341. ArrayLayout::empty()
  342. } else {
  343. // SAFETY: By the safety requirements of this function, `capacity * size_of::<T>()` is
  344. // smaller than `isize::MAX`.
  345. unsafe { ArrayLayout::new_unchecked(capacity) }
  346. };
  347. // INVARIANT: For ZSTs, we store an empty `ArrayLayout`, all other type invariants are
  348. // covered by the safety requirements of this function.
  349. Self {
  350. // SAFETY: By the safety requirements, `ptr` is either dangling or pointing to a valid
  351. // memory allocation, allocated with `A`.
  352. ptr: unsafe { NonNull::new_unchecked(ptr) },
  353. layout,
  354. len: length,
  355. _p: PhantomData::<A>,
  356. }
  357. }
  358. /// Consumes the `Vec<T, A>` and returns its raw components `pointer`, `length` and `capacity`.
  359. ///
  360. /// This will not run the destructor of the contained elements and for non-ZSTs the allocation
  361. /// will stay alive indefinitely. Use [`Vec::from_raw_parts`] to recover the [`Vec`], drop the
  362. /// elements and free the allocation, if any.
  363. pub fn into_raw_parts(self) -> (*mut T, usize, usize) {
  364. let mut me = ManuallyDrop::new(self);
  365. let len = me.len();
  366. let capacity = me.capacity();
  367. let ptr = me.as_mut_ptr();
  368. (ptr, len, capacity)
  369. }
  370. /// Ensures that the capacity exceeds the length by at least `additional` elements.
  371. ///
  372. /// # Examples
  373. ///
  374. /// ```
  375. /// let mut v = KVec::new();
  376. /// v.push(1, GFP_KERNEL)?;
  377. ///
  378. /// v.reserve(10, GFP_KERNEL)?;
  379. /// let cap = v.capacity();
  380. /// assert!(cap >= 10);
  381. ///
  382. /// v.reserve(10, GFP_KERNEL)?;
  383. /// let new_cap = v.capacity();
  384. /// assert_eq!(new_cap, cap);
  385. ///
  386. /// # Ok::<(), Error>(())
  387. /// ```
  388. pub fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError> {
  389. let len = self.len();
  390. let cap = self.capacity();
  391. if cap - len >= additional {
  392. return Ok(());
  393. }
  394. if Self::is_zst() {
  395. // The capacity is already `usize::MAX` for ZSTs, we can't go higher.
  396. return Err(AllocError);
  397. }
  398. // We know that `cap <= isize::MAX` because of the type invariants of `Self`. So the
  399. // multiplication by two won't overflow.
  400. let new_cap = core::cmp::max(cap * 2, len.checked_add(additional).ok_or(AllocError)?);
  401. let layout = ArrayLayout::new(new_cap).map_err(|_| AllocError)?;
  402. // SAFETY:
  403. // - `ptr` is valid because it's either `None` or comes from a previous call to
  404. // `A::realloc`.
  405. // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
  406. let ptr = unsafe {
  407. A::realloc(
  408. Some(self.ptr.cast()),
  409. layout.into(),
  410. self.layout.into(),
  411. flags,
  412. )?
  413. };
  414. // INVARIANT:
  415. // - `layout` is some `ArrayLayout::<T>`,
  416. // - `ptr` has been created by `A::realloc` from `layout`.
  417. self.ptr = ptr.cast();
  418. self.layout = layout;
  419. Ok(())
  420. }
  421. }
  422. impl<T: Clone, A: Allocator> Vec<T, A> {
  423. /// Extend the vector by `n` clones of `value`.
  424. pub fn extend_with(&mut self, n: usize, value: T, flags: Flags) -> Result<(), AllocError> {
  425. if n == 0 {
  426. return Ok(());
  427. }
  428. self.reserve(n, flags)?;
  429. let spare = self.spare_capacity_mut();
  430. for item in spare.iter_mut().take(n - 1) {
  431. item.write(value.clone());
  432. }
  433. // We can write the last element directly without cloning needlessly.
  434. spare[n - 1].write(value);
  435. // SAFETY:
  436. // - `self.len() + n < self.capacity()` due to the call to reserve above,
  437. // - the loop and the line above initialized the next `n` elements.
  438. unsafe { self.set_len(self.len() + n) };
  439. Ok(())
  440. }
  441. /// Pushes clones of the elements of slice into the [`Vec`] instance.
  442. ///
  443. /// # Examples
  444. ///
  445. /// ```
  446. /// let mut v = KVec::new();
  447. /// v.push(1, GFP_KERNEL)?;
  448. ///
  449. /// v.extend_from_slice(&[20, 30, 40], GFP_KERNEL)?;
  450. /// assert_eq!(&v, &[1, 20, 30, 40]);
  451. ///
  452. /// v.extend_from_slice(&[50, 60], GFP_KERNEL)?;
  453. /// assert_eq!(&v, &[1, 20, 30, 40, 50, 60]);
  454. /// # Ok::<(), Error>(())
  455. /// ```
  456. pub fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError> {
  457. self.reserve(other.len(), flags)?;
  458. for (slot, item) in core::iter::zip(self.spare_capacity_mut(), other) {
  459. slot.write(item.clone());
  460. }
  461. // SAFETY:
  462. // - `other.len()` spare entries have just been initialized, so it is safe to increase
  463. // the length by the same number.
  464. // - `self.len() + other.len() <= self.capacity()` is guaranteed by the preceding `reserve`
  465. // call.
  466. unsafe { self.set_len(self.len() + other.len()) };
  467. Ok(())
  468. }
  469. /// Create a new `Vec<T, A>` and extend it by `n` clones of `value`.
  470. pub fn from_elem(value: T, n: usize, flags: Flags) -> Result<Self, AllocError> {
  471. let mut v = Self::with_capacity(n, flags)?;
  472. v.extend_with(n, value, flags)?;
  473. Ok(v)
  474. }
  475. }
  476. impl<T, A> Drop for Vec<T, A>
  477. where
  478. A: Allocator,
  479. {
  480. fn drop(&mut self) {
  481. // SAFETY: `self.as_mut_ptr` is guaranteed to be valid by the type invariant.
  482. unsafe {
  483. ptr::drop_in_place(core::ptr::slice_from_raw_parts_mut(
  484. self.as_mut_ptr(),
  485. self.len,
  486. ))
  487. };
  488. // SAFETY:
  489. // - `self.ptr` was previously allocated with `A`.
  490. // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
  491. unsafe { A::free(self.ptr.cast(), self.layout.into()) };
  492. }
  493. }
  494. impl<T, A, const N: usize> From<Box<[T; N], A>> for Vec<T, A>
  495. where
  496. A: Allocator,
  497. {
  498. fn from(b: Box<[T; N], A>) -> Vec<T, A> {
  499. let len = b.len();
  500. let ptr = Box::into_raw(b);
  501. // SAFETY:
  502. // - `b` has been allocated with `A`,
  503. // - `ptr` fulfills the alignment requirements for `T`,
  504. // - `ptr` points to memory with at least a size of `size_of::<T>() * len`,
  505. // - all elements within `b` are initialized values of `T`,
  506. // - `len` does not exceed `isize::MAX`.
  507. unsafe { Vec::from_raw_parts(ptr as _, len, len) }
  508. }
  509. }
  510. impl<T> Default for KVec<T> {
  511. #[inline]
  512. fn default() -> Self {
  513. Self::new()
  514. }
  515. }
  516. impl<T: fmt::Debug, A: Allocator> fmt::Debug for Vec<T, A> {
  517. fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  518. fmt::Debug::fmt(&**self, f)
  519. }
  520. }
  521. impl<T, A> Deref for Vec<T, A>
  522. where
  523. A: Allocator,
  524. {
  525. type Target = [T];
  526. #[inline]
  527. fn deref(&self) -> &[T] {
  528. // SAFETY: The memory behind `self.as_ptr()` is guaranteed to contain `self.len`
  529. // initialized elements of type `T`.
  530. unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
  531. }
  532. }
  533. impl<T, A> DerefMut for Vec<T, A>
  534. where
  535. A: Allocator,
  536. {
  537. #[inline]
  538. fn deref_mut(&mut self) -> &mut [T] {
  539. // SAFETY: The memory behind `self.as_ptr()` is guaranteed to contain `self.len`
  540. // initialized elements of type `T`.
  541. unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
  542. }
  543. }
  544. impl<T: Eq, A> Eq for Vec<T, A> where A: Allocator {}
  545. impl<T, I: SliceIndex<[T]>, A> Index<I> for Vec<T, A>
  546. where
  547. A: Allocator,
  548. {
  549. type Output = I::Output;
  550. #[inline]
  551. fn index(&self, index: I) -> &Self::Output {
  552. Index::index(&**self, index)
  553. }
  554. }
  555. impl<T, I: SliceIndex<[T]>, A> IndexMut<I> for Vec<T, A>
  556. where
  557. A: Allocator,
  558. {
  559. #[inline]
  560. fn index_mut(&mut self, index: I) -> &mut Self::Output {
  561. IndexMut::index_mut(&mut **self, index)
  562. }
  563. }
  564. macro_rules! impl_slice_eq {
  565. ($([$($vars:tt)*] $lhs:ty, $rhs:ty,)*) => {
  566. $(
  567. impl<T, U, $($vars)*> PartialEq<$rhs> for $lhs
  568. where
  569. T: PartialEq<U>,
  570. {
  571. #[inline]
  572. fn eq(&self, other: &$rhs) -> bool { self[..] == other[..] }
  573. }
  574. )*
  575. }
  576. }
  577. impl_slice_eq! {
  578. [A1: Allocator, A2: Allocator] Vec<T, A1>, Vec<U, A2>,
  579. [A: Allocator] Vec<T, A>, &[U],
  580. [A: Allocator] Vec<T, A>, &mut [U],
  581. [A: Allocator] &[T], Vec<U, A>,
  582. [A: Allocator] &mut [T], Vec<U, A>,
  583. [A: Allocator] Vec<T, A>, [U],
  584. [A: Allocator] [T], Vec<U, A>,
  585. [A: Allocator, const N: usize] Vec<T, A>, [U; N],
  586. [A: Allocator, const N: usize] Vec<T, A>, &[U; N],
  587. }
  588. impl<'a, T, A> IntoIterator for &'a Vec<T, A>
  589. where
  590. A: Allocator,
  591. {
  592. type Item = &'a T;
  593. type IntoIter = slice::Iter<'a, T>;
  594. fn into_iter(self) -> Self::IntoIter {
  595. self.iter()
  596. }
  597. }
  598. impl<'a, T, A: Allocator> IntoIterator for &'a mut Vec<T, A>
  599. where
  600. A: Allocator,
  601. {
  602. type Item = &'a mut T;
  603. type IntoIter = slice::IterMut<'a, T>;
  604. fn into_iter(self) -> Self::IntoIter {
  605. self.iter_mut()
  606. }
  607. }
  608. /// An [`Iterator`] implementation for [`Vec`] that moves elements out of a vector.
  609. ///
  610. /// This structure is created by the [`Vec::into_iter`] method on [`Vec`] (provided by the
  611. /// [`IntoIterator`] trait).
  612. ///
  613. /// # Examples
  614. ///
  615. /// ```
  616. /// let v = kernel::kvec![0, 1, 2]?;
  617. /// let iter = v.into_iter();
  618. ///
  619. /// # Ok::<(), Error>(())
  620. /// ```
  621. pub struct IntoIter<T, A: Allocator> {
  622. ptr: *mut T,
  623. buf: NonNull<T>,
  624. len: usize,
  625. layout: ArrayLayout<T>,
  626. _p: PhantomData<A>,
  627. }
  628. impl<T, A> IntoIter<T, A>
  629. where
  630. A: Allocator,
  631. {
  632. fn into_raw_parts(self) -> (*mut T, NonNull<T>, usize, usize) {
  633. let me = ManuallyDrop::new(self);
  634. let ptr = me.ptr;
  635. let buf = me.buf;
  636. let len = me.len;
  637. let cap = me.layout.len();
  638. (ptr, buf, len, cap)
  639. }
  640. /// Same as `Iterator::collect` but specialized for `Vec`'s `IntoIter`.
  641. ///
  642. /// # Examples
  643. ///
  644. /// ```
  645. /// let v = kernel::kvec![1, 2, 3]?;
  646. /// let mut it = v.into_iter();
  647. ///
  648. /// assert_eq!(it.next(), Some(1));
  649. ///
  650. /// let v = it.collect(GFP_KERNEL);
  651. /// assert_eq!(v, [2, 3]);
  652. ///
  653. /// # Ok::<(), Error>(())
  654. /// ```
  655. ///
  656. /// # Implementation details
  657. ///
  658. /// Currently, we can't implement `FromIterator`. There are a couple of issues with this trait
  659. /// in the kernel, namely:
  660. ///
  661. /// - Rust's specialization feature is unstable. This prevents us to optimize for the special
  662. /// case where `I::IntoIter` equals `Vec`'s `IntoIter` type.
  663. /// - We also can't use `I::IntoIter`'s type ID either to work around this, since `FromIterator`
  664. /// doesn't require this type to be `'static`.
  665. /// - `FromIterator::from_iter` does return `Self` instead of `Result<Self, AllocError>`, hence
  666. /// we can't properly handle allocation failures.
  667. /// - Neither `Iterator::collect` nor `FromIterator::from_iter` can handle additional allocation
  668. /// flags.
  669. ///
  670. /// Instead, provide `IntoIter::collect`, such that we can at least convert a `IntoIter` into a
  671. /// `Vec` again.
  672. ///
  673. /// Note that `IntoIter::collect` doesn't require `Flags`, since it re-uses the existing backing
  674. /// buffer. However, this backing buffer may be shrunk to the actual count of elements.
  675. pub fn collect(self, flags: Flags) -> Vec<T, A> {
  676. let old_layout = self.layout;
  677. let (mut ptr, buf, len, mut cap) = self.into_raw_parts();
  678. let has_advanced = ptr != buf.as_ptr();
  679. if has_advanced {
  680. // Copy the contents we have advanced to at the beginning of the buffer.
  681. //
  682. // SAFETY:
  683. // - `ptr` is valid for reads of `len * size_of::<T>()` bytes,
  684. // - `buf.as_ptr()` is valid for writes of `len * size_of::<T>()` bytes,
  685. // - `ptr` and `buf.as_ptr()` are not be subject to aliasing restrictions relative to
  686. // each other,
  687. // - both `ptr` and `buf.ptr()` are properly aligned.
  688. unsafe { ptr::copy(ptr, buf.as_ptr(), len) };
  689. ptr = buf.as_ptr();
  690. // SAFETY: `len` is guaranteed to be smaller than `self.layout.len()`.
  691. let layout = unsafe { ArrayLayout::<T>::new_unchecked(len) };
  692. // SAFETY: `buf` points to the start of the backing buffer and `len` is guaranteed to be
  693. // smaller than `cap`. Depending on `alloc` this operation may shrink the buffer or leaves
  694. // it as it is.
  695. ptr = match unsafe {
  696. A::realloc(Some(buf.cast()), layout.into(), old_layout.into(), flags)
  697. } {
  698. // If we fail to shrink, which likely can't even happen, continue with the existing
  699. // buffer.
  700. Err(_) => ptr,
  701. Ok(ptr) => {
  702. cap = len;
  703. ptr.as_ptr().cast()
  704. }
  705. };
  706. }
  707. // SAFETY: If the iterator has been advanced, the advanced elements have been copied to
  708. // the beginning of the buffer and `len` has been adjusted accordingly.
  709. //
  710. // - `ptr` is guaranteed to point to the start of the backing buffer.
  711. // - `cap` is either the original capacity or, after shrinking the buffer, equal to `len`.
  712. // - `alloc` is guaranteed to be unchanged since `into_iter` has been called on the original
  713. // `Vec`.
  714. unsafe { Vec::from_raw_parts(ptr, len, cap) }
  715. }
  716. }
  717. impl<T, A> Iterator for IntoIter<T, A>
  718. where
  719. A: Allocator,
  720. {
  721. type Item = T;
  722. /// # Examples
  723. ///
  724. /// ```
  725. /// let v = kernel::kvec![1, 2, 3]?;
  726. /// let mut it = v.into_iter();
  727. ///
  728. /// assert_eq!(it.next(), Some(1));
  729. /// assert_eq!(it.next(), Some(2));
  730. /// assert_eq!(it.next(), Some(3));
  731. /// assert_eq!(it.next(), None);
  732. ///
  733. /// # Ok::<(), Error>(())
  734. /// ```
  735. fn next(&mut self) -> Option<T> {
  736. if self.len == 0 {
  737. return None;
  738. }
  739. let current = self.ptr;
  740. // SAFETY: We can't overflow; decreasing `self.len` by one every time we advance `self.ptr`
  741. // by one guarantees that.
  742. unsafe { self.ptr = self.ptr.add(1) };
  743. self.len -= 1;
  744. // SAFETY: `current` is guaranteed to point at a valid element within the buffer.
  745. Some(unsafe { current.read() })
  746. }
  747. /// # Examples
  748. ///
  749. /// ```
  750. /// let v: KVec<u32> = kernel::kvec![1, 2, 3]?;
  751. /// let mut iter = v.into_iter();
  752. /// let size = iter.size_hint().0;
  753. ///
  754. /// iter.next();
  755. /// assert_eq!(iter.size_hint().0, size - 1);
  756. ///
  757. /// iter.next();
  758. /// assert_eq!(iter.size_hint().0, size - 2);
  759. ///
  760. /// iter.next();
  761. /// assert_eq!(iter.size_hint().0, size - 3);
  762. ///
  763. /// # Ok::<(), Error>(())
  764. /// ```
  765. fn size_hint(&self) -> (usize, Option<usize>) {
  766. (self.len, Some(self.len))
  767. }
  768. }
  769. impl<T, A> Drop for IntoIter<T, A>
  770. where
  771. A: Allocator,
  772. {
  773. fn drop(&mut self) {
  774. // SAFETY: `self.ptr` is guaranteed to be valid by the type invariant.
  775. unsafe { ptr::drop_in_place(ptr::slice_from_raw_parts_mut(self.ptr, self.len)) };
  776. // SAFETY:
  777. // - `self.buf` was previously allocated with `A`.
  778. // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
  779. unsafe { A::free(self.buf.cast(), self.layout.into()) };
  780. }
  781. }
  782. impl<T, A> IntoIterator for Vec<T, A>
  783. where
  784. A: Allocator,
  785. {
  786. type Item = T;
  787. type IntoIter = IntoIter<T, A>;
  788. /// Consumes the `Vec<T, A>` and creates an `Iterator`, which moves each value out of the
  789. /// vector (from start to end).
  790. ///
  791. /// # Examples
  792. ///
  793. /// ```
  794. /// let v = kernel::kvec![1, 2]?;
  795. /// let mut v_iter = v.into_iter();
  796. ///
  797. /// let first_element: Option<u32> = v_iter.next();
  798. ///
  799. /// assert_eq!(first_element, Some(1));
  800. /// assert_eq!(v_iter.next(), Some(2));
  801. /// assert_eq!(v_iter.next(), None);
  802. ///
  803. /// # Ok::<(), Error>(())
  804. /// ```
  805. ///
  806. /// ```
  807. /// let v = kernel::kvec![];
  808. /// let mut v_iter = v.into_iter();
  809. ///
  810. /// let first_element: Option<u32> = v_iter.next();
  811. ///
  812. /// assert_eq!(first_element, None);
  813. ///
  814. /// # Ok::<(), Error>(())
  815. /// ```
  816. #[inline]
  817. fn into_iter(self) -> Self::IntoIter {
  818. let buf = self.ptr;
  819. let layout = self.layout;
  820. let (ptr, len, _) = self.into_raw_parts();
  821. IntoIter {
  822. ptr,
  823. buf,
  824. len,
  825. layout,
  826. _p: PhantomData::<A>,
  827. }
  828. }
  829. }