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