rust_examples/rustonomicon/
vec.rs

1use std::alloc::{self, Layout};
2use std::marker::PhantomData;
3use std::mem::{self};
4use std::ops::{Deref, DerefMut};
5use std::ptr::{self, NonNull};
6
7struct Vec<T> {
8    raw: RawVec<T>,
9    len: usize,
10}
11unsafe impl<T: Send> Send for Vec<T> {}
12unsafe impl<T: Sync> Sync for Vec<T> {}
13impl<T> Vec<T> {
14    pub fn new() -> Self {
15        Vec {
16            raw: RawVec::new(),
17            len: 0,
18        }
19    }
20    pub fn cap(&self) -> usize {
21        self.raw.cap
22    }
23    pub fn ptr(&self) -> *mut T {
24        self.raw.ptr.as_ptr()
25    }
26
27    pub fn push(&mut self, elem: T) {
28        if self.len == self.cap() {
29            self.raw.grow();
30        }
31        unsafe { self.ptr().add(self.len).write(elem) }
32        self.len += 1;
33    }
34    pub fn pop(&mut self) -> Option<T> {
35        if self.len == 0 {
36            None
37        } else {
38            self.len -= 1;
39            unsafe { Some(self.ptr().add(self.len).read()) }
40        }
41    }
42    pub fn insert(&mut self, index: usize, elem: T) {
43        assert!(index <= self.len, "index out of bounds");
44        if self.len == self.cap() {
45            self.raw.grow();
46        }
47        unsafe {
48            ptr::copy(
49                self.ptr().add(index),
50                self.ptr().add(index + 1),
51                self.len - index,
52            );
53            ptr::write(self.ptr().add(index), elem);
54        }
55        self.len += 1;
56    }
57    pub fn remove(&mut self, index: usize) -> T {
58        assert!(index < self.len, "index out of bounds");
59        unsafe {
60            self.len -= 1;
61            let result = ptr::read(self.ptr().add(index));
62            ptr::copy(
63                self.ptr().add(index + 1),
64                self.ptr().add(index),
65                self.len - index,
66            );
67            result
68        }
69    }
70    pub fn drain(&mut self) -> Drain<T> {
71        let iter = unsafe { RawValIter::new(&self) };
72        //    drop()的实际逻辑将由Drain做，所以self虽然还存在，但是已不能再做pop()操作。
73        // 因为lifetime绑定，所以self和Drain的销毁时机是一致的。
74        // 因为mut，self本身不能再被其他人使用
75        self.len = 0;
76        Drain {
77            vec: PhantomData,
78            iter,
79        }
80    }
81}
82impl<T> Drop for Vec<T> {
83    fn drop(&mut self) {
84        while let Some(_) = self.pop() {}
85    }
86}
87impl<T> Deref for Vec<T> {
88    type Target = [T];
89
90    fn deref(&self) -> &Self::Target {
91        unsafe { std::slice::from_raw_parts(self.ptr(), self.len) }
92    }
93}
94impl<T> DerefMut for Vec<T> {
95    fn deref_mut(&mut self) -> &mut Self::Target {
96        unsafe { std::slice::from_raw_parts_mut(self.ptr(), self.len) }
97    }
98}
99struct RawValIter<T> {
100    start: *const T,
101    end: *const T,
102}
103impl<T> RawValIter<T> {
104    unsafe fn new(slice: &[T]) -> RawValIter<T> {
105        RawValIter {
106            start: slice.as_ptr(),
107            end: if mem::size_of::<T>() == 0 {
108                // 因为ZST，所以end如果用T的size来计算则和start一致，也就是说丢失了len信息
109                let result = (slice.as_ptr() as usize + slice.len()) as *const _;
110                result
111            } else if slice.len() == 0 {
112                slice.as_ptr()
113            } else {
114                slice.as_ptr().add(slice.len())
115            },
116        }
117    }
118}
119struct IntoIter<T> {
120    // 用于内存回收
121    _buffer: RawVec<T>,
122    iter: RawValIter<T>,
123}
124impl<T> Drop for IntoIter<T> {
125    fn drop(&mut self) {
126        for _ in &mut *self {}
127    }
128}
129impl<T> IntoIterator for Vec<T> {
130    type Item = T;
131
132    type IntoIter = IntoIter<T>;
133
134    fn into_iter(self) -> Self::IntoIter {
135        unsafe {
136            IntoIter {
137                _buffer: ptr::read(&self.raw),
138                iter: RawValIter::new(&self),
139            }
140        }
141    }
142}
143impl<T> Iterator for RawValIter<T> {
144    type Item = T;
145
146    fn next(&mut self) -> Option<Self::Item> {
147        unsafe {
148            if self.start == self.end {
149                None
150            } else {
151                unsafe {
152                    if mem::size_of::<T>() == 0 {
153                        self.start = (self.start as usize + 1) as *const _;
154                        Some(ptr::read(NonNull::<T>::dangling().as_ptr()))
155                    } else {
156                        let result = Some(ptr::read(self.start));
157                        self.start = self.start.offset(1);
158                        result
159                    }
160                }
161            }
162        }
163    }
164    fn size_hint(&self) -> (usize, Option<usize>) {
165        let element_size = mem::size_of::<T>();
166        let len = (self.end as usize - self.start as usize)
167            / if element_size == 0 { 1 } else { element_size };
168        (len, Some(len))
169    }
170}
171impl<T> Iterator for IntoIter<T> {
172    type Item = T;
173
174    fn next(&mut self) -> Option<Self::Item> {
175        self.iter.next()
176    }
177    fn size_hint(&self) -> (usize, Option<usize>) {
178        self.iter.size_hint()
179    }
180}
181impl<T> DoubleEndedIterator for RawValIter<T> {
182    fn next_back(&mut self) -> Option<Self::Item> {
183        if self.start == self.end {
184            None
185        } else {
186            unsafe {
187                if mem::size_of::<T>() == 0 {
188                    self.end = (self.end as usize - 1) as *const _;
189                    Some(ptr::read(NonNull::<T>::dangling().as_ptr()))
190                } else {
191                    self.end = self.end.offset(-1);
192                    let result = ptr::read(self.end);
193                    Some(result)
194                }
195            }
196        }
197    }
198}
199impl<T> DoubleEndedIterator for IntoIter<T> {
200    fn next_back(&mut self) -> Option<Self::Item> {
201        self.iter.next_back()
202    }
203}
204struct RawVec<T> {
205    ptr: NonNull<T>,
206    cap: usize,
207}
208unsafe impl<T: Send> Send for RawVec<T> {}
209unsafe impl<T: Sync> Sync for RawVec<T> {}
210impl<T> RawVec<T> {
211    fn grow(&mut self) {
212        assert!(mem::size_of::<T>() != 0, "capacity overflow");
213        let new_cap = if self.cap == 0 { 1 } else { self.cap * 2 };
214        let new_layout = Layout::array::<T>(new_cap).unwrap();
215        assert!(
216            new_layout.size() < i32::MAX as usize,
217            "Allocation too large"
218        );
219        let new_ptr = if self.cap == 0 {
220            unsafe { alloc::alloc(new_layout) }
221        } else {
222            let old_layout = Layout::array::<T>(self.cap).unwrap();
223            unsafe { alloc::realloc(self.ptr.as_ptr() as *mut u8, old_layout, new_cap) }
224        };
225        self.ptr = match NonNull::new(new_ptr as *mut T) {
226            Some(ptr) => ptr,
227            None => alloc::handle_alloc_error(new_layout),
228        };
229        self.cap = new_cap;
230    }
231
232    fn new() -> RawVec<T> {
233        // assert!(mem::size_of::<T>() != 0, "We're not ready to handle ZSTs");
234        let cap = if mem::size_of::<T>() == 0 {
235            usize::MAX
236        } else {
237            0
238        };
239        RawVec {
240            ptr: NonNull::dangling(),
241            cap,
242        }
243    }
244}
245impl<T> Drop for RawVec<T> {
246    fn drop(&mut self) {
247        if self.cap != 0 {
248            let layout = Layout::array::<T>(self.cap).unwrap();
249            unsafe {
250                alloc::dealloc(self.ptr.as_ptr() as *mut u8, layout);
251            }
252        }
253    }
254}
255pub struct Drain<'a, T> {
256    vec: PhantomData<&'a mut T>,
257    iter: RawValIter<T>,
258}
259impl<'a, T> Iterator for Drain<'a, T> {
260    type Item = T;
261
262    fn next(&mut self) -> Option<Self::Item> {
263        self.iter.next()
264    }
265    fn size_hint(&self) -> (usize, Option<usize>) {
266        self.iter.size_hint()
267    }
268}
269impl<'a, T> Drop for Drain<'a, T> {
270    fn drop(&mut self) {
271        for _ in &mut *self {}
272    }
273}
274#[cfg(test)]
275mod test {
276    #[test]
277    fn basic() {
278        let mut vec = Vec::<i32>::new();
279        assert_eq!(vec.pop(), None);
280        vec.push(1);
281        vec.push(2);
282        vec.insert(0, 3);
283        let vec_slice = &vec;
284        assert_eq!(vec_slice.len(), 3);
285        assert_eq!(vec_slice[0], 3);
286        assert_eq!(vec.remove(0), 3);
287        let vec_slice = &vec;
288        assert_eq!(vec_slice.len(), 2);
289        assert_eq!(vec_slice[0], 1);
290        assert_eq!(vec_slice[1], 2);
291        assert_eq!(vec.len(), 2);
292        assert_eq!(vec.pop(), Some(2));
293        assert_eq!(vec.pop(), Some(1));
294    }
295}
rust_examples/rustonomicon/vec.rs

rust_examples/rustonomicon/
vec.rs
