最终代码(The Final Code)

  1. use std::alloc::{self, Layout};
  2. use std::marker::PhantomData;
  3. use std::mem;
  4. use std::ops::{Deref, DerefMut};
  5. use std::ptr::{self, NonNull};
  7. struct RawVec<T> {
  8.     ptr: NonNull<T>,
  9.     cap: usize,
  10.     _marker: PhantomData<T>,
  11. }
  13. unsafe impl<T: Send> Send for RawVec<T> {}
  14. unsafe impl<T: Sync> Sync for RawVec<T> {}
  16. impl<T> RawVec<T> {
  17.     fn new() -> Self {
  18.         // !0 is usize::MAX. This branch should be stripped at compile time.
  19.         let cap = if mem::size_of::<T>() == 0 { !0 } else { 0 };
  21.         // `NonNull::dangling()` doubles as "unallocated" and "zero-sized allocation"
  22.         RawVec {
  23.             ptr: NonNull::dangling(),
  24.             cap: cap,
  25.             _marker: PhantomData,
  26.         }
  27.     }
  29.     fn grow(&mut self) {
  30.         // since we set the capacity to usize::MAX when T has size 0,
  31.         // getting to here necessarily means the Vec is overfull.
  32.         assert!(mem::size_of::<T>() != 0, "capacity overflow");
  34.         let (new_cap, new_layout) = if self.cap == 0 {
  35.             (1, Layout::array::<T>(1).unwrap())
  36.         } else {
  37.             // This can't overflow because we ensure self.cap <= isize::MAX.
  38.             let new_cap = 2 * self.cap;
  39.             // `Layout::array` checks that the number of bytes is <= usize::MAX,
  40.             // but this is redundant since old_layout.size() <= isize::MAX,
  41.             // so the `unwrap` should never fail.
  42.             let new_layout = Layout::array::<T>(new_cap).unwrap();
  43.             (new_cap, new_layout)
  44.         };
  46.         // Ensure that the new allocation doesn't exceed `isize::MAX` bytes.
  47.         assert!(
  48.             new_layout.size() <= isize::MAX as usize,
  49.             "Allocation too large"
  50.         );
  52.         let new_ptr = if self.cap == 0 {
  53.             unsafe { alloc::alloc(new_layout) }
  54.         } else {
  55.             let old_layout = Layout::array::<T>(self.cap).unwrap();
  56.             let old_ptr = self.ptr.as_ptr() as *mut u8;
  57.             unsafe { alloc::realloc(old_ptr, old_layout, new_layout.size()) }
  58.         };
  60.         // If allocation fails, `new_ptr` will be null, in which case we abort.
  61.         self.ptr = match NonNull::new(new_ptr as *mut T) {
  62.             Some(p) => p,
  63.             None => alloc::handle_alloc_error(new_layout),
  64.         };
  65.         self.cap = new_cap;
  66.     }
  67. }
  69. impl<T> Drop for RawVec<T> {
  70.     fn drop(&mut self) {
  71.         let elem_size = mem::size_of::<T>();
  73.         if self.cap != 0 && elem_size != 0 {
  74.             unsafe {
  75.                 alloc::dealloc(
  76.                     self.ptr.as_ptr() as *mut u8,
  77.                     Layout::array::<T>(self.cap).unwrap(),
  78.                 );
  79.             }
  80.         }
  81.     }
  82. }
  84. pub struct Vec<T> {
  85.     buf: RawVec<T>,
  86.     len: usize,
  87. }
  89. impl<T> Vec<T> {
  90.     fn ptr(&self) -> *mut T {
  91.         self.buf.ptr.as_ptr()
  92.     }
  94.     fn cap(&self) -> usize {
  95.         self.buf.cap
  96.     }
  98.     pub fn new() -> Self {
  99.         Vec {
  100.             buf: RawVec::new(),
  101.             len: 0,
  102.         }
  103.     }
  104.     pub fn push(&mut self, elem: T) {
  105.         if self.len == self.cap() {
  106.             self.buf.grow();
  107.         }
  109.         unsafe {
  110.             ptr::write(self.ptr().add(self.len), elem);
  111.         }
  113.         // Can't overflow, we'll OOM first.
  114.         self.len += 1;
  115.     }
  117.     pub fn pop(&mut self) -> Option<T> {
  118.         if self.len == 0 {
  119.             None
  120.         } else {
  121.             self.len -= 1;
  122.             unsafe { Some(ptr::read(self.ptr().add(self.len))) }
  123.         }
  124.     }
  126.     pub fn insert(&mut self, index: usize, elem: T) {
  127.         assert!(index <= self.len, "index out of bounds");
  128.         if self.cap() == self.len { self.buf.grow(); }
  130.         unsafe {
  131.             ptr::copy(
  132.                 self.ptr().add(index),
  133.                 self.ptr().add(index + 1),
  134.                 self.len - index,
  135.             );
  136.             ptr::write(self.ptr().add(index), elem);
  137.             self.len += 1;
  138.         }
  139.     }
  141.     pub fn remove(&mut self, index: usize) -> T {
  142.         assert!(index < self.len, "index out of bounds");
  143.         unsafe {
  144.             self.len -= 1;
  145.             let result = ptr::read(self.ptr().add(index));
  146.             ptr::copy(
  147.                 self.ptr().add(index + 1),
  148.                 self.ptr().add(index),
  149.                 self.len - index,
  150.             );
  151.             result
  152.         }
  153.     }
  155.     pub fn into_iter(self) -> IntoIter<T> {
  156.         unsafe {
  157.             let iter = RawValIter::new(&self);
  158.             let buf = ptr::read(&self.buf);
  159.             mem::forget(self);
  161.             IntoIter {
  162.                 iter: iter,
  163.                 _buf: buf,
  164.             }
  165.         }
  166.     }
  168.     pub fn drain(&mut self) -> Drain<T> {
  169.         unsafe {
  170.             let iter = RawValIter::new(&self);
  172.             // this is a mem::forget safety thing. If Drain is forgotten, we just
  173.             // leak the whole Vec's contents. Also we need to do this *eventually*
  174.             // anyway, so why not do it now?
  175.             self.len = 0;
  177.             Drain {
  178.                 iter: iter,
  179.                 vec: PhantomData,
  180.             }
  181.         }
  182.     }
  183. }
  185. impl<T> Drop for Vec<T> {
  186.     fn drop(&mut self) {
  187.         while let Some(_) = self.pop() {}
  188.         // deallocation is handled by RawVec
  189.     }
  190. }
  192. impl<T> Deref for Vec<T> {
  193.     type Target = [T];
  194.     fn deref(&self) -> &[T] {
  195.         unsafe { std::slice::from_raw_parts(self.ptr(), self.len) }
  196.     }
  197. }
  199. impl<T> DerefMut for Vec<T> {
  200.     fn deref_mut(&mut self) -> &mut [T] {
  201.         unsafe { std::slice::from_raw_parts_mut(self.ptr(), self.len) }
  202.     }
  203. }
  205. struct RawValIter<T> {
  206.     start: *const T,
  207.     end: *const T,
  208. }
  210. impl<T> RawValIter<T> {
  211.     unsafe fn new(slice: &[T]) -> Self {
  212.         RawValIter {
  213.             start: slice.as_ptr(),
  214.             end: if mem::size_of::<T>() == 0 {
  215.                 ((slice.as_ptr() as usize) + slice.len()) as *const _
  216.             } else if slice.len() == 0 {
  217.                 slice.as_ptr()
  218.             } else {
  219.                 slice.as_ptr().add(slice.len())
  220.             },
  221.         }
  222.     }
  223. }
  225. impl<T> Iterator for RawValIter<T> {
  226.     type Item = T;
  227.     fn next(&mut self) -> Option<T> {
  228.         if self.start == self.end {
  229.             None
  230.         } else {
  231.             unsafe {
  232.                 let result = ptr::read(self.start);
  233.                 self.start = if mem::size_of::<T>() == 0 {
  234.                     (self.start as usize + 1) as *const _
  235.                 } else {
  236.                     self.start.offset(1)
  237.                 };
  238.                 Some(result)
  239.             }
  240.         }
  241.     }
  243.     fn size_hint(&self) -> (usize, Option<usize>) {
  244.         let elem_size = mem::size_of::<T>();
  245.         let len = (self.end as usize - self.start as usize) /
  246.                   if elem_size == 0 { 1 } else { elem_size };
  247.         (len, Some(len))
  248.     }
  249. }
  251. impl<T> DoubleEndedIterator for RawValIter<T> {
  252.     fn next_back(&mut self) -> Option<T> {
  253.         if self.start == self.end {
  254.             None
  255.         } else {
  256.             unsafe {
  257.                 self.end = if mem::size_of::<T>() == 0 {
  258.                     (self.end as usize - 1) as *const _
  259.                 } else {
  260.                     self.end.offset(-1)
  261.                 };
  262.                 Some(ptr::read(self.end))
  263.             }
  264.         }
  265.     }
  266. }
  268. pub struct IntoIter<T> {
  269.     _buf: RawVec<T>, // we don't actually care about this. Just need it to live.
  270.     iter: RawValIter<T>,
  271. }
  273. impl<T> Iterator for IntoIter<T> {
  274.     type Item = T;
  275.     fn next(&mut self) -> Option<T> {
  276.         self.iter.next()
  277.     }
  278.     fn size_hint(&self) -> (usize, Option<usize>) {
  279.         self.iter.size_hint()
  280.     }
  281. }
  283. impl<T> DoubleEndedIterator for IntoIter<T> {
  284.     fn next_back(&mut self) -> Option<T> {
  285.         self.iter.next_back()
  286.     }
  287. }
  289. impl<T> Drop for IntoIter<T> {
  290.     fn drop(&mut self) {
  291.         for _ in &mut *self {}
  292.     }
  293. }
  295. pub struct Drain<'a, T: 'a> {
  296.     vec: PhantomData<&'a mut Vec<T>>,
  297.     iter: RawValIter<T>,
  298. }
  300. impl<'a, T> Iterator for Drain<'a, T> {
  301.     type Item = T;
  302.     fn next(&mut self) -> Option<T> {
  303.         self.iter.next()
  304.     }
  305.     fn size_hint(&self) -> (usize, Option<usize>) {
  306.         self.iter.size_hint()
  307.     }
  308. }
  310. impl<'a, T> DoubleEndedIterator for Drain<'a, T> {
  311.     fn next_back(&mut self) -> Option<T> {
  312.         self.iter.next_back()
  313.     }
  314. }
  316. impl<'a, T> Drop for Drain<'a, T> {
  317.     fn drop(&mut self) {
  318.         // pre-drain the iter
  319.         for _ in &mut *self {}
  320.     }
  321. }
  323. #
  324. # fn main() {
  325. #     tests::create_push_pop();
  326. #     tests::iter_test();
  327. #     tests::test_drain();
  328. #     tests::test_zst();
  329. #     println!("All tests finished OK");
  330. # }
  331. #
  332. # mod tests {
  333. #     use super::*;
  334. #
  335. #     pub fn create_push_pop() {
  336. #         let mut v = Vec::new();
  337. #         v.push(1);
  338. #         assert_eq!(1, v.len());
  339. #         assert_eq!(1, v[0]);
  340. #         for i in v.iter_mut() {
  341. #             *i += 1;
  342. #         }
  343. #         v.insert(0, 5);
  344. #         let x = v.pop();
  345. #         assert_eq!(Some(2), x);
  346. #         assert_eq!(1, v.len());
  347. #         v.push(10);
  348. #         let x = v.remove(0);
  349. #         assert_eq!(5, x);
  350. #         assert_eq!(1, v.len());
  351. #     }
  352. #
  353. #     pub fn iter_test() {
  354. #         let mut v = Vec::new();
  355. #         for i in 0..10 {
  356. #             v.push(Box::new(i))
  357. #         }
  358. #         let mut iter = v.into_iter();
  359. #         let first = iter.next().unwrap();
  360. #         let last = iter.next_back().unwrap();
  361. #         drop(iter);
  362. #         assert_eq!(0, *first);
  363. #         assert_eq!(9, *last);
  364. #     }
  365. #
  366. #     pub fn test_drain() {
  367. #         let mut v = Vec::new();
  368. #         for i in 0..10 {
  369. #             v.push(Box::new(i))
  370. #         }
  371. #         {
  372. #             let mut drain = v.drain();
  373. #             let first = drain.next().unwrap();
  374. #             let last = drain.next_back().unwrap();
  375. #             assert_eq!(0, *first);
  376. #             assert_eq!(9, *last);
  377. #         }
  378. #         assert_eq!(0, v.len());
  379. #         v.push(Box::new(1));
  380. #         assert_eq!(1, *v.pop().unwrap());
  381. #     }
  382. #
  383. #     pub fn test_zst() {
  384. #         let mut v = Vec::new();
  385. #         for _i in 0..10 {
  386. #             v.push(())
  387. #         }
  388. #
  389. #         let mut count = 0;
  390. #
  391. #         for _ in v.into_iter() {
  392. #             count += 1
  393. #         }
  394. #
  395. #         assert_eq!(10, count);
  396. #     }
  397. # }