RawVec

我们实际上在这里达到了一个有趣的状态：我们在 Vec 和 IntoIter 中重复了指定缓冲区和释放其内存的逻辑。现在我们已经实现了它，并且确定了实际的逻辑重复，这是一个进行一些逻辑压缩的好时机。

我们将抽象出(ptr, cap)对，并为它们编写分配、增长和释放的逻辑：

  1. struct RawVec<T> {
  2.     ptr: NonNull<T>,
  3.     cap: usize,
  4.     _marker: PhantomData<T>,
  5. }
  7. unsafe impl<T: Send> Send for RawVec<T> {}
  8. unsafe impl<T: Sync> Sync for RawVec<T> {}
  10. impl<T> RawVec<T> {
  11.     fn new() -> Self {
  12.         assert!(mem::size_of::<T>() != 0, "TODO: implement ZST support");
  13.         RawVec {
  14.             ptr: NonNull::dangling(),
  15.             cap: 0,
  16.             _marker: PhantomData,
  17.         }
  18.     }
  20.     fn grow(&mut self) {
  21.         let (new_cap, new_layout) = if self.cap == 0 {
  22.             (1, Layout::array::<T>(1).unwrap())
  23.         } else {
  24.             // 保证新申请的内存没有超出 `isize::MAX` 字节
  25.             let new_cap = 2 * self.cap;
  27.             // `Layout::array` 会检查申请的空间是否小于等于 usize::MAX，
  28.             // 但是因为 old_layout.size() <= isize::MAX，
  29.             // 所以这里的 unwrap 永远不可能失败
  30.             let new_layout = Layout::array::<T>(new_cap).unwrap();
  31.             (new_cap, new_layout)
  32.         };
  34.         // 保证新申请的内存没有超出 `isize::MAX` 字节
  35.         assert!(new_layout.size() <= isize::MAX as usize, "Allocation too large");
  37.         let new_ptr = if self.cap == 0 {
  38.             unsafe { alloc::alloc(new_layout) }
  39.         } else {
  40.             let old_layout = Layout::array::<T>(self.cap).unwrap();
  41.             let old_ptr = self.ptr.as_ptr() as *mut u8;
  42.             unsafe { alloc::realloc(old_ptr, old_layout, new_layout.size()) }
  43.         };
  45.         // 如果分配失败，`new_ptr` 就会成为空指针，我们需要对应 abort 的操作
  46.         self.ptr = match NonNull::new(new_ptr as *mut T) {
  47.             Some(p) => p,
  48.             None => alloc::handle_alloc_error(new_layout),
  49.         };
  50.         self.cap = new_cap;
  51.     }
  52. }
  54. impl<T> Drop for RawVec<T> {
  55.     fn drop(&mut self) {
  56.         if self.cap != 0 {
  57.             let layout = Layout::array::<T>(self.cap).unwrap();
  58.             unsafe {
  59.                 alloc::dealloc(self.ptr.as_ptr() as *mut u8, layout);
  60.             }
  61.         }
  62.     }
  63. }

随后，把 Vec 改成这样：

  1. pub struct Vec<T> {
  2.     buf: RawVec<T>,
  3.     len: usize,
  4. }
  6. impl<T> Vec<T> {
  7.     fn ptr(&self) -> *mut T {
  8.         self.buf.ptr.as_ptr()
  9.     }
  11.     fn cap(&self) -> usize {
  12.         self.buf.cap
  13.     }
  15.     pub fn new() -> Self {
  16.         Vec {
  17.             buf: RawVec::new(),
  18.             len: 0,
  19.         }
  20.     }
  22.     // push/pop/insert/remove 这些操作做了小小的改变，如下所示:
  23.     // * `self.ptr.as_ptr() -> self.ptr()`
  24.     // * `self.cap -> self.cap()`
  25.     // * `self.grow() -> self.buf.grow()`
  26. }
  28. impl<T> Drop for Vec<T> {
  29.     fn drop(&mut self) {
  30.         while let Some(_) = self.pop() {}
  31.         // RawVec 来负责释放内存
  32.     }
  33. }

最后，我们可以把 IntoIter 改得相当简单：

  1. pub struct IntoIter<T> {
  2.     _buf: RawVec<T>, // 我们实际上并不关心这个，只需要他们保证分配的空间不被释放
  3.     start: *const T,
  4.     end: *const T,
  5. }
  7. // next 和 next_back 保持不变，因为它们并没有用到 buf
  9. impl<T> Drop for IntoIter<T> {
  10.     fn drop(&mut self) {
  11.         // 我们只需要确保 Vec 中所有元素都被读取了，
  12.         // 在这之后这些元素会被自动清理
  13.         for _ in &mut *self {}
  14.     }
  15. }
  17. impl<T> IntoIterator for Vec<T> {
  18.     type Item = T;
  19.     type IntoIter = IntoIter<T>;
  20.     fn into_iter(self) -> IntoIter<T> {
  21.         unsafe {
  22.             // 需要使用 ptr::read 非安全地把 buf 移出，因为它没有实现 Copy，
  23.             // 而且 Vec 实现了 Drop Trait (因此我们不能销毁它)
  24.             let buf = ptr::read(&self.buf);
  25.             let len = self.len;
  26.             mem::forget(self);
  28.             IntoIter {
  29.                 start: buf.ptr.as_ptr(),
  30.                 end: if buf.cap == 0 {
  31.                     // 不能通过这个指针获取偏移，除非已经分配了内存
  32.                     buf.ptr.as_ptr()
  33.                 } else {
  34.                     buf.ptr.as_ptr().add(len)
  35.                 },
  36.                 _buf: buf,
  37.             }
  38.         }
  39.     }
  40. }

是不是好多了！