1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

use core::iter::FusedIterator;
use core::ptr::{self, NonNull};
use core::{fmt, mem};

use super::{count, Iter, VecDeque};

/// `VecDeque` 的元素上的 draining 迭代器。
///
/// 该 `struct` 是通过 [`VecDeque`] 上的 [`drain`] 方法创建的。
/// 有关更多信息，请参见其文档。
///
/// [`drain`]: VecDeque::drain
#[stable(feature = "drain", since = "1.6.0")]
pub struct Drain<'a, T: 'a> {
    pub(crate) after_tail: usize,
    pub(crate) after_head: usize,
    pub(crate) iter: Iter<'a, T>,
    pub(crate) deque: NonNull<VecDeque<T>>,
}

#[stable(feature = "collection_debug", since = "1.17.0")]
impl<T: fmt::Debug> fmt::Debug for Drain<'_, T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("Drain")
            .field(&self.after_tail)
            .field(&self.after_head)
            .field(&self.iter)
            .finish()
    }
}

#[stable(feature = "drain", since = "1.6.0")]
unsafe impl<T: Sync> Sync for Drain<'_, T> {}
#[stable(feature = "drain", since = "1.6.0")]
unsafe impl<T: Send> Send for Drain<'_, T> {}

#[stable(feature = "drain", since = "1.6.0")]
impl<T> Drop for Drain<'_, T> {
    fn drop(&mut self) {
        struct DropGuard<'r, 'a, T>(&'r mut Drain<'a, T>);

        impl<'r, 'a, T> Drop for DropGuard<'r, 'a, T> {
            fn drop(&mut self) {
                self.0.for_each(drop);

                let source_deque = unsafe { self.0.deque.as_mut() };

                // T = source_deque_tail; H = source_deque_head; t = drain_tail; h = drain_head
                //
                //        T   t   h   H
                // [. . . o o x x o o . . .]
                //
                let orig_tail = source_deque.tail;
                let drain_tail = source_deque.head;
                let drain_head = self.0.after_tail;
                let orig_head = self.0.after_head;

                let tail_len = count(orig_tail, drain_tail, source_deque.cap());
                let head_len = count(drain_head, orig_head, source_deque.cap());

                // 恢复原始 head 值
                source_deque.head = orig_head;

                match (tail_len, head_len) {
                    (0, 0) => {
                        source_deque.head = 0;
                        source_deque.tail = 0;
                    }
                    (0, _) => {
                        source_deque.tail = drain_head;
                    }
                    (_, 0) => {
                        source_deque.head = drain_tail;
                    }
                    _ => unsafe {
                        if tail_len <= head_len {
                            source_deque.tail = source_deque.wrap_sub(drain_head, tail_len);
                            source_deque.wrap_copy(source_deque.tail, orig_tail, tail_len);
                        } else {
                            source_deque.head = source_deque.wrap_add(drain_tail, head_len);
                            source_deque.wrap_copy(drain_tail, drain_head, head_len);
                        }
                    },
                }
            }
        }

        while let Some(item) = self.next() {
            let guard = DropGuard(self);
            drop(item);
            mem::forget(guard);
        }

        DropGuard(self);
    }
}

#[stable(feature = "drain", since = "1.6.0")]
impl<T> Iterator for Drain<'_, T> {
    type Item = T;

    #[inline]
    fn next(&mut self) -> Option<T> {
        self.iter.next().map(|elt| unsafe { ptr::read(elt) })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

#[stable(feature = "drain", since = "1.6.0")]
impl<T> DoubleEndedIterator for Drain<'_, T> {
    #[inline]
    fn next_back(&mut self) -> Option<T> {
        self.iter.next_back().map(|elt| unsafe { ptr::read(elt) })
    }
}

#[stable(feature = "drain", since = "1.6.0")]
impl<T> ExactSizeIterator for Drain<'_, T> {}

#[stable(feature = "fused", since = "1.26.0")]
impl<T> FusedIterator for Drain<'_, T> {}