library/alloc/src/collections/vec_deque/drain.rs

   1 use core::iter::FusedIterator;
   2 use core::ptr::{self, NonNull};
   3 use core::{fmt, mem};
   4
   5 use crate::alloc::{Allocator, Global};
   6
   7 use super::{count, Iter, VecDeque};
   8
   9 /// A draining iterator over the elements of a `VecDeque`.
  10 ///
  11 /// This `struct` is created by the [`drain`] method on [`VecDeque`]. See its
  12 /// documentation for more.
  13 ///
  14 /// [`drain`]: VecDeque::drain
  15 #[stable(feature = "drain", since = "1.6.0")]
  16 pub struct Drain<
  17     'a,
  18     T: 'a,
  19     #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
  20 > {
  21     after_tail: usize,
  22     after_head: usize,
  23     iter: Iter<'a, T>,
  24     deque: NonNull<VecDeque<T, A>>,
  25 }
  26
  27 impl<'a, T, A: Allocator> Drain<'a, T, A> {
  28     pub(super) unsafe fn new(
  29         after_tail: usize,
  30         after_head: usize,
  31         iter: Iter<'a, T>,
  32         deque: NonNull<VecDeque<T, A>>,
  33     ) -> Self {
  34         Drain { after_tail, after_head, iter, deque }
  35     }
  36 }
  37
  38 #[stable(feature = "collection_debug", since = "1.17.0")]
  39 impl<T: fmt::Debug, A: Allocator> fmt::Debug for Drain<'_, T, A> {
  40     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  41         f.debug_tuple("Drain")
  42             .field(&self.after_tail)
  43             .field(&self.after_head)
  44             .field(&self.iter)
  45             .finish()
  46     }
  47 }
  48
  49 #[stable(feature = "drain", since = "1.6.0")]
  50 unsafe impl<T: Sync, A: Allocator + Sync> Sync for Drain<'_, T, A> {}
  51 #[stable(feature = "drain", since = "1.6.0")]
  52 unsafe impl<T: Send, A: Allocator + Send> Send for Drain<'_, T, A> {}
  53
  54 #[stable(feature = "drain", since = "1.6.0")]
  55 impl<T, A: Allocator> Drop for Drain<'_, T, A> {
  56     fn drop(&mut self) {
  57         struct DropGuard<'r, 'a, T, A: Allocator>(&'r mut Drain<'a, T, A>);
  58
  59         impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> {
  60             fn drop(&mut self) {
  61                 self.0.for_each(drop);
  62
  63                 let source_deque = unsafe { self.0.deque.as_mut() };
  64
  65                 // T = source_deque_tail; H = source_deque_head; t = drain_tail; h = drain_head
  66                 //
  67                 //        T   t   h   H
  68                 // [. . . o o x x o o . . .]
  69                 //
  70                 let orig_tail = source_deque.tail;
  71                 let drain_tail = source_deque.head;
  72                 let drain_head = self.0.after_tail;
  73                 let orig_head = self.0.after_head;
  74
  75                 let tail_len = count(orig_tail, drain_tail, source_deque.cap());
  76                 let head_len = count(drain_head, orig_head, source_deque.cap());
  77
  78                 // Restore the original head value
  79                 source_deque.head = orig_head;
  80
  81                 match (tail_len, head_len) {
  82                     (0, 0) => {
  83                         source_deque.head = 0;
  84                         source_deque.tail = 0;
  85                     }
  86                     (0, _) => {
  87                         source_deque.tail = drain_head;
  88                     }
  89                     (_, 0) => {
  90                         source_deque.head = drain_tail;
  91                     }
  92                     _ => unsafe {
  93                         if tail_len <= head_len {
  94                             source_deque.tail = source_deque.wrap_sub(drain_head, tail_len);
  95                             source_deque.wrap_copy(source_deque.tail, orig_tail, tail_len);
  96                         } else {
  97                             source_deque.head = source_deque.wrap_add(drain_tail, head_len);
  98                             source_deque.wrap_copy(drain_tail, drain_head, head_len);
  99                         }
 100                     },
 101                 }
 102             }
 103         }
 104
 105         while let Some(item) = self.next() {
 106             let guard = DropGuard(self);
 107             drop(item);
 108             mem::forget(guard);
 109         }
 110
 111         DropGuard(self);
 112     }
 113 }
 114
 115 #[stable(feature = "drain", since = "1.6.0")]
 116 impl<T, A: Allocator> Iterator for Drain<'_, T, A> {
 117     type Item = T;
 118
 119     #[inline]
 120     fn next(&mut self) -> Option<T> {
 121         self.iter.next().map(|elt| unsafe { ptr::read(elt) })
 122     }
 123
 124     #[inline]
 125     fn size_hint(&self) -> (usize, Option<usize>) {
 126         self.iter.size_hint()
 127     }
 128 }
 129
 130 #[stable(feature = "drain", since = "1.6.0")]
 131 impl<T, A: Allocator> DoubleEndedIterator for Drain<'_, T, A> {
 132     #[inline]
 133     fn next_back(&mut self) -> Option<T> {
 134         self.iter.next_back().map(|elt| unsafe { ptr::read(elt) })
 135     }
 136 }
 137
 138 #[stable(feature = "drain", since = "1.6.0")]
 139 impl<T, A: Allocator> ExactSizeIterator for Drain<'_, T, A> {}
 140
 141 #[stable(feature = "fused", since = "1.26.0")]
 142 impl<T, A: Allocator> FusedIterator for Drain<'_, T, A> {}