rc_vec/rc_vec/
drain.rs

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> { }