drain.rs 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186
  1. // SPDX-License-Identifier: Apache-2.0 OR MIT
  2. use crate::alloc::{Allocator, Global};
  3. use core::fmt;
  4. use core::iter::{FusedIterator, TrustedLen};
  5. use core::mem;
  6. use core::ptr::{self, NonNull};
  7. use core::slice::{self};
  8. use super::Vec;
  9. /// A draining iterator for `Vec<T>`.
  10. ///
  11. /// This `struct` is created by [`Vec::drain`].
  12. /// See its documentation for more.
  13. ///
  14. /// # Example
  15. ///
  16. /// ```
  17. /// let mut v = vec![0, 1, 2];
  18. /// let iter: std::vec::Drain<_> = v.drain(..);
  19. /// ```
  20. #[stable(feature = "drain", since = "1.6.0")]
  21. pub struct Drain<
  22. 'a,
  23. T: 'a,
  24. #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + 'a = Global,
  25. > {
  26. /// Index of tail to preserve
  27. pub(super) tail_start: usize,
  28. /// Length of tail
  29. pub(super) tail_len: usize,
  30. /// Current remaining range to remove
  31. pub(super) iter: slice::Iter<'a, T>,
  32. pub(super) vec: NonNull<Vec<T, A>>,
  33. }
  34. #[stable(feature = "collection_debug", since = "1.17.0")]
  35. impl<T: fmt::Debug, A: Allocator> fmt::Debug for Drain<'_, T, A> {
  36. fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  37. f.debug_tuple("Drain").field(&self.iter.as_slice()).finish()
  38. }
  39. }
  40. impl<'a, T, A: Allocator> Drain<'a, T, A> {
  41. /// Returns the remaining items of this iterator as a slice.
  42. ///
  43. /// # Examples
  44. ///
  45. /// ```
  46. /// let mut vec = vec!['a', 'b', 'c'];
  47. /// let mut drain = vec.drain(..);
  48. /// assert_eq!(drain.as_slice(), &['a', 'b', 'c']);
  49. /// let _ = drain.next().unwrap();
  50. /// assert_eq!(drain.as_slice(), &['b', 'c']);
  51. /// ```
  52. #[must_use]
  53. #[stable(feature = "vec_drain_as_slice", since = "1.46.0")]
  54. pub fn as_slice(&self) -> &[T] {
  55. self.iter.as_slice()
  56. }
  57. /// Returns a reference to the underlying allocator.
  58. #[unstable(feature = "allocator_api", issue = "32838")]
  59. #[must_use]
  60. #[inline]
  61. pub fn allocator(&self) -> &A {
  62. unsafe { self.vec.as_ref().allocator() }
  63. }
  64. }
  65. #[stable(feature = "vec_drain_as_slice", since = "1.46.0")]
  66. impl<'a, T, A: Allocator> AsRef<[T]> for Drain<'a, T, A> {
  67. fn as_ref(&self) -> &[T] {
  68. self.as_slice()
  69. }
  70. }
  71. #[stable(feature = "drain", since = "1.6.0")]
  72. unsafe impl<T: Sync, A: Sync + Allocator> Sync for Drain<'_, T, A> {}
  73. #[stable(feature = "drain", since = "1.6.0")]
  74. unsafe impl<T: Send, A: Send + Allocator> Send for Drain<'_, T, A> {}
  75. #[stable(feature = "drain", since = "1.6.0")]
  76. impl<T, A: Allocator> Iterator for Drain<'_, T, A> {
  77. type Item = T;
  78. #[inline]
  79. fn next(&mut self) -> Option<T> {
  80. self.iter.next().map(|elt| unsafe { ptr::read(elt as *const _) })
  81. }
  82. fn size_hint(&self) -> (usize, Option<usize>) {
  83. self.iter.size_hint()
  84. }
  85. }
  86. #[stable(feature = "drain", since = "1.6.0")]
  87. impl<T, A: Allocator> DoubleEndedIterator for Drain<'_, T, A> {
  88. #[inline]
  89. fn next_back(&mut self) -> Option<T> {
  90. self.iter.next_back().map(|elt| unsafe { ptr::read(elt as *const _) })
  91. }
  92. }
  93. #[stable(feature = "drain", since = "1.6.0")]
  94. impl<T, A: Allocator> Drop for Drain<'_, T, A> {
  95. fn drop(&mut self) {
  96. /// Moves back the un-`Drain`ed elements to restore the original `Vec`.
  97. struct DropGuard<'r, 'a, T, A: Allocator>(&'r mut Drain<'a, T, A>);
  98. impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> {
  99. fn drop(&mut self) {
  100. if self.0.tail_len > 0 {
  101. unsafe {
  102. let source_vec = self.0.vec.as_mut();
  103. // memmove back untouched tail, update to new length
  104. let start = source_vec.len();
  105. let tail = self.0.tail_start;
  106. if tail != start {
  107. let src = source_vec.as_ptr().add(tail);
  108. let dst = source_vec.as_mut_ptr().add(start);
  109. ptr::copy(src, dst, self.0.tail_len);
  110. }
  111. source_vec.set_len(start + self.0.tail_len);
  112. }
  113. }
  114. }
  115. }
  116. let iter = mem::replace(&mut self.iter, (&mut []).iter());
  117. let drop_len = iter.len();
  118. let mut vec = self.vec;
  119. if mem::size_of::<T>() == 0 {
  120. // ZSTs have no identity, so we don't need to move them around, we only need to drop the correct amount.
  121. // this can be achieved by manipulating the Vec length instead of moving values out from `iter`.
  122. unsafe {
  123. let vec = vec.as_mut();
  124. let old_len = vec.len();
  125. vec.set_len(old_len + drop_len + self.tail_len);
  126. vec.truncate(old_len + self.tail_len);
  127. }
  128. return;
  129. }
  130. // ensure elements are moved back into their appropriate places, even when drop_in_place panics
  131. let _guard = DropGuard(self);
  132. if drop_len == 0 {
  133. return;
  134. }
  135. // as_slice() must only be called when iter.len() is > 0 because
  136. // vec::Splice modifies vec::Drain fields and may grow the vec which would invalidate
  137. // the iterator's internal pointers. Creating a reference to deallocated memory
  138. // is invalid even when it is zero-length
  139. let drop_ptr = iter.as_slice().as_ptr();
  140. unsafe {
  141. // drop_ptr comes from a slice::Iter which only gives us a &[T] but for drop_in_place
  142. // a pointer with mutable provenance is necessary. Therefore we must reconstruct
  143. // it from the original vec but also avoid creating a &mut to the front since that could
  144. // invalidate raw pointers to it which some unsafe code might rely on.
  145. let vec_ptr = vec.as_mut().as_mut_ptr();
  146. let drop_offset = drop_ptr.sub_ptr(vec_ptr);
  147. let to_drop = ptr::slice_from_raw_parts_mut(vec_ptr.add(drop_offset), drop_len);
  148. ptr::drop_in_place(to_drop);
  149. }
  150. }
  151. }
  152. #[stable(feature = "drain", since = "1.6.0")]
  153. impl<T, A: Allocator> ExactSizeIterator for Drain<'_, T, A> {
  154. fn is_empty(&self) -> bool {
  155. self.iter.is_empty()
  156. }
  157. }
  158. #[unstable(feature = "trusted_len", issue = "37572")]
  159. unsafe impl<T, A: Allocator> TrustedLen for Drain<'_, T, A> {}
  160. #[stable(feature = "fused", since = "1.26.0")]
  161. impl<T, A: Allocator> FusedIterator for Drain<'_, T, A> {}