1 // SPDX-License-Identifier: GPL-2.0 2 3 //! Implementation of [`Vec`]. 4 5 use super::{ 6 allocator::{ 7 KVmalloc, 8 Kmalloc, 9 Vmalloc, 10 VmallocPageIter, // 11 }, 12 flags::__GFP_ZERO, 13 layout::ArrayLayout, 14 AllocError, 15 Allocator, 16 Box, 17 Flags, 18 NumaNode, // 19 }; 20 21 use crate::{ 22 fmt, 23 page::{ 24 AsPageIter, 25 PAGE_SIZE, // 26 }, // 27 }; 28 29 use core::{ 30 borrow::{ 31 Borrow, 32 BorrowMut, // 33 }, 34 marker::PhantomData, 35 mem::{ 36 ManuallyDrop, 37 MaybeUninit, // 38 }, 39 ops::{ 40 Deref, 41 DerefMut, 42 Index, 43 IndexMut, // 44 }, 45 ptr::{ 46 self, 47 NonNull, // 48 }, 49 slice::{ 50 self, 51 SliceIndex, // 52 }, // 53 }; 54 55 use pin_init::Zeroable; 56 57 mod errors; 58 pub use self::errors::{InsertError, PushError, RemoveError}; 59 60 /// Create a [`KVec`] containing the arguments. 61 /// 62 /// New memory is allocated with `GFP_KERNEL`. 63 /// 64 /// # Examples 65 /// 66 /// ``` 67 /// let mut v = kernel::kvec![]; 68 /// v.push(1, GFP_KERNEL)?; 69 /// assert_eq!(v, [1]); 70 /// 71 /// let mut v = kernel::kvec![1; 3]?; 72 /// v.push(4, GFP_KERNEL)?; 73 /// assert_eq!(v, [1, 1, 1, 4]); 74 /// 75 /// let mut v = kernel::kvec![1, 2, 3]?; 76 /// v.push(4, GFP_KERNEL)?; 77 /// assert_eq!(v, [1, 2, 3, 4]); 78 /// 79 /// # Ok::<(), Error>(()) 80 /// ``` 81 #[macro_export] 82 macro_rules! kvec { 83 () => ( 84 $crate::alloc::KVec::new() 85 ); 86 ($elem:expr; $n:expr) => ( 87 $crate::alloc::KVec::from_elem($elem, $n, GFP_KERNEL) 88 ); 89 ($($x:expr),+ $(,)?) => ( 90 match $crate::alloc::KBox::new_uninit(GFP_KERNEL) { 91 Ok(b) => Ok($crate::alloc::KVec::from($crate::alloc::KBox::write(b, [$($x),+]))), 92 Err(e) => Err(e), 93 } 94 ); 95 } 96 97 /// The kernel's [`Vec`] type. 98 /// 99 /// A contiguous growable array type with contents allocated with the kernel's allocators (e.g. 100 /// [`Kmalloc`], [`Vmalloc`] or [`KVmalloc`]), written `Vec<T, A>`. 101 /// 102 /// For non-zero-sized values, a [`Vec`] will use the given allocator `A` for its allocation. For 103 /// the most common allocators the type aliases [`KVec`], [`VVec`] and [`KVVec`] exist. 104 /// 105 /// For zero-sized types the [`Vec`]'s pointer must be `dangling_mut::<T>`; no memory is allocated. 106 /// 107 /// Generally, [`Vec`] consists of a pointer that represents the vector's backing buffer, the 108 /// capacity of the vector (the number of elements that currently fit into the vector), its length 109 /// (the number of elements that are currently stored in the vector) and the `Allocator` type used 110 /// to allocate (and free) the backing buffer. 111 /// 112 /// A [`Vec`] can be deconstructed into and (re-)constructed from its previously named raw parts 113 /// and manually modified. 114 /// 115 /// [`Vec`]'s backing buffer gets, if required, automatically increased (re-allocated) when elements 116 /// are added to the vector. 117 /// 118 /// # Invariants 119 /// 120 /// - `self.ptr` is always properly aligned and either points to memory allocated with `A` or, for 121 /// zero-sized types, is a dangling, well aligned pointer. 122 /// 123 /// - `self.len` always represents the exact number of elements stored in the vector. 124 /// 125 /// - `self.layout` represents the absolute number of elements that can be stored within the vector 126 /// without re-allocation. For ZSTs `self.layout`'s capacity is zero. However, it is legal for the 127 /// backing buffer to be larger than `layout`. 128 /// 129 /// - `self.len()` is always less than or equal to `self.capacity()`. 130 /// 131 /// - The `Allocator` type `A` of the vector is the exact same `Allocator` type the backing buffer 132 /// was allocated with (and must be freed with). 133 pub struct Vec<T, A: Allocator> { 134 ptr: NonNull<T>, 135 /// Represents the actual buffer size as `cap` times `size_of::<T>` bytes. 136 /// 137 /// Note: This isn't quite the same as `Self::capacity`, which in contrast returns the number of 138 /// elements we can still store without reallocating. 139 layout: ArrayLayout<T>, 140 len: usize, 141 _p: PhantomData<A>, 142 } 143 144 /// Type alias for [`Vec`] with a [`Kmalloc`] allocator. 145 /// 146 /// # Examples 147 /// 148 /// ``` 149 /// let mut v = KVec::new(); 150 /// v.push(1, GFP_KERNEL)?; 151 /// assert_eq!(&v, &[1]); 152 /// 153 /// # Ok::<(), Error>(()) 154 /// ``` 155 pub type KVec<T> = Vec<T, Kmalloc>; 156 157 /// Type alias for [`Vec`] with a [`Vmalloc`] allocator. 158 /// 159 /// # Examples 160 /// 161 /// ``` 162 /// let mut v = VVec::new(); 163 /// v.push(1, GFP_KERNEL)?; 164 /// assert_eq!(&v, &[1]); 165 /// 166 /// # Ok::<(), Error>(()) 167 /// ``` 168 pub type VVec<T> = Vec<T, Vmalloc>; 169 170 /// Type alias for [`Vec`] with a [`KVmalloc`] allocator. 171 /// 172 /// # Examples 173 /// 174 /// ``` 175 /// let mut v = KVVec::new(); 176 /// v.push(1, GFP_KERNEL)?; 177 /// assert_eq!(&v, &[1]); 178 /// 179 /// # Ok::<(), Error>(()) 180 /// ``` 181 pub type KVVec<T> = Vec<T, KVmalloc>; 182 183 // SAFETY: `Vec` is `Send` if `T` is `Send` because `Vec` owns its elements. 184 unsafe impl<T, A> Send for Vec<T, A> 185 where 186 T: Send, 187 A: Allocator, 188 { 189 } 190 191 // SAFETY: `Vec` is `Sync` if `T` is `Sync` because `Vec` owns its elements. 192 unsafe impl<T, A> Sync for Vec<T, A> 193 where 194 T: Sync, 195 A: Allocator, 196 { 197 } 198 199 impl<T, A> Vec<T, A> 200 where 201 A: Allocator, 202 { 203 #[inline] 204 const fn is_zst() -> bool { 205 core::mem::size_of::<T>() == 0 206 } 207 208 /// Returns the number of elements that can be stored within the vector without allocating 209 /// additional memory. 210 pub const fn capacity(&self) -> usize { 211 if const { Self::is_zst() } { 212 usize::MAX 213 } else { 214 self.layout.len() 215 } 216 } 217 218 /// Returns the number of elements stored within the vector. 219 #[inline] 220 pub const fn len(&self) -> usize { 221 self.len 222 } 223 224 /// Increments `self.len` by `additional`. 225 /// 226 /// # Safety 227 /// 228 /// - `additional` must be less than or equal to `self.capacity - self.len`. 229 /// - All elements within the interval [`self.len`,`self.len + additional`) must be initialized. 230 #[inline] 231 pub const unsafe fn inc_len(&mut self, additional: usize) { 232 // Guaranteed by the type invariant to never underflow. 233 debug_assert!(additional <= self.capacity() - self.len()); 234 // INVARIANT: By the safety requirements of this method this represents the exact number of 235 // elements stored within `self`. 236 self.len += additional; 237 } 238 239 /// Decreases `self.len` by `count`. 240 /// 241 /// Returns a mutable slice to the elements forgotten by the vector. It is the caller's 242 /// responsibility to drop these elements if necessary. 243 /// 244 /// # Safety 245 /// 246 /// - `count` must be less than or equal to `self.len`. 247 unsafe fn dec_len(&mut self, count: usize) -> &mut [T] { 248 debug_assert!(count <= self.len()); 249 // INVARIANT: We relinquish ownership of the elements within the range `[self.len - count, 250 // self.len)`, hence the updated value of `set.len` represents the exact number of elements 251 // stored within `self`. 252 self.len -= count; 253 // SAFETY: The memory after `self.len()` is guaranteed to contain `count` initialized 254 // elements of type `T`. 255 unsafe { slice::from_raw_parts_mut(self.as_mut_ptr().add(self.len), count) } 256 } 257 258 /// Returns a slice of the entire vector. 259 /// 260 /// # Examples 261 /// 262 /// ``` 263 /// let mut v = KVec::new(); 264 /// v.push(1, GFP_KERNEL)?; 265 /// v.push(2, GFP_KERNEL)?; 266 /// assert_eq!(v.as_slice(), &[1, 2]); 267 /// # Ok::<(), Error>(()) 268 /// ``` 269 #[inline] 270 pub fn as_slice(&self) -> &[T] { 271 self 272 } 273 274 /// Returns a mutable slice of the entire vector. 275 #[inline] 276 pub fn as_mut_slice(&mut self) -> &mut [T] { 277 self 278 } 279 280 /// Returns a mutable raw pointer to the vector's backing buffer, or, if `T` is a ZST, a 281 /// dangling raw pointer. 282 #[inline] 283 pub fn as_mut_ptr(&mut self) -> *mut T { 284 self.ptr.as_ptr() 285 } 286 287 /// Returns a raw pointer to the vector's backing buffer, or, if `T` is a ZST, a dangling raw 288 /// pointer. 289 #[inline] 290 pub const fn as_ptr(&self) -> *const T { 291 self.ptr.as_ptr() 292 } 293 294 /// Returns `true` if the vector contains no elements, `false` otherwise. 295 /// 296 /// # Examples 297 /// 298 /// ``` 299 /// let mut v = KVec::new(); 300 /// assert!(v.is_empty()); 301 /// 302 /// v.push(1, GFP_KERNEL); 303 /// assert!(!v.is_empty()); 304 /// ``` 305 #[inline] 306 pub const fn is_empty(&self) -> bool { 307 self.len() == 0 308 } 309 310 /// Creates a new, empty `Vec<T, A>`. 311 /// 312 /// This method does not allocate by itself. 313 #[inline] 314 pub const fn new() -> Self { 315 // INVARIANT: Since this is a new, empty `Vec` with no backing memory yet, 316 // - `ptr` is a properly aligned dangling pointer for type `T`, 317 // - `layout` is an empty `ArrayLayout` (zero capacity) 318 // - `len` is zero, since no elements can be or have been stored, 319 // - `A` is always valid. 320 Self { 321 ptr: NonNull::dangling(), 322 layout: ArrayLayout::empty(), 323 len: 0, 324 _p: PhantomData::<A>, 325 } 326 } 327 328 /// Returns a slice of `MaybeUninit<T>` for the remaining spare capacity of the vector. 329 pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] { 330 // SAFETY: 331 // - `self.len` is smaller than `self.capacity` by the type invariant and hence, the 332 // resulting pointer is guaranteed to be part of the same allocated object. 333 // - `self.len` can not overflow `isize`. 334 let ptr = unsafe { self.as_mut_ptr().add(self.len) }.cast::<MaybeUninit<T>>(); 335 336 // SAFETY: The memory between `self.len` and `self.capacity` is guaranteed to be allocated 337 // and valid, but uninitialized. 338 unsafe { slice::from_raw_parts_mut(ptr, self.capacity() - self.len) } 339 } 340 341 /// Appends an element to the back of the [`Vec`] instance. 342 /// 343 /// # Examples 344 /// 345 /// ``` 346 /// let mut v = KVec::new(); 347 /// v.push(1, GFP_KERNEL)?; 348 /// assert_eq!(&v, &[1]); 349 /// 350 /// v.push(2, GFP_KERNEL)?; 351 /// assert_eq!(&v, &[1, 2]); 352 /// # Ok::<(), Error>(()) 353 /// ``` 354 pub fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> { 355 self.reserve(1, flags)?; 356 // SAFETY: The call to `reserve` was successful, so the capacity is at least one greater 357 // than the length. 358 unsafe { self.push_within_capacity_unchecked(v) }; 359 Ok(()) 360 } 361 362 /// Appends an element to the back of the [`Vec`] instance without reallocating. 363 /// 364 /// Fails if the vector does not have capacity for the new element. 365 /// 366 /// # Examples 367 /// 368 /// ``` 369 /// let mut v = KVec::with_capacity(10, GFP_KERNEL)?; 370 /// for i in 0..10 { 371 /// v.push_within_capacity(i)?; 372 /// } 373 /// 374 /// assert!(v.push_within_capacity(10).is_err()); 375 /// # Ok::<(), Error>(()) 376 /// ``` 377 pub fn push_within_capacity(&mut self, v: T) -> Result<(), PushError<T>> { 378 if self.len() < self.capacity() { 379 // SAFETY: The length is less than the capacity. 380 unsafe { self.push_within_capacity_unchecked(v) }; 381 Ok(()) 382 } else { 383 Err(PushError(v)) 384 } 385 } 386 387 /// Appends an element to the back of the [`Vec`] instance without reallocating. 388 /// 389 /// # Safety 390 /// 391 /// The length must be less than the capacity. 392 unsafe fn push_within_capacity_unchecked(&mut self, v: T) { 393 let spare = self.spare_capacity_mut(); 394 395 // SAFETY: By the safety requirements, `spare` is non-empty. 396 unsafe { spare.get_unchecked_mut(0) }.write(v); 397 398 // SAFETY: We just initialised the first spare entry, so it is safe to increase the length 399 // by 1. We also know that the new length is <= capacity because the caller guarantees that 400 // the length is less than the capacity at the beginning of this function. 401 unsafe { self.inc_len(1) }; 402 } 403 404 /// Inserts an element at the given index in the [`Vec`] instance. 405 /// 406 /// Fails if the vector does not have capacity for the new element. Panics if the index is out 407 /// of bounds. 408 /// 409 /// # Examples 410 /// 411 /// ``` 412 /// use kernel::alloc::kvec::InsertError; 413 /// 414 /// let mut v = KVec::with_capacity(5, GFP_KERNEL)?; 415 /// for i in 0..5 { 416 /// v.insert_within_capacity(0, i)?; 417 /// } 418 /// 419 /// assert!(matches!(v.insert_within_capacity(0, 5), Err(InsertError::OutOfCapacity(_)))); 420 /// assert!(matches!(v.insert_within_capacity(1000, 5), Err(InsertError::IndexOutOfBounds(_)))); 421 /// assert_eq!(v, [4, 3, 2, 1, 0]); 422 /// # Ok::<(), Error>(()) 423 /// ``` 424 pub fn insert_within_capacity( 425 &mut self, 426 index: usize, 427 element: T, 428 ) -> Result<(), InsertError<T>> { 429 let len = self.len(); 430 if index > len { 431 return Err(InsertError::IndexOutOfBounds(element)); 432 } 433 434 if len >= self.capacity() { 435 return Err(InsertError::OutOfCapacity(element)); 436 } 437 438 // SAFETY: This is in bounds since `index <= len < capacity`. 439 let p = unsafe { self.as_mut_ptr().add(index) }; 440 // INVARIANT: This breaks the Vec invariants by making `index` contain an invalid element, 441 // but we restore the invariants below. 442 // SAFETY: Both the src and dst ranges end no later than one element after the length. 443 // Since the length is less than the capacity, both ranges are in bounds of the allocation. 444 unsafe { ptr::copy(p, p.add(1), len - index) }; 445 // INVARIANT: This restores the Vec invariants. 446 // SAFETY: The pointer is in-bounds of the allocation. 447 unsafe { ptr::write(p, element) }; 448 // SAFETY: Index `len` contains a valid element due to the above copy and write. 449 unsafe { self.inc_len(1) }; 450 Ok(()) 451 } 452 453 /// Removes the last element from a vector and returns it, or `None` if it is empty. 454 /// 455 /// # Examples 456 /// 457 /// ``` 458 /// let mut v = KVec::new(); 459 /// v.push(1, GFP_KERNEL)?; 460 /// v.push(2, GFP_KERNEL)?; 461 /// assert_eq!(&v, &[1, 2]); 462 /// 463 /// assert_eq!(v.pop(), Some(2)); 464 /// assert_eq!(v.pop(), Some(1)); 465 /// assert_eq!(v.pop(), None); 466 /// # Ok::<(), Error>(()) 467 /// ``` 468 pub fn pop(&mut self) -> Option<T> { 469 if self.is_empty() { 470 return None; 471 } 472 473 let removed: *mut T = { 474 // SAFETY: We just checked that the length is at least one. 475 let slice = unsafe { self.dec_len(1) }; 476 // SAFETY: The argument to `dec_len` was 1 so this returns a slice of length 1. 477 unsafe { slice.get_unchecked_mut(0) } 478 }; 479 480 // SAFETY: The guarantees of `dec_len` allow us to take ownership of this value. 481 Some(unsafe { removed.read() }) 482 } 483 484 /// Removes the element at the given index. 485 /// 486 /// # Examples 487 /// 488 /// ``` 489 /// let mut v = kernel::kvec![1, 2, 3]?; 490 /// assert_eq!(v.remove(1)?, 2); 491 /// assert_eq!(v, [1, 3]); 492 /// # Ok::<(), Error>(()) 493 /// ``` 494 pub fn remove(&mut self, i: usize) -> Result<T, RemoveError> { 495 let value = { 496 let value_ref = self.get(i).ok_or(RemoveError)?; 497 // INVARIANT: This breaks the invariants by invalidating the value at index `i`, but we 498 // restore the invariants below. 499 // SAFETY: The value at index `i` is valid, because otherwise we would have already 500 // failed with `RemoveError`. 501 unsafe { ptr::read(value_ref) } 502 }; 503 504 // SAFETY: We checked that `i` is in-bounds. 505 let p = unsafe { self.as_mut_ptr().add(i) }; 506 507 // INVARIANT: After this call, the invalid value is at the last slot, so the Vec invariants 508 // are restored after the below call to `dec_len(1)`. 509 // SAFETY: `p.add(1).add(self.len - i - 1)` is `i+1+len-i-1 == len` elements after the 510 // beginning of the vector, so this is in-bounds of the vector's allocation. 511 unsafe { ptr::copy(p.add(1), p, self.len - i - 1) }; 512 513 // SAFETY: Since the check at the beginning of this call did not fail with `RemoveError`, 514 // the length is at least one. 515 unsafe { self.dec_len(1) }; 516 517 Ok(value) 518 } 519 520 /// Creates a new [`Vec`] instance with at least the given capacity. 521 /// 522 /// # Examples 523 /// 524 /// ``` 525 /// let v = KVec::<u32>::with_capacity(20, GFP_KERNEL)?; 526 /// 527 /// assert!(v.capacity() >= 20); 528 /// # Ok::<(), Error>(()) 529 /// ``` 530 pub fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError> { 531 let mut v = Vec::new(); 532 533 v.reserve(capacity, flags)?; 534 535 Ok(v) 536 } 537 538 /// Creates a new [`Vec`] with `n` zero-initialized elements. 539 /// 540 /// # Examples 541 /// 542 /// ``` 543 /// let v = KVec::<u32>::zeroed(20, GFP_KERNEL)?; 544 /// 545 /// assert!(v.iter().all(|&x| x == 0)); 546 /// # Ok::<(), Error>(()) 547 /// ``` 548 pub fn zeroed(n: usize, flags: Flags) -> Result<Self, AllocError> 549 where 550 T: Zeroable, 551 { 552 let mut v = Self::with_capacity(n, flags | __GFP_ZERO)?; 553 554 // SAFETY: 555 // - `n <= capacity - len`: `with_capacity(n)` guarantees capacity >= n, len is 0. 556 // - All elements in `[0, n)` are initialized: `__GFP_ZERO` zeroes the allocation, 557 // and `T: Zeroable` guarantees all-zeroes is a valid bit pattern. 558 unsafe { v.inc_len(n) }; 559 Ok(v) 560 } 561 562 /// Creates a `Vec<T, A>` from a pointer, a length and a capacity using the allocator `A`. 563 /// 564 /// # Examples 565 /// 566 /// ``` 567 /// let mut v = kernel::kvec![1, 2, 3]?; 568 /// v.reserve(1, GFP_KERNEL)?; 569 /// 570 /// let (mut ptr, mut len, cap) = v.into_raw_parts(); 571 /// 572 /// // SAFETY: We've just reserved memory for another element. 573 /// unsafe { ptr.add(len).write(4) }; 574 /// len += 1; 575 /// 576 /// // SAFETY: We only wrote an additional element at the end of the `KVec`'s buffer and 577 /// // correspondingly increased the length of the `KVec` by one. Otherwise, we construct it 578 /// // from the exact same raw parts. 579 /// let v = unsafe { KVec::from_raw_parts(ptr, len, cap) }; 580 /// 581 /// assert_eq!(v, [1, 2, 3, 4]); 582 /// 583 /// # Ok::<(), Error>(()) 584 /// ``` 585 /// 586 /// # Safety 587 /// 588 /// If `T` is a ZST: 589 /// 590 /// - `ptr` must be a dangling, well aligned pointer. 591 /// 592 /// Otherwise: 593 /// 594 /// - `ptr` must have been allocated with the allocator `A`. 595 /// - `ptr` must satisfy or exceed the alignment requirements of `T`. 596 /// - `ptr` must point to memory with a size of at least `size_of::<T>() * capacity` bytes. 597 /// - The allocated size in bytes must not be larger than `isize::MAX`. 598 /// - `length` must be less than or equal to `capacity`. 599 /// - The first `length` elements must be initialized values of type `T`. 600 /// 601 /// It is also valid to create an empty `Vec` passing a dangling pointer for `ptr` and zero for 602 /// `cap` and `len`. 603 pub unsafe fn from_raw_parts(ptr: *mut T, length: usize, capacity: usize) -> Self { 604 let layout = if Self::is_zst() { 605 ArrayLayout::empty() 606 } else { 607 // SAFETY: By the safety requirements of this function, `capacity * size_of::<T>()` is 608 // smaller than `isize::MAX`. 609 unsafe { ArrayLayout::new_unchecked(capacity) } 610 }; 611 612 // INVARIANT: For ZSTs, we store an empty `ArrayLayout`, all other type invariants are 613 // covered by the safety requirements of this function. 614 Self { 615 // SAFETY: By the safety requirements, `ptr` is either dangling or pointing to a valid 616 // memory allocation, allocated with `A`. 617 ptr: unsafe { NonNull::new_unchecked(ptr) }, 618 layout, 619 len: length, 620 _p: PhantomData::<A>, 621 } 622 } 623 624 /// Consumes the `Vec<T, A>` and returns its raw components `pointer`, `length` and `capacity`. 625 /// 626 /// This will not run the destructor of the contained elements and for non-ZSTs the allocation 627 /// will stay alive indefinitely. Use [`Vec::from_raw_parts`] to recover the [`Vec`], drop the 628 /// elements and free the allocation, if any. 629 pub fn into_raw_parts(self) -> (*mut T, usize, usize) { 630 let mut me = ManuallyDrop::new(self); 631 let len = me.len(); 632 let capacity = me.capacity(); 633 let ptr = me.as_mut_ptr(); 634 (ptr, len, capacity) 635 } 636 637 /// Clears the vector, removing all values. 638 /// 639 /// Note that this method has no effect on the allocated capacity 640 /// of the vector. 641 /// 642 /// # Examples 643 /// 644 /// ``` 645 /// let mut v = kernel::kvec![1, 2, 3]?; 646 /// 647 /// v.clear(); 648 /// 649 /// assert!(v.is_empty()); 650 /// # Ok::<(), Error>(()) 651 /// ``` 652 #[inline] 653 pub fn clear(&mut self) { 654 self.truncate(0); 655 } 656 657 /// Ensures that the capacity exceeds the length by at least `additional` elements. 658 /// 659 /// # Examples 660 /// 661 /// ``` 662 /// let mut v = KVec::new(); 663 /// v.push(1, GFP_KERNEL)?; 664 /// 665 /// v.reserve(10, GFP_KERNEL)?; 666 /// let cap = v.capacity(); 667 /// assert!(cap >= v.len() + 10); 668 /// 669 /// v.reserve(10, GFP_KERNEL)?; 670 /// let new_cap = v.capacity(); 671 /// assert_eq!(new_cap, cap); 672 /// 673 /// # Ok::<(), Error>(()) 674 /// ``` 675 pub fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError> { 676 let len = self.len(); 677 let cap = self.capacity(); 678 679 if cap - len >= additional { 680 return Ok(()); 681 } 682 683 if Self::is_zst() { 684 // The capacity is already `usize::MAX` for ZSTs, we can't go higher. 685 return Err(AllocError); 686 } 687 688 // We know that `cap <= isize::MAX` because of the type invariants of `Self`. So the 689 // multiplication by two won't overflow. 690 let new_cap = core::cmp::max(cap * 2, len.checked_add(additional).ok_or(AllocError)?); 691 let layout = ArrayLayout::new(new_cap).map_err(|_| AllocError)?; 692 693 // SAFETY: 694 // - `ptr` is valid because it's either `None` or comes from a previous call to 695 // `A::realloc`. 696 // - `self.layout` matches the `ArrayLayout` of the preceding allocation. 697 let ptr = unsafe { 698 A::realloc( 699 Some(self.ptr.cast()), 700 layout.into(), 701 self.layout.into(), 702 flags, 703 NumaNode::NO_NODE, 704 )? 705 }; 706 707 // INVARIANT: 708 // - `layout` is some `ArrayLayout::<T>`, 709 // - `ptr` has been created by `A::realloc` from `layout`. 710 self.ptr = ptr.cast(); 711 self.layout = layout; 712 713 Ok(()) 714 } 715 716 /// Shortens the vector, setting the length to `len` and drops the removed values. 717 /// If `len` is greater than or equal to the current length, this does nothing. 718 /// 719 /// This has no effect on the capacity and will not allocate. 720 /// 721 /// # Examples 722 /// 723 /// ``` 724 /// let mut v = kernel::kvec![1, 2, 3]?; 725 /// v.truncate(1); 726 /// assert_eq!(v.len(), 1); 727 /// assert_eq!(&v, &[1]); 728 /// 729 /// # Ok::<(), Error>(()) 730 /// ``` 731 pub fn truncate(&mut self, len: usize) { 732 if let Some(count) = self.len().checked_sub(len) { 733 // SAFETY: `count` is `self.len() - len` so it is guaranteed to be less than or 734 // equal to `self.len()`. 735 let ptr: *mut [T] = unsafe { self.dec_len(count) }; 736 737 // SAFETY: the contract of `dec_len` guarantees that the elements in `ptr` are 738 // valid elements whose ownership has been transferred to the caller. 739 unsafe { ptr::drop_in_place(ptr) }; 740 } 741 } 742 743 /// Takes ownership of all items in this vector without consuming the allocation. 744 /// 745 /// # Examples 746 /// 747 /// ``` 748 /// let mut v = kernel::kvec![0, 1, 2, 3]?; 749 /// 750 /// for (i, j) in v.drain_all().enumerate() { 751 /// assert_eq!(i, j); 752 /// } 753 /// 754 /// assert!(v.capacity() >= 4); 755 /// # Ok::<(), Error>(()) 756 /// ``` 757 pub fn drain_all(&mut self) -> DrainAll<'_, T> { 758 // SAFETY: This does not underflow the length. 759 let elems = unsafe { self.dec_len(self.len()) }; 760 // INVARIANT: The first `len` elements of the spare capacity are valid values, and as we 761 // just set the length to zero, we may transfer ownership to the `DrainAll` object. 762 DrainAll { 763 elements: elems.iter_mut(), 764 } 765 } 766 767 /// Removes all elements that don't match the provided closure. 768 /// 769 /// # Examples 770 /// 771 /// ``` 772 /// let mut v = kernel::kvec![1, 2, 3, 4]?; 773 /// v.retain(|i| *i % 2 == 0); 774 /// assert_eq!(v, [2, 4]); 775 /// # Ok::<(), Error>(()) 776 /// ``` 777 pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) { 778 let mut num_kept = 0; 779 let mut next_to_check = 0; 780 while let Some(to_check) = self.get_mut(next_to_check) { 781 if f(to_check) { 782 self.swap(num_kept, next_to_check); 783 num_kept += 1; 784 } 785 next_to_check += 1; 786 } 787 self.truncate(num_kept); 788 } 789 } 790 // TODO: This is a temporary KVVec-specific implementation. It should be replaced with a generic 791 // `shrink_to()` for `impl<T, A: Allocator> Vec<T, A>` that uses `A::realloc()` once the 792 // underlying allocators properly support shrinking via realloc. 793 impl<T> Vec<T, KVmalloc> { 794 /// Shrinks the capacity of the vector with a lower bound. 795 /// 796 /// The capacity will remain at least as large as both the length and the supplied value. 797 /// If the current capacity is less than the lower limit, this is a no-op. 798 /// 799 /// For `kmalloc` allocations, this delegates to `realloc()`, which decides whether 800 /// shrinking is worthwhile. For `vmalloc` allocations, shrinking only occurs if the 801 /// operation would free at least one page of memory, and performs a deep copy since 802 /// `vrealloc` does not yet support in-place shrinking. 803 /// 804 /// # Examples 805 /// 806 /// ``` 807 /// // Allocate enough capacity to span multiple pages. 808 /// let elements_per_page = kernel::page::PAGE_SIZE / core::mem::size_of::<u32>(); 809 /// let mut v = KVVec::with_capacity(elements_per_page * 4, GFP_KERNEL)?; 810 /// v.push(1, GFP_KERNEL)?; 811 /// v.push(2, GFP_KERNEL)?; 812 /// 813 /// v.shrink_to(0, GFP_KERNEL)?; 814 /// # Ok::<(), Error>(()) 815 /// ``` 816 pub fn shrink_to(&mut self, min_capacity: usize, flags: Flags) -> Result<(), AllocError> { 817 let target_cap = core::cmp::max(self.len(), min_capacity); 818 819 if self.capacity() <= target_cap { 820 return Ok(()); 821 } 822 823 if Self::is_zst() { 824 return Ok(()); 825 } 826 827 // For kmalloc allocations, delegate to realloc() and let the allocator decide 828 // whether shrinking is worthwhile. 829 // 830 // SAFETY: `self.ptr` points to a valid `KVmalloc` allocation. 831 if !unsafe { bindings::is_vmalloc_addr(self.ptr.as_ptr().cast()) } { 832 let new_layout = ArrayLayout::<T>::new(target_cap).map_err(|_| AllocError)?; 833 834 // SAFETY: 835 // - `self.ptr` is valid and was previously allocated with `KVmalloc`. 836 // - `self.layout` matches the `ArrayLayout` of the preceding allocation. 837 let ptr = unsafe { 838 KVmalloc::realloc( 839 Some(self.ptr.cast()), 840 new_layout.into(), 841 self.layout.into(), 842 flags, 843 NumaNode::NO_NODE, 844 )? 845 }; 846 847 self.ptr = ptr.cast(); 848 self.layout = new_layout; 849 return Ok(()); 850 } 851 852 // Only shrink if we would free at least one page. 853 let current_size = self.capacity() * core::mem::size_of::<T>(); 854 let target_size = target_cap * core::mem::size_of::<T>(); 855 let current_pages = current_size.div_ceil(PAGE_SIZE); 856 let target_pages = target_size.div_ceil(PAGE_SIZE); 857 858 if current_pages <= target_pages { 859 return Ok(()); 860 } 861 862 if target_cap == 0 { 863 if !self.layout.is_empty() { 864 // SAFETY: 865 // - `self.ptr` was previously allocated with `KVmalloc`. 866 // - `self.layout` matches the `ArrayLayout` of the preceding allocation. 867 unsafe { KVmalloc::free(self.ptr.cast(), self.layout.into()) }; 868 } 869 self.ptr = NonNull::dangling(); 870 self.layout = ArrayLayout::empty(); 871 return Ok(()); 872 } 873 874 // SAFETY: `target_cap <= self.capacity()` and original capacity was valid. 875 let new_layout = unsafe { ArrayLayout::<T>::new_unchecked(target_cap) }; 876 877 let new_ptr = KVmalloc::alloc(new_layout.into(), flags, NumaNode::NO_NODE)?; 878 879 // SAFETY: 880 // - `self.as_ptr()` is valid for reads of `self.len()` elements of `T`. 881 // - `new_ptr` is valid for writes of at least `target_cap >= self.len()` elements. 882 // - The two allocations do not overlap since `new_ptr` is freshly allocated. 883 // - Both pointers are properly aligned for `T`. 884 unsafe { 885 ptr::copy_nonoverlapping(self.as_ptr(), new_ptr.as_ptr().cast::<T>(), self.len()) 886 }; 887 888 // SAFETY: 889 // - `self.ptr` was previously allocated with `KVmalloc`. 890 // - `self.layout` matches the `ArrayLayout` of the preceding allocation. 891 unsafe { KVmalloc::free(self.ptr.cast(), self.layout.into()) }; 892 893 self.ptr = new_ptr.cast::<T>(); 894 self.layout = new_layout; 895 896 Ok(()) 897 } 898 } 899 900 impl<T: Clone, A: Allocator> Vec<T, A> { 901 /// Extend the vector by `n` clones of `value`. 902 /// 903 /// # Examples 904 /// 905 /// ``` 906 /// let mut v = KVec::new(); 907 /// v.push(1, GFP_KERNEL)?; 908 /// 909 /// v.extend_with(3, 5, GFP_KERNEL)?; 910 /// assert_eq!(&v, &[1, 5, 5, 5]); 911 /// 912 /// v.extend_with(2, 8, GFP_KERNEL)?; 913 /// assert_eq!(&v, &[1, 5, 5, 5, 8, 8]); 914 /// 915 /// v.extend_with(0, 3, GFP_KERNEL)?; 916 /// assert_eq!(&v, &[1, 5, 5, 5, 8, 8]); 917 /// 918 /// # Ok::<(), Error>(()) 919 /// ``` 920 pub fn extend_with(&mut self, n: usize, value: T, flags: Flags) -> Result<(), AllocError> { 921 if n == 0 { 922 return Ok(()); 923 } 924 925 self.reserve(n, flags)?; 926 927 let spare = self.spare_capacity_mut(); 928 929 for item in spare.iter_mut().take(n - 1) { 930 item.write(value.clone()); 931 } 932 933 // We can write the last element directly without cloning needlessly. 934 spare[n - 1].write(value); 935 936 // SAFETY: 937 // - `self.len() + n <= self.capacity()` due to the call to reserve above, 938 // - the loop and the line above initialized the next `n` elements. 939 unsafe { self.inc_len(n) }; 940 941 Ok(()) 942 } 943 944 /// Pushes clones of the elements of slice into the [`Vec`] instance. 945 /// 946 /// # Examples 947 /// 948 /// ``` 949 /// let mut v = KVec::new(); 950 /// v.push(1, GFP_KERNEL)?; 951 /// 952 /// v.extend_from_slice(&[20, 30, 40], GFP_KERNEL)?; 953 /// assert_eq!(&v, &[1, 20, 30, 40]); 954 /// 955 /// v.extend_from_slice(&[50, 60], GFP_KERNEL)?; 956 /// assert_eq!(&v, &[1, 20, 30, 40, 50, 60]); 957 /// # Ok::<(), Error>(()) 958 /// ``` 959 pub fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError> { 960 self.reserve(other.len(), flags)?; 961 for (slot, item) in core::iter::zip(self.spare_capacity_mut(), other) { 962 slot.write(item.clone()); 963 } 964 965 // SAFETY: 966 // - `other.len()` spare entries have just been initialized, so it is safe to increase 967 // the length by the same number. 968 // - `self.len() + other.len() <= self.capacity()` is guaranteed by the preceding `reserve` 969 // call. 970 unsafe { self.inc_len(other.len()) }; 971 Ok(()) 972 } 973 974 /// Create a new `Vec<T, A>` and extend it by `n` clones of `value`. 975 pub fn from_elem(value: T, n: usize, flags: Flags) -> Result<Self, AllocError> { 976 let mut v = Self::with_capacity(n, flags)?; 977 978 v.extend_with(n, value, flags)?; 979 980 Ok(v) 981 } 982 983 /// Resizes the [`Vec`] so that `len` is equal to `new_len`. 984 /// 985 /// If `new_len` is smaller than `len`, the `Vec` is [`Vec::truncate`]d. 986 /// If `new_len` is larger, each new slot is filled with clones of `value`. 987 /// 988 /// # Examples 989 /// 990 /// ``` 991 /// let mut v = kernel::kvec![1, 2, 3]?; 992 /// v.resize(1, 42, GFP_KERNEL)?; 993 /// assert_eq!(&v, &[1]); 994 /// 995 /// v.resize(3, 42, GFP_KERNEL)?; 996 /// assert_eq!(&v, &[1, 42, 42]); 997 /// 998 /// # Ok::<(), Error>(()) 999 /// ``` 1000 pub fn resize(&mut self, new_len: usize, value: T, flags: Flags) -> Result<(), AllocError> { 1001 match new_len.checked_sub(self.len()) { 1002 Some(n) => self.extend_with(n, value, flags), 1003 None => { 1004 self.truncate(new_len); 1005 Ok(()) 1006 } 1007 } 1008 } 1009 } 1010 1011 impl<T, A> Drop for Vec<T, A> 1012 where 1013 A: Allocator, 1014 { 1015 fn drop(&mut self) { 1016 // SAFETY: `self.as_mut_ptr` is guaranteed to be valid by the type invariant. 1017 unsafe { 1018 ptr::drop_in_place(core::ptr::slice_from_raw_parts_mut( 1019 self.as_mut_ptr(), 1020 self.len, 1021 )) 1022 }; 1023 1024 // SAFETY: 1025 // - `self.ptr` was previously allocated with `A`. 1026 // - `self.layout` matches the `ArrayLayout` of the preceding allocation. 1027 unsafe { A::free(self.ptr.cast(), self.layout.into()) }; 1028 } 1029 } 1030 1031 impl<T, A, const N: usize> From<Box<[T; N], A>> for Vec<T, A> 1032 where 1033 A: Allocator, 1034 { 1035 fn from(b: Box<[T; N], A>) -> Vec<T, A> { 1036 let len = b.len(); 1037 let ptr = Box::into_raw(b); 1038 1039 // SAFETY: 1040 // - `b` has been allocated with `A`, 1041 // - `ptr` fulfills the alignment requirements for `T`, 1042 // - `ptr` points to memory with at least a size of `size_of::<T>() * len`, 1043 // - all elements within `b` are initialized values of `T`, 1044 // - `len` does not exceed `isize::MAX`. 1045 unsafe { Vec::from_raw_parts(ptr.cast(), len, len) } 1046 } 1047 } 1048 1049 impl<T, A: Allocator> Default for Vec<T, A> { 1050 #[inline] 1051 fn default() -> Self { 1052 Self::new() 1053 } 1054 } 1055 1056 impl<T: fmt::Debug, A: Allocator> fmt::Debug for Vec<T, A> { 1057 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 1058 fmt::Debug::fmt(&**self, f) 1059 } 1060 } 1061 1062 impl<T, A> Deref for Vec<T, A> 1063 where 1064 A: Allocator, 1065 { 1066 type Target = [T]; 1067 1068 #[inline] 1069 fn deref(&self) -> &[T] { 1070 // SAFETY: The memory behind `self.as_ptr()` is guaranteed to contain `self.len` 1071 // initialized elements of type `T`. 1072 unsafe { slice::from_raw_parts(self.as_ptr(), self.len) } 1073 } 1074 } 1075 1076 impl<T, A> DerefMut for Vec<T, A> 1077 where 1078 A: Allocator, 1079 { 1080 #[inline] 1081 fn deref_mut(&mut self) -> &mut [T] { 1082 // SAFETY: The memory behind `self.as_ptr()` is guaranteed to contain `self.len` 1083 // initialized elements of type `T`. 1084 unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) } 1085 } 1086 } 1087 1088 /// # Examples 1089 /// 1090 /// ``` 1091 /// # use core::borrow::Borrow; 1092 /// struct Foo<B: Borrow<[u32]>>(B); 1093 /// 1094 /// // Owned array. 1095 /// let owned_array = Foo([1, 2, 3]); 1096 /// 1097 /// // Owned vector. 1098 /// let owned_vec = Foo(KVec::from_elem(0, 3, GFP_KERNEL)?); 1099 /// 1100 /// let arr = [1, 2, 3]; 1101 /// // Borrowed slice from `arr`. 1102 /// let borrowed_slice = Foo(&arr[..]); 1103 /// # Ok::<(), Error>(()) 1104 /// ``` 1105 impl<T, A> Borrow<[T]> for Vec<T, A> 1106 where 1107 A: Allocator, 1108 { 1109 fn borrow(&self) -> &[T] { 1110 self.as_slice() 1111 } 1112 } 1113 1114 /// # Examples 1115 /// 1116 /// ``` 1117 /// # use core::borrow::BorrowMut; 1118 /// struct Foo<B: BorrowMut<[u32]>>(B); 1119 /// 1120 /// // Owned array. 1121 /// let owned_array = Foo([1, 2, 3]); 1122 /// 1123 /// // Owned vector. 1124 /// let owned_vec = Foo(KVec::from_elem(0, 3, GFP_KERNEL)?); 1125 /// 1126 /// let mut arr = [1, 2, 3]; 1127 /// // Borrowed slice from `arr`. 1128 /// let borrowed_slice = Foo(&mut arr[..]); 1129 /// # Ok::<(), Error>(()) 1130 /// ``` 1131 impl<T, A> BorrowMut<[T]> for Vec<T, A> 1132 where 1133 A: Allocator, 1134 { 1135 fn borrow_mut(&mut self) -> &mut [T] { 1136 self.as_mut_slice() 1137 } 1138 } 1139 1140 impl<T: Eq, A> Eq for Vec<T, A> where A: Allocator {} 1141 1142 impl<T, I: SliceIndex<[T]>, A> Index<I> for Vec<T, A> 1143 where 1144 A: Allocator, 1145 { 1146 type Output = I::Output; 1147 1148 #[inline] 1149 fn index(&self, index: I) -> &Self::Output { 1150 Index::index(&**self, index) 1151 } 1152 } 1153 1154 impl<T, I: SliceIndex<[T]>, A> IndexMut<I> for Vec<T, A> 1155 where 1156 A: Allocator, 1157 { 1158 #[inline] 1159 fn index_mut(&mut self, index: I) -> &mut Self::Output { 1160 IndexMut::index_mut(&mut **self, index) 1161 } 1162 } 1163 1164 macro_rules! impl_slice_eq { 1165 ($([$($vars:tt)*] $lhs:ty, $rhs:ty,)*) => { 1166 $( 1167 impl<T, U, $($vars)*> PartialEq<$rhs> for $lhs 1168 where 1169 T: PartialEq<U>, 1170 { 1171 #[inline] 1172 fn eq(&self, other: &$rhs) -> bool { self[..] == other[..] } 1173 } 1174 )* 1175 } 1176 } 1177 1178 impl_slice_eq! { 1179 [A1: Allocator, A2: Allocator] Vec<T, A1>, Vec<U, A2>, 1180 [A: Allocator] Vec<T, A>, &[U], 1181 [A: Allocator] Vec<T, A>, &mut [U], 1182 [A: Allocator] &[T], Vec<U, A>, 1183 [A: Allocator] &mut [T], Vec<U, A>, 1184 [A: Allocator] Vec<T, A>, [U], 1185 [A: Allocator] [T], Vec<U, A>, 1186 [A: Allocator, const N: usize] Vec<T, A>, [U; N], 1187 [A: Allocator, const N: usize] Vec<T, A>, &[U; N], 1188 } 1189 1190 impl<'a, T, A> IntoIterator for &'a Vec<T, A> 1191 where 1192 A: Allocator, 1193 { 1194 type Item = &'a T; 1195 type IntoIter = slice::Iter<'a, T>; 1196 1197 fn into_iter(self) -> Self::IntoIter { 1198 self.iter() 1199 } 1200 } 1201 1202 impl<'a, T, A: Allocator> IntoIterator for &'a mut Vec<T, A> 1203 where 1204 A: Allocator, 1205 { 1206 type Item = &'a mut T; 1207 type IntoIter = slice::IterMut<'a, T>; 1208 1209 fn into_iter(self) -> Self::IntoIter { 1210 self.iter_mut() 1211 } 1212 } 1213 1214 /// # Examples 1215 /// 1216 /// ``` 1217 /// use kernel::{ 1218 /// alloc::allocator::VmallocPageIter, 1219 /// page::{ 1220 /// AsPageIter, 1221 /// PAGE_SIZE, // 1222 /// }, // 1223 /// }; 1224 /// 1225 /// let mut vec = VVec::<u8>::new(); 1226 /// 1227 /// assert!(vec.page_iter().next().is_none()); 1228 /// 1229 /// vec.reserve(PAGE_SIZE, GFP_KERNEL)?; 1230 /// 1231 /// let page = vec.page_iter().next().expect("At least one page should be available.\n"); 1232 /// 1233 /// // SAFETY: There is no concurrent read or write to the same page. 1234 /// unsafe { page.fill_zero_raw(0, PAGE_SIZE)? }; 1235 /// # Ok::<(), Error>(()) 1236 /// ``` 1237 impl<T> AsPageIter for VVec<T> { 1238 type Iter<'a> 1239 = VmallocPageIter<'a> 1240 where 1241 T: 'a; 1242 1243 fn page_iter(&mut self) -> Self::Iter<'_> { 1244 let ptr = self.ptr.cast(); 1245 let size = self.layout.size(); 1246 1247 // SAFETY: 1248 // - `ptr` is a valid pointer to the beginning of a `Vmalloc` allocation. 1249 // - `ptr` is guaranteed to be valid for the lifetime of `'a`. 1250 // - `size` is the size of the `Vmalloc` allocation `ptr` points to. 1251 unsafe { VmallocPageIter::new(ptr, size) } 1252 } 1253 } 1254 1255 /// An [`Iterator`] implementation for [`Vec`] that moves elements out of a vector. 1256 /// 1257 /// This structure is created by the [`Vec::into_iter`] method on [`Vec`] (provided by the 1258 /// [`IntoIterator`] trait). 1259 /// 1260 /// # Examples 1261 /// 1262 /// ``` 1263 /// let v = kernel::kvec![0, 1, 2]?; 1264 /// let iter = v.into_iter(); 1265 /// 1266 /// # Ok::<(), Error>(()) 1267 /// ``` 1268 pub struct IntoIter<T, A: Allocator> { 1269 ptr: *mut T, 1270 buf: NonNull<T>, 1271 len: usize, 1272 layout: ArrayLayout<T>, 1273 _p: PhantomData<A>, 1274 } 1275 1276 impl<T, A> IntoIter<T, A> 1277 where 1278 A: Allocator, 1279 { 1280 fn into_raw_parts(self) -> (*mut T, NonNull<T>, usize, usize) { 1281 let me = ManuallyDrop::new(self); 1282 let ptr = me.ptr; 1283 let buf = me.buf; 1284 let len = me.len; 1285 let cap = me.layout.len(); 1286 (ptr, buf, len, cap) 1287 } 1288 1289 /// Same as `Iterator::collect` but specialized for `Vec`'s `IntoIter`. 1290 /// 1291 /// # Examples 1292 /// 1293 /// ``` 1294 /// let v = kernel::kvec![1, 2, 3]?; 1295 /// let mut it = v.into_iter(); 1296 /// 1297 /// assert_eq!(it.next(), Some(1)); 1298 /// 1299 /// let v = it.collect(GFP_KERNEL); 1300 /// assert_eq!(v, [2, 3]); 1301 /// 1302 /// # Ok::<(), Error>(()) 1303 /// ``` 1304 /// 1305 /// # Implementation details 1306 /// 1307 /// Currently, we can't implement `FromIterator`. There are a couple of issues with this trait 1308 /// in the kernel, namely: 1309 /// 1310 /// - Rust's specialization feature is unstable. This prevents us to optimize for the special 1311 /// case where `I::IntoIter` equals `Vec`'s `IntoIter` type. 1312 /// - We also can't use `I::IntoIter`'s type ID either to work around this, since `FromIterator` 1313 /// doesn't require this type to be `'static`. 1314 /// - `FromIterator::from_iter` does return `Self` instead of `Result<Self, AllocError>`, hence 1315 /// we can't properly handle allocation failures. 1316 /// - Neither `Iterator::collect` nor `FromIterator::from_iter` can handle additional allocation 1317 /// flags. 1318 /// 1319 /// Instead, provide `IntoIter::collect`, such that we can at least convert a `IntoIter` into a 1320 /// `Vec` again. 1321 /// 1322 /// Note that `IntoIter::collect` doesn't require `Flags`, since it re-uses the existing backing 1323 /// buffer. However, this backing buffer may be shrunk to the actual count of elements. 1324 pub fn collect(self, flags: Flags) -> Vec<T, A> { 1325 let old_layout = self.layout; 1326 let (mut ptr, buf, len, mut cap) = self.into_raw_parts(); 1327 let has_advanced = ptr != buf.as_ptr(); 1328 1329 if has_advanced { 1330 // Copy the contents we have advanced to at the beginning of the buffer. 1331 // 1332 // SAFETY: 1333 // - `ptr` is valid for reads of `len * size_of::<T>()` bytes, 1334 // - `buf.as_ptr()` is valid for writes of `len * size_of::<T>()` bytes, 1335 // - `ptr` and `buf.as_ptr()` are not be subject to aliasing restrictions relative to 1336 // each other, 1337 // - both `ptr` and `buf.ptr()` are properly aligned. 1338 unsafe { ptr::copy(ptr, buf.as_ptr(), len) }; 1339 ptr = buf.as_ptr(); 1340 1341 // SAFETY: `len` is guaranteed to be smaller than `self.layout.len()` by the type 1342 // invariant. 1343 let layout = unsafe { ArrayLayout::<T>::new_unchecked(len) }; 1344 1345 // SAFETY: `buf` points to the start of the backing buffer and `len` is guaranteed by 1346 // the type invariant to be smaller than `cap`. Depending on `realloc` this operation 1347 // may shrink the buffer or leave it as it is. 1348 ptr = match unsafe { 1349 A::realloc( 1350 Some(buf.cast()), 1351 layout.into(), 1352 old_layout.into(), 1353 flags, 1354 NumaNode::NO_NODE, 1355 ) 1356 } { 1357 // If we fail to shrink, which likely can't even happen, continue with the existing 1358 // buffer. 1359 Err(_) => ptr, 1360 Ok(ptr) => { 1361 cap = len; 1362 ptr.as_ptr().cast() 1363 } 1364 }; 1365 } 1366 1367 // SAFETY: If the iterator has been advanced, the advanced elements have been copied to 1368 // the beginning of the buffer and `len` has been adjusted accordingly. 1369 // 1370 // - `ptr` is guaranteed to point to the start of the backing buffer. 1371 // - `cap` is either the original capacity or, after shrinking the buffer, equal to `len`. 1372 // - `alloc` is guaranteed to be unchanged since `into_iter` has been called on the original 1373 // `Vec`. 1374 unsafe { Vec::from_raw_parts(ptr, len, cap) } 1375 } 1376 } 1377 1378 impl<T, A> Iterator for IntoIter<T, A> 1379 where 1380 A: Allocator, 1381 { 1382 type Item = T; 1383 1384 /// # Examples 1385 /// 1386 /// ``` 1387 /// let v = kernel::kvec![1, 2, 3]?; 1388 /// let mut it = v.into_iter(); 1389 /// 1390 /// assert_eq!(it.next(), Some(1)); 1391 /// assert_eq!(it.next(), Some(2)); 1392 /// assert_eq!(it.next(), Some(3)); 1393 /// assert_eq!(it.next(), None); 1394 /// 1395 /// # Ok::<(), Error>(()) 1396 /// ``` 1397 fn next(&mut self) -> Option<T> { 1398 if self.len == 0 { 1399 return None; 1400 } 1401 1402 let current = self.ptr; 1403 1404 // SAFETY: We can't overflow; decreasing `self.len` by one every time we advance `self.ptr` 1405 // by one guarantees that. 1406 unsafe { self.ptr = self.ptr.add(1) }; 1407 1408 self.len -= 1; 1409 1410 // SAFETY: `current` is guaranteed to point at a valid element within the buffer. 1411 Some(unsafe { current.read() }) 1412 } 1413 1414 /// # Examples 1415 /// 1416 /// ``` 1417 /// let v: KVec<u32> = kernel::kvec![1, 2, 3]?; 1418 /// let mut iter = v.into_iter(); 1419 /// let size = iter.size_hint().0; 1420 /// 1421 /// iter.next(); 1422 /// assert_eq!(iter.size_hint().0, size - 1); 1423 /// 1424 /// iter.next(); 1425 /// assert_eq!(iter.size_hint().0, size - 2); 1426 /// 1427 /// iter.next(); 1428 /// assert_eq!(iter.size_hint().0, size - 3); 1429 /// 1430 /// # Ok::<(), Error>(()) 1431 /// ``` 1432 fn size_hint(&self) -> (usize, Option<usize>) { 1433 (self.len, Some(self.len)) 1434 } 1435 } 1436 1437 impl<T, A> Drop for IntoIter<T, A> 1438 where 1439 A: Allocator, 1440 { 1441 fn drop(&mut self) { 1442 // SAFETY: `self.ptr` is guaranteed to be valid by the type invariant. 1443 unsafe { ptr::drop_in_place(ptr::slice_from_raw_parts_mut(self.ptr, self.len)) }; 1444 1445 // SAFETY: 1446 // - `self.buf` was previously allocated with `A`. 1447 // - `self.layout` matches the `ArrayLayout` of the preceding allocation. 1448 unsafe { A::free(self.buf.cast(), self.layout.into()) }; 1449 } 1450 } 1451 1452 impl<T, A> IntoIterator for Vec<T, A> 1453 where 1454 A: Allocator, 1455 { 1456 type Item = T; 1457 type IntoIter = IntoIter<T, A>; 1458 1459 /// Consumes the `Vec<T, A>` and creates an `Iterator`, which moves each value out of the 1460 /// vector (from start to end). 1461 /// 1462 /// # Examples 1463 /// 1464 /// ``` 1465 /// let v = kernel::kvec![1, 2]?; 1466 /// let mut v_iter = v.into_iter(); 1467 /// 1468 /// let first_element: Option<u32> = v_iter.next(); 1469 /// 1470 /// assert_eq!(first_element, Some(1)); 1471 /// assert_eq!(v_iter.next(), Some(2)); 1472 /// assert_eq!(v_iter.next(), None); 1473 /// 1474 /// # Ok::<(), Error>(()) 1475 /// ``` 1476 /// 1477 /// ``` 1478 /// let v = kernel::kvec![]; 1479 /// let mut v_iter = v.into_iter(); 1480 /// 1481 /// let first_element: Option<u32> = v_iter.next(); 1482 /// 1483 /// assert_eq!(first_element, None); 1484 /// 1485 /// # Ok::<(), Error>(()) 1486 /// ``` 1487 #[inline] 1488 fn into_iter(self) -> Self::IntoIter { 1489 let buf = self.ptr; 1490 let layout = self.layout; 1491 let (ptr, len, _) = self.into_raw_parts(); 1492 1493 IntoIter { 1494 ptr, 1495 buf, 1496 len, 1497 layout, 1498 _p: PhantomData::<A>, 1499 } 1500 } 1501 } 1502 1503 /// An iterator that owns all items in a vector, but does not own its allocation. 1504 /// 1505 /// # Invariants 1506 /// 1507 /// Every `&mut T` returned by the iterator references a `T` that the iterator may take ownership 1508 /// of. 1509 pub struct DrainAll<'vec, T> { 1510 elements: slice::IterMut<'vec, T>, 1511 } 1512 1513 impl<'vec, T> Iterator for DrainAll<'vec, T> { 1514 type Item = T; 1515 1516 fn next(&mut self) -> Option<T> { 1517 let elem: *mut T = self.elements.next()?; 1518 // SAFETY: By the type invariants, we may take ownership of this value. 1519 Some(unsafe { elem.read() }) 1520 } 1521 1522 fn size_hint(&self) -> (usize, Option<usize>) { 1523 self.elements.size_hint() 1524 } 1525 } 1526 1527 impl<'vec, T> Drop for DrainAll<'vec, T> { 1528 fn drop(&mut self) { 1529 if core::mem::needs_drop::<T>() { 1530 let iter = core::mem::take(&mut self.elements); 1531 let ptr: *mut [T] = iter.into_slice(); 1532 // SAFETY: By the type invariants, we own these values so we may destroy them. 1533 unsafe { ptr::drop_in_place(ptr) }; 1534 } 1535 } 1536 } 1537 1538 #[cfg(CONFIG_RUST_KVEC_KUNIT_TEST)] 1539 #[macros::kunit_tests(rust_kvec)] 1540 mod tests { 1541 use super::*; 1542 use crate::prelude::*; 1543 1544 #[test] 1545 fn test_kvec_retain() { 1546 /// Verify correctness for one specific function. 1547 #[expect(clippy::needless_range_loop)] 1548 fn verify(c: &[bool]) { 1549 let mut vec1: KVec<usize> = KVec::with_capacity(c.len(), GFP_KERNEL).unwrap(); 1550 let mut vec2: KVec<usize> = KVec::with_capacity(c.len(), GFP_KERNEL).unwrap(); 1551 1552 for i in 0..c.len() { 1553 vec1.push_within_capacity(i).unwrap(); 1554 if c[i] { 1555 vec2.push_within_capacity(i).unwrap(); 1556 } 1557 } 1558 1559 vec1.retain(|i| c[*i]); 1560 1561 assert_eq!(vec1, vec2); 1562 } 1563 1564 /// Add one to a binary integer represented as a boolean array. 1565 fn add(value: &mut [bool]) { 1566 let mut carry = true; 1567 for v in value { 1568 let new_v = carry != *v; 1569 carry = carry && *v; 1570 *v = new_v; 1571 } 1572 } 1573 1574 // This boolean array represents a function from index to boolean. We check that `retain` 1575 // behaves correctly for all possible boolean arrays of every possible length less than 1576 // ten. 1577 let mut func = KVec::with_capacity(10, GFP_KERNEL).unwrap(); 1578 for len in 0..10 { 1579 for _ in 0u32..1u32 << len { 1580 verify(&func); 1581 add(&mut func); 1582 } 1583 func.push_within_capacity(false).unwrap(); 1584 } 1585 } 1586 1587 #[test] 1588 fn test_kvvec_shrink_to() { 1589 use crate::page::PAGE_SIZE; 1590 1591 // Create a vector with capacity spanning multiple pages. 1592 let mut v = KVVec::<u8>::with_capacity(PAGE_SIZE * 4, GFP_KERNEL).unwrap(); 1593 1594 // Add a few elements. 1595 v.push(1, GFP_KERNEL).unwrap(); 1596 v.push(2, GFP_KERNEL).unwrap(); 1597 v.push(3, GFP_KERNEL).unwrap(); 1598 1599 let initial_capacity = v.capacity(); 1600 assert!(initial_capacity >= PAGE_SIZE * 4); 1601 1602 // Shrink to a capacity that would free at least one page. 1603 v.shrink_to(PAGE_SIZE, GFP_KERNEL).unwrap(); 1604 1605 // Capacity should have been reduced. 1606 assert!(v.capacity() < initial_capacity); 1607 assert!(v.capacity() >= PAGE_SIZE); 1608 1609 // Elements should be preserved. 1610 assert_eq!(v.len(), 3); 1611 assert_eq!(v[0], 1); 1612 assert_eq!(v[1], 2); 1613 assert_eq!(v[2], 3); 1614 1615 // Shrink to zero (should shrink to len). 1616 v.shrink_to(0, GFP_KERNEL).unwrap(); 1617 1618 // Capacity should be at least the length. 1619 assert!(v.capacity() >= v.len()); 1620 1621 // Elements should still be preserved. 1622 assert_eq!(v.len(), 3); 1623 assert_eq!(v[0], 1); 1624 assert_eq!(v[1], 2); 1625 assert_eq!(v[2], 3); 1626 } 1627 1628 #[test] 1629 fn test_kvvec_shrink_to_empty() { 1630 use crate::page::PAGE_SIZE; 1631 1632 // Create a vector with large capacity but no elements. 1633 let mut v = KVVec::<u8>::with_capacity(PAGE_SIZE * 4, GFP_KERNEL).unwrap(); 1634 1635 assert!(v.is_empty()); 1636 1637 // Shrink empty vector to zero. 1638 v.shrink_to(0, GFP_KERNEL).unwrap(); 1639 1640 // Should have freed the allocation. 1641 assert_eq!(v.capacity(), 0); 1642 assert!(v.is_empty()); 1643 } 1644 1645 #[test] 1646 fn test_kvvec_shrink_to_no_op() { 1647 use crate::page::PAGE_SIZE; 1648 1649 // Create a small vector. 1650 let mut v = KVVec::<u8>::with_capacity(PAGE_SIZE, GFP_KERNEL).unwrap(); 1651 v.push(1, GFP_KERNEL).unwrap(); 1652 1653 let capacity_before = v.capacity(); 1654 1655 // Try to shrink to a capacity larger than current - should be no-op. 1656 v.shrink_to(capacity_before + 100, GFP_KERNEL).unwrap(); 1657 1658 assert_eq!(v.capacity(), capacity_before); 1659 assert_eq!(v.len(), 1); 1660 assert_eq!(v[0], 1); 1661 } 1662 1663 #[test] 1664 fn test_kvvec_shrink_to_respects_min_capacity() { 1665 use crate::page::PAGE_SIZE; 1666 1667 // Create a vector with large capacity. 1668 let mut v = KVVec::<u8>::with_capacity(PAGE_SIZE * 4, GFP_KERNEL).unwrap(); 1669 1670 // Add some elements. 1671 for i in 0..10u8 { 1672 v.push(i, GFP_KERNEL).unwrap(); 1673 } 1674 1675 // Shrink to a min_capacity larger than length. 1676 let min_cap = PAGE_SIZE * 2; 1677 v.shrink_to(min_cap, GFP_KERNEL).unwrap(); 1678 1679 // Capacity should be at least min_capacity. 1680 assert!(v.capacity() >= min_cap); 1681 1682 // All elements preserved. 1683 assert_eq!(v.len(), 10); 1684 for i in 0..10u8 { 1685 assert_eq!(v[i as usize], i); 1686 } 1687 } 1688 } 1689