最终的代码

  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 等价于 usize::MAX， 这一段分支代码在编译期间就可以计算出结果返回的结果，返回给 cap
  19.         let cap = if mem::size_of::<T>() == 0 { !0 } else { 0 };
  21.         // `NonNull::dangling()` 有双重含义:
  22.         // `未分配内存 (unallocated)`, `零尺寸 (zero-sized allocation)`
  23.         RawVec {
  24.             ptr: NonNull::dangling(),
  25.             cap: cap,
  26.             _marker: PhantomData,
  27.         }
  28.     }
  30.     fn grow(&mut self) {
  31.         // 因为当 T 的尺寸为 0 时，我们设置了 cap 为 usize::MAX，
  32.         // 这一步成立便意味着 Vec 溢出了.
  33.         assert!(mem::size_of::<T>() != 0, "capacity overflow");
  35.         let (new_cap, new_layout) = if self.cap == 0 {
  36.             (1, Layout::array::<T>(1).unwrap())
  37.         } else {
  38.             // 保证新申请的内存没有超出 `isize::MAX` 字节
  39.             let new_cap = 2 * self.cap;
  41.             // `Layout::array` 会检查申请的空间是否小于等于 usize::MAX，
  42.             // 但是因为 old_layout.size() <= isize::MAX，
  43.             // 所以这里的 unwrap 永远不可能失败
  44.             let new_layout = Layout::array::<T>(new_cap).unwrap();
  45.             (new_cap, new_layout)
  46.         };
  48.         // 保证新申请的内存没有超出 `isize::MAX` 字节
  49.         assert!(
  50.             new_layout.size() <= isize::MAX as usize,
  51.             "Allocation too large"
  52.         );
  54.         let new_ptr = if self.cap == 0 {
  55.             unsafe { alloc::alloc(new_layout) }
  56.         } else {
  57.             let old_layout = Layout::array::<T>(self.cap).unwrap();
  58.             let old_ptr = self.ptr.as_ptr() as *mut u8;
  59.             unsafe { alloc::realloc(old_ptr, old_layout, new_layout.size()) }
  60.         };
  62.         // 如果分配失败，`new_ptr` 就会成为空指针，我们需要处理该意外情况
  63.         self.ptr = match NonNull::new(new_ptr as *mut T) {
  64.             Some(p) => p,
  65.             None => alloc::handle_alloc_error(new_layout),
  66.         };
  67.         self.cap = new_cap;
  68.     }
  69. }
  71. impl<T> Drop for RawVec<T> {
  72.     fn drop(&mut self) {
  73.         let elem_size = mem::size_of::<T>();
  75.         if self.cap != 0 && elem_size != 0 {
  76.             unsafe {
  77.                 alloc::dealloc(
  78.                     self.ptr.as_ptr() as *mut u8,
  79.                     Layout::array::<T>(self.cap).unwrap(),
  80.                 );
  81.             }
  82.         }
  83.     }
  84. }
  86. pub struct Vec<T> {
  87.     buf: RawVec<T>,
  88.     len: usize,
  89. }
  91. impl<T> Vec<T> {
  92.     fn ptr(&self) -> *mut T {
  93.         self.buf.ptr.as_ptr()
  94.     }
  96.     fn cap(&self) -> usize {
  97.         self.buf.cap
  98.     }
  100.     pub fn new() -> Self {
  101.         Vec {
  102.             buf: RawVec::new(),
  103.             len: 0,
  104.         }
  105.     }
  106.     pub fn push(&mut self, elem: T) {
  107.         if self.len == self.cap() {
  108.             self.buf.grow();
  109.         }
  111.         unsafe {
  112.             ptr::write(self.ptr().add(self.len), elem);
  113.         }
  115.         // 不会溢出，会先 OOM
  116.         self.len += 1;
  117.     }
  119.     pub fn pop(&mut self) -> Option<T> {
  120.         if self.len == 0 {
  121.             None
  122.         } else {
  123.             self.len -= 1;
  124.             unsafe { Some(ptr::read(self.ptr().add(self.len))) }
  125.         }
  126.     }
  128.     pub fn insert(&mut self, index: usize, elem: T) {
  129.         assert!(index <= self.len, "index out of bounds");
  130.         if self.cap() == self.len {
  131.             self.buf.grow();
  132.         }
  134.         unsafe {
  135.             ptr::copy(
  136.                 self.ptr().add(index),
  137.                 self.ptr().add(index + 1),
  138.                 self.len - index,
  139.             );
  140.             ptr::write(self.ptr().add(index), elem);
  141.             self.len += 1;
  142.         }
  143.     }
  145.     pub fn remove(&mut self, index: usize) -> T {
  146.         assert!(index < self.len, "index out of bounds");
  147.         unsafe {
  148.             self.len -= 1;
  149.             let result = ptr::read(self.ptr().add(index));
  150.             ptr::copy(
  151.                 self.ptr().add(index + 1),
  152.                 self.ptr().add(index),
  153.                 self.len - index,
  154.             );
  155.             result
  156.         }
  157.     }
  159.     pub fn drain(&mut self) -> Drain<T> {
  160.         unsafe {
  161.             let iter = RawValIter::new(&self);
  163.             // 这里事关 mem::forget 的安全。
  164.             // 如果 Drain 被 forget，我们就会泄露整个 Vec 的内存
  165.             // 既然我们始终要做这一步，为何不在这里完成呢？
  166.             self.len = 0;
  168.             Drain {
  169.                 iter: iter,
  170.                 vec: PhantomData,
  171.             }
  172.         }
  173.     }
  174. }
  176. impl<T> Drop for Vec<T> {
  177.     fn drop(&mut self) {
  178.         while let Some(_) = self.pop() {}
  179.         // RawVec 来负责释放内存
  180.     }
  181. }
  183. impl<T> Deref for Vec<T> {
  184.     type Target = [T];
  185.     fn deref(&self) -> &[T] {
  186.         unsafe { std::slice::from_raw_parts(self.ptr(), self.len) }
  187.     }
  188. }
  190. impl<T> DerefMut for Vec<T> {
  191.     fn deref_mut(&mut self) -> &mut [T] {
  192.         unsafe { std::slice::from_raw_parts_mut(self.ptr(), self.len) }
  193.     }
  194. }
  196. impl<T> IntoIterator for Vec<T> {
  197.     type Item = T;
  198.     type IntoIter = IntoIter<T>;
  199.     fn into_iter(self) -> IntoIter<T> {
  200.         unsafe {
  201.             let iter = RawValIter::new(&self);
  202.             let buf = ptr::read(&self.buf);
  203.             mem::forget(self);
  204.             IntoIter {
  205.                 iter: iter,
  206.                 _buf: buf,
  207.             }
  208.         }
  209.     }
  210. }
  212. struct RawValIter<T> {
  213.     start: *const T,
  214.     end: *const T,
  215. }
  217. impl<T> RawValIter<T> {
  218.     unsafe fn new(slice: &[T]) -> Self {
  219.         RawValIter {
  220.             start: slice.as_ptr(),
  221.             end: if mem::size_of::<T>() == 0 {
  222.                 ((slice.as_ptr() as usize) + slice.len()) as *const _
  223.             } else if slice.len() == 0 {
  224.                 slice.as_ptr()
  225.             } else {
  226.                 slice.as_ptr().add(slice.len())
  227.             },
  228.         }
  229.     }
  230. }
  232. impl<T> Iterator for RawValIter<T> {
  233.     type Item = T;
  234.     fn next(&mut self) -> Option<T> {
  235.         if self.start == self.end {
  236.             None
  237.         } else {
  238.             unsafe {
  239.                 if mem::size_of::<T>() == 0 {
  240.                     self.start = (self.start as usize + 1) as *const _;
  241.                     Some(ptr::read(NonNull::<T>::dangling().as_ptr()))
  242.                 } else {
  243.                     let old_ptr = self.start;
  244.                     self.start = self.start.offset(1);
  245.                     Some(ptr::read(old_ptr))
  246.                 }
  247.             }
  248.         }
  249.     }
  251.     fn size_hint(&self) -> (usize, Option<usize>) {
  252.         let elem_size = mem::size_of::<T>();
  253.         let len = (self.end as usize - self.start as usize)
  254.                   / if elem_size == 0 { 1 } else { elem_size };
  255.         (len, Some(len))
  256.     }
  257. }
  259. impl<T> DoubleEndedIterator for RawValIter<T> {
  260.     fn next_back(&mut self) -> Option<T> {
  261.         if self.start == self.end {
  262.             None
  263.         } else {
  264.             unsafe {
  265.                 if mem::size_of::<T>() == 0 {
  266.                     self.end = (self.end as usize - 1) as *const _;
  267.                     Some(ptr::read(NonNull::<T>::dangling().as_ptr()))
  268.                 } else {
  269.                     self.end = self.end.offset(-1);
  270.                     Some(ptr::read(self.end))
  271.                 }
  272.             }
  273.         }
  274.     }
  275. }
  277. pub struct IntoIter<T> {
  278.     _buf: RawVec<T>,    // 我们实际上并不关心这个，只需要他们保证分配的空间不被释放
  279.     iter: RawValIter<T>,
  280. }
  282. impl<T> Iterator for IntoIter<T> {
  283.     type Item = T;
  284.     fn next(&mut self) -> Option<T> {
  285.         self.iter.next()
  286.     }
  287.     fn size_hint(&self) -> (usize, Option<usize>) {
  288.         self.iter.size_hint()
  289.     }
  290. }
  292. impl<T> DoubleEndedIterator for IntoIter<T> {
  293.     fn next_back(&mut self) -> Option<T> {
  294.         self.iter.next_back()
  295.     }
  296. }
  298. impl<T> Drop for IntoIter<T> {
  299.     fn drop(&mut self) {
  300.         for _ in &mut *self {}
  301.     }
  302. }
  304. pub struct Drain<'a, T: 'a> {
  305.     vec: PhantomData<&'a mut Vec<T>>,
  306.     iter: RawValIter<T>,
  307. }
  309. impl<'a, T> Iterator for Drain<'a, T> {
  310.     type Item = T;
  311.     fn next(&mut self) -> Option<T> {
  312.         self.iter.next()
  313.     }
  314.     fn size_hint(&self) -> (usize, Option<usize>) {
  315.         self.iter.size_hint()
  316.     }
  317. }
  319. impl<'a, T> DoubleEndedIterator for Drain<'a, T> {
  320.     fn next_back(&mut self) -> Option<T> {
  321.         self.iter.next_back()
  322.     }
  323. }
  325. impl<'a, T> Drop for Drain<'a, T> {
  326.     fn drop(&mut self) {
  327.         // 消耗drain
  328.         for _ in &mut *self {}
  329.     }
  330. }
  331. #
  332. # fn main() {
  333. #     tests::create_push_pop();
  334. #     tests::iter_test();
  335. #     tests::test_drain();
  336. #     tests::test_zst();
  337. #     println!("All tests finished OK");
  338. # }
  339. #
  340. # mod tests {
  341. #     use super::*;
  342. #
  343. #     pub fn create_push_pop() {
  344. #         let mut v = Vec::new();
  345. #         v.push(1);
  346. #         assert_eq!(1, v.len());
  347. #         assert_eq!(1, v[0]);
  348. #         for i in v.iter_mut() {
  349. #             *i += 1;
  350. #         }
  351. #         v.insert(0, 5);
  352. #         let x = v.pop();
  353. #         assert_eq!(Some(2), x);
  354. #         assert_eq!(1, v.len());
  355. #         v.push(10);
  356. #         let x = v.remove(0);
  357. #         assert_eq!(5, x);
  358. #         assert_eq!(1, v.len());
  359. #     }
  360. #
  361. #     pub fn iter_test() {
  362. #         let mut v = Vec::new();
  363. #         for i in 0..10 {
  364. #             v.push(Box::new(i))
  365. #         }
  366. #         let mut iter = v.into_iter();
  367. #         let first = iter.next().unwrap();
  368. #         let last = iter.next_back().unwrap();
  369. #         drop(iter);
  370. #         assert_eq!(0, *first);
  371. #         assert_eq!(9, *last);
  372. #     }
  373. #
  374. #     pub fn test_drain() {
  375. #         let mut v = Vec::new();
  376. #         for i in 0..10 {
  377. #             v.push(Box::new(i))
  378. #         }
  379. #         {
  380. #             let mut drain = v.drain();
  381. #             let first = drain.next().unwrap();
  382. #             let last = drain.next_back().unwrap();
  383. #             assert_eq!(0, *first);
  384. #             assert_eq!(9, *last);
  385. #         }
  386. #         assert_eq!(0, v.len());
  387. #         v.push(Box::new(1));
  388. #         assert_eq!(1, *v.pop().unwrap());
  389. #     }
  390. #
  391. #     pub fn test_zst() {
  392. #         let mut v = Vec::new();
  393. #         for _i in 0..10 {
  394. #             v.push(())
  395. #         }
  396. #
  397. #         let mut count = 0;
  398. #
  399. #         for _ in v.into_iter() {
  400. #             count += 1
  401. #         }
  402. #
  403. #         assert_eq!(10, count);
  404. #     }
  405. # }