1use core::{
2 fmt::Debug,
3 iter::FusedIterator,
4 mem::take,
5 ptr::{self, NonNull},
6 slice,
7};
8use crate::is_zst::IsZst;
9use rc_vec_proc_macro::rc_impl_gen_arc_impl;
10
11use super::{ArcVec, RcVec};
12
13#[rc_impl_gen_arc_impl]
14pub struct RcVecDrain<'a, T: 'a> {
15 pub(super) tail_start: usize,
16 pub(super) tail_len: usize,
17 pub(super) iter: slice::Iter<'a, T>,
18 pub(super) vec: NonNull<RcVec<T>>,
19}
20
21impl<'a, T: Debug + 'a> Debug for RcVecDrain<'a, T> {
22 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
23 f.debug_tuple("RcVecDrain")
24 .field(&self.iter.as_slice())
25 .finish()
26 }
27}
28
29impl<'a, T: Debug + 'a> Debug for ArcVecDrain<'a, T> {
30 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
31 f.debug_tuple("ArcVecDrain")
32 .field(&self.iter.as_slice())
33 .finish()
34 }
35}
36
37#[rc_impl_gen_arc_impl]
38impl<'a, T: 'a> RcVecDrain<'a, T> {
39 pub fn as_slice(&self) -> &[T] {
40 self.iter.as_slice()
41 }
42}
43
44#[rc_impl_gen_arc_impl]
45impl<'a, T: 'a> Iterator for RcVecDrain<'a, T> {
46 type Item = T;
47
48 fn next(&mut self) -> Option<Self::Item> {
49 self.iter.next().map(|ele| unsafe { ptr::read(ele) })
50 }
51
52 fn size_hint(&self) -> (usize, Option<usize>) {
53 self.iter.size_hint()
54 }
55}
56
57#[rc_impl_gen_arc_impl]
58impl<'a, T: 'a> DoubleEndedIterator for RcVecDrain<'a, T> {
59 fn next_back(&mut self) -> Option<Self::Item> {
60 self.iter.next_back().map(|ele| unsafe { ptr::read(ele) })
61 }
62}
63
64#[rc_impl_gen_arc_impl]
65struct RcVecMoveGuard<'r, 'a, T>(&'r mut RcVecDrain<'a, T>);
66
67#[rc_impl_gen_arc_impl]
68impl<'r, 'a, T> Drop for RcVecMoveGuard<'r, 'a, T> {
69 fn drop(&mut self) {
70 unsafe {
71 let src_vec = self.0.vec.as_mut();
72 let start = src_vec.len();
73 let tail = self.0.tail_start;
74
75 if tail != start {
76 let dest = src_vec.as_mut_ptr().add(start);
77 src_vec.as_ptr().add(tail)
78 .copy_to(dest, self.0.tail_len);
79 }
80
81 src_vec.set_len(start+self.0.tail_len);
82 }
83 }
84}
85
86#[rc_impl_gen_arc_impl]
87impl<'a, T: 'a> Drop for RcVecDrain<'a, T> {
88 fn drop(&mut self) {
89 let iter = take(&mut self.iter);
90 let drop_len = iter.len();
91
92 let mut vec = self.vec;
93
94 if T::ZST {
95 unsafe {
96 let vec = vec.as_mut();
97 let old_len = vec.len();
98
99 vec.set_len(old_len + drop_len + self.tail_len);
100 vec.truncate(old_len + self.tail_len);
101 }
102
103 return;
104 }
105
106 let _guard = RcVecMoveGuard(self);
107
108 if drop_len == 0 {
109 return;
110 }
111
112 let drop_ptr = iter.as_slice().as_ptr();
113
114 unsafe {
115 let vec_ptr = vec.as_mut().as_ptr();
116 let drop_offset = drop_ptr.offset_from(vec_ptr).try_into().unwrap();
117 let range = drop_offset..drop_offset+drop_len;
118 vec.as_mut().raw.drop_elems_from_range(range);
119 }
120 }
121}
122
123#[rc_impl_gen_arc_impl]
124impl<T> ExactSizeIterator for RcVecDrain<'_, T> { }
125
126#[rc_impl_gen_arc_impl]
127impl<T> FusedIterator for RcVecDrain<'_, T> { }