原文链接：https://doc.rust-lang.org/nomicon/vec-zsts.html

最终代码

  1. #![feature(ptr_internals)]
  2. #![feature(allocator_api)]
  4. use std::ptr::{Unique, NonNull, self};
  5. use std::mem;
  6. use std::ops::{Deref, DerefMut};
  7. use std::marker::PhantomData;
  8. use std::alloc::{Alloc, GlobalAlloc, Layout, Global, handle_alloc_error};
  10. struct RawVec<T> {
  11.     ptr: Unique<T>,
  12.     cap: usize,
  13. }
  15. impl<T> RawVec<T> {
  16.     fn new() -> Self {
  17.         // !0就是usize::MAX。这段分支代码在编译期就可以计算出结果。
  18.         let cap = if mem::size_of::<T>() == 0 { !0 } else { 0 };
  20.         // Unique::empty()有着“未分配”和“零尺寸分配”的双重含义
  21.         RawVec { ptr: Unique::empty(), cap: cap }
  22.     }
  24.     fn grow(&mut self) {
  25.         unsafe {
  26.             let elem_size = mem::size_of::<T>();
  28.             // 因为当elem_size为0时我们设置了cap为usize::MAX，
  29.             // 这一步成立意味着Vec的容量溢出了
  30.             assert!(elem_size != 0, "capacity overflow");
  32.             let (new_cap, ptr) = if self.cap == 0 {
  33.                 let ptr = Global.alloc(Layout::array::<T>(1).unwrap());
  34.                 (1, ptr)
  35.             } else {
  36.                 let new_cap = 2 * self.cap;
  37.                 let c: NonNull<T> = self.ptr.into();
  38.                 let ptr = Global.realloc(c.cast(),
  39.                                          Layout::array::<T>(self.cap).unwrap(),
  40.                                          Layout::array::<T>(new_cap).unwrap().size());
  41.                 (new_cap, ptr)
  42.             };
  44.             // 如果分配或再分配失败，oom
  45.             if ptr.is_err() {
  46.                 handle_alloc_error(Layout::from_size_align_unchecked(
  47.                     new_cap * elem_size,
  48.                     mem::align_of::<T>(),
  49.                 ))
  50.             }
  51.             let ptr = ptr.unwrap();
  53.             self.ptr = Unique::new_unchecked(ptr.as_ptr() as *mut _);
  54.             self.cap = new_cap;
  55.         }
  56.     }
  57. }
  59. impl<T> Drop for RawVec<T> {
  60.     fn drop(&mut self) {
  61.         let elem_size = mem::size_of::<T>();
  62.         if self.cap != 0 && elem_size != 0 {
  63.             unsafe {
  64.                 let c: NonNull<T> = self.ptr.into();
  65.                 Global.dealloc(c.cast(),
  66.                                Layout::array::<T>(self.cap).unwrap());
  67.             }
  68.         }
  69.     }
  70. }
  72. pub struct Vec<T> {
  73.     buf: RawVec<T>,
  74.     len: usize,
  75. }
  77. impl<T> Vec<T> {
  78.     fn ptr(&self) -> *mut T { self.buf.ptr.as_ptr() }
  80.     fn cap(&self) -> usize { self.buf.cap }
  82.     pub fn new() -> Self {
  83.         Vec { buf: RawVec::new(), len: 0 }
  84.     }
  85.     pub fn push(&mut self, elem: T) {
  86.         if self.len == self.cap() { self.buf.grow(); }
  88.         unsafe {
  89.             ptr::write(self.ptr().offset(self.len as isize), elem);
  90.         }
  92.         // 不会溢出，会先OOM
  93.         self.len += 1;
  94.     }
  96.     pub fn pop(&mut self) -> Option<T> {
  97.         if self.len == 0 {
  98.             None
  99.         } else {
  100.             self.len -= 1;
  101.             unsafe {
  102.                 Some(ptr::read(self.ptr().offset(self.len as isize)))
  103.             }
  104.         }
  105.     }
  107.     pub fn insert(&mut self, index: usize, elem: T) {
  108.         assert!(index <= self.len, "index out of bounds");
  109.         if self.cap() == self.len { self.buf.grow(); }
  111.         unsafe {
  112.             if index < self.len {
  113.                 ptr::copy(self.ptr().offset(index as isize),
  114.                           self.ptr().offset(index as isize + 1),
  115.                           self.len - index);
  116.             }
  117.             ptr::write(self.ptr().offset(index as isize), elem);
  118.             self.len += 1;
  119.         }
  120.     }
  122.     pub fn remove(&mut self, index: usize) -> T {
  123.         assert!(index < self.len, "index out of bounds");
  124.         unsafe {
  125.             self.len -= 1;
  126.             let result = ptr::read(self.ptr().offset(index as isize));
  127.             ptr::copy(self.ptr().offset(index as isize + 1),
  128.                       self.ptr().offset(index as isize),
  129.                       self.len - index);
  130.             result
  131.         }
  132.     }
  134.     pub fn into_iter(self) -> IntoIter<T> {
  135.         unsafe {
  136.             let iter = RawValIter::new(&self);
  137.             let buf = ptr::read(&self.buf);
  138.             mem::forget(self);
  140.             IntoIter {
  141.                 iter: iter,
  142.                 _buf: buf,
  143.             }
  144.         }
  145.     }
  147.     pub fn drain(&mut self) -> Drain<T> {
  148.         unsafe {
  149.             let iter = RawValIter::new(&self);
  151.             // 这一步是为了mem::forget的安全。如果Drain被forget，我们会泄露整个Vec的内容
  152.             // 同时，既然我们无论如何都会做这一步，为什么不现在做呢？
  153.             self.len = 0;
  155.             Drain {
  156.                 iter: iter,
  157.                 vec: PhantomData,
  158.             }
  159.         }
  160.     }
  161. }
  163. impl<T> Drop for Vec<T> {
  164.     fn drop(&mut self) {
  165.         while let Some(_) = self.pop() {}
  166.         // 分配由RawVec负责
  167.     }
  168. }
  170. impl<T> Deref for Vec<T> {
  171.     type Target = [T];
  172.     fn deref(&self) -> &[T] {
  173.         unsafe {
  174.             ::std::slice::from_raw_parts(self.ptr(), self.len)
  175.         }
  176.     }
  177. }
  179. impl<T> DerefMut for Vec<T> {
  180.     fn deref_mut(&mut self) -> &mut [T] {
  181.         unsafe {
  182.             ::std::slice::from_raw_parts_mut(self.ptr(), self.len)
  183.         }
  184.     }
  185. }
  191. struct RawValIter<T> {
  192.     start: *const T,
  193.     end: *const T,
  194. }
  196. impl<T> RawValIter<T> {
  197.     unsafe fn new(slice: &[T]) -> Self {
  198.         RawValIter {
  199.             start: slice.as_ptr(),
  200.             end: if mem::size_of::<T>() == 0 {
  201.                 ((slice.as_ptr() as usize) + slice.len()) as *const _
  202.             } else if slice.len() == 0 {
  203.                 slice.as_ptr()
  204.             } else {
  205.                 slice.as_ptr().offset(slice.len() as isize)
  206.             }
  207.         }
  208.     }
  209. }
  211. impl<T> Iterator for RawValIter<T> {
  212.     type Item = T;
  213.     fn next(&mut self) -> Option<T> {
  214.         if self.start == self.end {
  215.             None
  216.         } else {
  217.             unsafe {
  218.                 let result = ptr::read(self.start);
  219.                 self.start = if mem::size_of::<T>() == 0 {
  220.                     (self.start as usize + 1) as *const _
  221.                 } else {
  222.                     self.start.offset(1)
  223.                 };
  224.                 Some(result)
  225.             }
  226.         }
  227.     }
  229.     fn size_hint(&self) -> (usize, Option<usize>) {
  230.         let elem_size = mem::size_of::<T>();
  231.         let len = (self.end as usize - self.start as usize)
  232.                   / if elem_size == 0 { 1 } else { elem_size };
  233.         (len, Some(len))
  234.     }
  235. }
  237. impl<T> DoubleEndedIterator for RawValIter<T> {
  238.     fn next_back(&mut self) -> Option<T> {
  239.         if self.start == self.end {
  240.             None
  241.         } else {
  242.             unsafe {
  243.                 self.end = if mem::size_of::<T>() == 0 {
  244.                     (self.end as usize - 1) as *const _
  245.                 } else {
  246.                     self.end.offset(-1)
  247.                 };
  248.                 Some(ptr::read(self.end))
  249.             }
  250.         }
  251.     }
  252. }
  257. pub struct IntoIter<T> {
  258.     _buf: RawVec<T>, // 我们并不关心这个，只是需要它们保持分配空间不被销毁
  259.     iter: RawValIter<T>,
  260. }
  262. impl<T> Iterator for IntoIter<T> {
  263.     type Item = T;
  264.     fn next(&mut self) -> Option<T> { self.iter.next() }
  265.     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
  266. }
  268. impl<T> DoubleEndedIterator for IntoIter<T> {
  269.     fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
  270. }
  272. impl<T> Drop for IntoIter<T> {
  273.     fn drop(&mut self) {
  274.         for _ in &mut *self {}
  275.     }
  276. }
  281. pub struct Drain<'a, T: 'a> {
  282.     vec: PhantomData<&'a mut Vec<T>>,
  283.     iter: RawValIter<T>,
  284. }
  286. impl<'a, T> Iterator for Drain<'a, T> {
  287.     type Item = T;
  288.     fn next(&mut self) -> Option<T> { self.iter.next() }
  289.     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
  290. }
  292. impl<'a, T> DoubleEndedIterator for Drain<'a, T> {
  293.     fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
  294. }
  296. impl<'a, T> Drop for Drain<'a, T> {
  297.     fn drop(&mut self) {
  298.         // pre-drain the iter
  299.         for _ in &mut self.iter {}
  300.     }
  301. }
  303. # fn main() {
  304. #     tests::create_push_pop();
  305. #     tests::iter_test();
  306. #     tests::test_drain();
  307. #     tests::test_zst();
  308. #     println!("All tests finished OK");
  309. # }
  311. # mod tests {
  312. #     use super::*;
  313. #     pub fn create_push_pop() {
  314. #         let mut v = Vec::new();
  315. #         v.push(1);
  316. #         assert_eq!(1, v.len());
  317. #         assert_eq!(1, v[0]);
  318. #         for i in v.iter_mut() {
  319. #             *i += 1;
  320. #         }
  321. #         v.insert(0, 5);
  322. #         let x = v.pop();
  323. #         assert_eq!(Some(2), x);
  324. #         assert_eq!(1, v.len());
  325. #         v.push(10);
  326. #         let x = v.remove(0);
  327. #         assert_eq!(5, x);
  328. #         assert_eq!(1, v.len());
  329. #     }
  330. # 
  331. #     pub fn iter_test() {
  332. #         let mut v = Vec::new();
  333. #         for i in 0..10 {
  334. #             v.push(Box::new(i))
  335. #         }
  336. #         let mut iter = v.into_iter();
  337. #         let first = iter.next().unwrap();
  338. #         let last = iter.next_back().unwrap();
  339. #         drop(iter);
  340. #         assert_eq!(0, *first);
  341. #         assert_eq!(9, *last);
  342. #     }
  343. # 
  344. #     pub fn test_drain() {
  345. #         let mut v = Vec::new();
  346. #         for i in 0..10 {
  347. #             v.push(Box::new(i))
  348. #         }
  349. #         {
  350. #             let mut drain = v.drain();
  351. #             let first = drain.next().unwrap();
  352. #             let last = drain.next_back().unwrap();
  353. #             assert_eq!(0, *first);
  354. #             assert_eq!(9, *last);
  355. #         }
  356. #         assert_eq!(0, v.len());
  357. #         v.push(Box::new(1));
  358. #         assert_eq!(1, *v.pop().unwrap());
  359. #     }
  360. # 
  361. #     pub fn test_zst() {
  362. #         let mut v = Vec::new();
  363. #         for _i in 0..10 {
  364. #             v.push(())
  365. #         }
  366. # 
  367. #         let mut count = 0;
  368. # 
  369. #         for _ in v.into_iter() {
  370. #             count += 1
  371. #         }
  372. # 
  373. #         assert_eq!(10, count);
  374. #     }
  375. # }