The Final Code

  1. #![feature(ptr_internals)]
  2. #![feature(allocator_api)]
  3. #![feature(alloc_layout_extra)]
  5. use std::ptr::{Unique, NonNull, self};
  6. use std::mem;
  7. use std::ops::{Deref, DerefMut};
  8. use std::marker::PhantomData;
  9. use std::alloc::{Alloc, GlobalAlloc, Layout, Global, handle_alloc_error};
  11. struct RawVec<T> {
  12.     ptr: Unique<T>,
  13.     cap: usize,
  14. }
  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.         // Unique::empty() doubles as "unallocated" and "zero-sized allocation"
  22.         RawVec { ptr: Unique::empty(), cap: cap }
  23.     }
  25.     fn grow(&mut self) {
  26.         unsafe {
  27.             let elem_size = mem::size_of::<T>();
  29.             // since we set the capacity to usize::MAX when elem_size is
  30.             // 0, getting to here necessarily means the Vec is overfull.
  31.             assert!(elem_size != 0, "capacity overflow");
  33.             let (new_cap, ptr) = if self.cap == 0 {
  34.                 let ptr = Global.alloc(Layout::array::<T>(1).unwrap());
  35.                 (1, ptr)
  36.             } else {
  37.                 let new_cap = 2 * self.cap;
  38.                 let c: NonNull<T> = self.ptr.into();
  39.                 let ptr = Global.realloc(c.cast(),
  40.                                          Layout::array::<T>(self.cap).unwrap(),
  41.                                          Layout::array::<T>(new_cap).unwrap().size());
  42.                 (new_cap, ptr)
  43.             };
  45.             // If allocate or reallocate fail, oom
  46.             if ptr.is_err() {
  47.                 handle_alloc_error(Layout::from_size_align_unchecked(
  48.                     new_cap * elem_size,
  49.                     mem::align_of::<T>(),
  50.                 ))
  51.             }
  52.             let ptr = ptr.unwrap();
  54.             self.ptr = Unique::new_unchecked(ptr.as_ptr() as *mut _);
  55.             self.cap = new_cap;
  56.         }
  57.     }
  58. }
  60. impl<T> Drop for RawVec<T> {
  61.     fn drop(&mut self) {
  62.         let elem_size = mem::size_of::<T>();
  63.         if self.cap != 0 && elem_size != 0 {
  64.             unsafe {
  65.                 let c: NonNull<T> = self.ptr.into();
  66.                 Global.dealloc(c.cast(),
  67.                                Layout::array::<T>(self.cap).unwrap());
  68.             }
  69.         }
  70.     }
  71. }
  73. pub struct Vec<T> {
  74.     buf: RawVec<T>,
  75.     len: usize,
  76. }
  78. impl<T> Vec<T> {
  79.     fn ptr(&self) -> *mut T { self.buf.ptr.as_ptr() }
  81.     fn cap(&self) -> usize { self.buf.cap }
  83.     pub fn new() -> Self {
  84.         Vec { buf: RawVec::new(), len: 0 }
  85.     }
  86.     pub fn push(&mut self, elem: T) {
  87.         if self.len == self.cap() { self.buf.grow(); }
  89.         unsafe {
  90.             ptr::write(self.ptr().offset(self.len as isize), elem);
  91.         }
  93.         // Can't fail, we'll OOM first.
  94.         self.len += 1;
  95.     }
  97.     pub fn pop(&mut self) -> Option<T> {
  98.         if self.len == 0 {
  99.             None
  100.         } else {
  101.             self.len -= 1;
  102.             unsafe {
  103.                 Some(ptr::read(self.ptr().offset(self.len as isize)))
  104.             }
  105.         }
  106.     }
  108.     pub fn insert(&mut self, index: usize, elem: T) {
  109.         assert!(index <= self.len, "index out of bounds");
  110.         if self.cap() == self.len { self.buf.grow(); }
  112.         unsafe {
  113.             if index < self.len {
  114.                 ptr::copy(self.ptr().offset(index as isize),
  115.                           self.ptr().offset(index as isize + 1),
  116.                           self.len - index);
  117.             }
  118.             ptr::write(self.ptr().offset(index as isize), elem);
  119.             self.len += 1;
  120.         }
  121.     }
  123.     pub fn remove(&mut self, index: usize) -> T {
  124.         assert!(index < self.len, "index out of bounds");
  125.         unsafe {
  126.             self.len -= 1;
  127.             let result = ptr::read(self.ptr().offset(index as isize));
  128.             ptr::copy(self.ptr().offset(index as isize + 1),
  129.                       self.ptr().offset(index as isize),
  130.                       self.len - index);
  131.             result
  132.         }
  133.     }
  135.     pub fn into_iter(self) -> IntoIter<T> {
  136.         unsafe {
  137.             let iter = RawValIter::new(&self);
  138.             let buf = ptr::read(&self.buf);
  139.             mem::forget(self);
  141.             IntoIter {
  142.                 iter: iter,
  143.                 _buf: buf,
  144.             }
  145.         }
  146.     }
  148.     pub fn drain(&mut self) -> Drain<T> {
  149.         unsafe {
  150.             let iter = RawValIter::new(&self);
  152.             // this is a mem::forget safety thing. If Drain is forgotten, we just
  153.             // leak the whole Vec's contents. Also we need to do this *eventually*
  154.             // anyway, so why not do it now?
  155.             self.len = 0;
  157.             Drain {
  158.                 iter: iter,
  159.                 vec: PhantomData,
  160.             }
  161.         }
  162.     }
  163. }
  165. impl<T> Drop for Vec<T> {
  166.     fn drop(&mut self) {
  167.         while let Some(_) = self.pop() {}
  168.         // allocation is handled by RawVec
  169.     }
  170. }
  172. impl<T> Deref for Vec<T> {
  173.     type Target = [T];
  174.     fn deref(&self) -> &[T] {
  175.         unsafe {
  176.             ::std::slice::from_raw_parts(self.ptr(), self.len)
  177.         }
  178.     }
  179. }
  181. impl<T> DerefMut for Vec<T> {
  182.     fn deref_mut(&mut self) -> &mut [T] {
  183.         unsafe {
  184.             ::std::slice::from_raw_parts_mut(self.ptr(), self.len)
  185.         }
  186.     }
  187. }
  193. struct RawValIter<T> {
  194.     start: *const T,
  195.     end: *const T,
  196. }
  198. impl<T> RawValIter<T> {
  199.     unsafe fn new(slice: &[T]) -> Self {
  200.         RawValIter {
  201.             start: slice.as_ptr(),
  202.             end: if mem::size_of::<T>() == 0 {
  203.                 ((slice.as_ptr() as usize) + slice.len()) as *const _
  204.             } else if slice.len() == 0 {
  205.                 slice.as_ptr()
  206.             } else {
  207.                 slice.as_ptr().offset(slice.len() as isize)
  208.             }
  209.         }
  210.     }
  211. }
  213. impl<T> Iterator for RawValIter<T> {
  214.     type Item = T;
  215.     fn next(&mut self) -> Option<T> {
  216.         if self.start == self.end {
  217.             None
  218.         } else {
  219.             unsafe {
  220.                 let result = ptr::read(self.start);
  221.                 self.start = if mem::size_of::<T>() == 0 {
  222.                     (self.start as usize + 1) as *const _
  223.                 } else {
  224.                     self.start.offset(1)
  225.                 };
  226.                 Some(result)
  227.             }
  228.         }
  229.     }
  231.     fn size_hint(&self) -> (usize, Option<usize>) {
  232.         let elem_size = mem::size_of::<T>();
  233.         let len = (self.end as usize - self.start as usize)
  234.                   / if elem_size == 0 { 1 } else { elem_size };
  235.         (len, Some(len))
  236.     }
  237. }
  239. impl<T> DoubleEndedIterator for RawValIter<T> {
  240.     fn next_back(&mut self) -> Option<T> {
  241.         if self.start == self.end {
  242.             None
  243.         } else {
  244.             unsafe {
  245.                 self.end = if mem::size_of::<T>() == 0 {
  246.                     (self.end as usize - 1) as *const _
  247.                 } else {
  248.                     self.end.offset(-1)
  249.                 };
  250.                 Some(ptr::read(self.end))
  251.             }
  252.         }
  253.     }
  254. }
  259. pub struct IntoIter<T> {
  260.     _buf: RawVec<T>, // we don't actually care about this. Just need it to live.
  261.     iter: RawValIter<T>,
  262. }
  264. impl<T> Iterator for IntoIter<T> {
  265.     type Item = T;
  266.     fn next(&mut self) -> Option<T> { self.iter.next() }
  267.     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
  268. }
  270. impl<T> DoubleEndedIterator for IntoIter<T> {
  271.     fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
  272. }
  274. impl<T> Drop for IntoIter<T> {
  275.     fn drop(&mut self) {
  276.         for _ in &mut *self {}
  277.     }
  278. }
  283. pub struct Drain<'a, T: 'a> {
  284.     vec: PhantomData<&'a mut Vec<T>>,
  285.     iter: RawValIter<T>,
  286. }
  288. impl<'a, T> Iterator for Drain<'a, T> {
  289.     type Item = T;
  290.     fn next(&mut self) -> Option<T> { self.iter.next() }
  291.     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
  292. }
  294. impl<'a, T> DoubleEndedIterator for Drain<'a, T> {
  295.     fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
  296. }
  298. impl<'a, T> Drop for Drain<'a, T> {
  299.     fn drop(&mut self) {
  300.         // pre-drain the iter
  301.         for _ in &mut self.iter {}
  302.     }
  303. }
  305. # fn main() {
  306. #     tests::create_push_pop();
  307. #     tests::iter_test();
  308. #     tests::test_drain();
  309. #     tests::test_zst();
  310. #     println!("All tests finished OK");
  311. # }
  313. # mod tests {
  314. #     use super::*;
  315. #     pub fn create_push_pop() {
  316. #         let mut v = Vec::new();
  317. #         v.push(1);
  318. #         assert_eq!(1, v.len());
  319. #         assert_eq!(1, v[0]);
  320. #         for i in v.iter_mut() {
  321. #             *i += 1;
  322. #         }
  323. #         v.insert(0, 5);
  324. #         let x = v.pop();
  325. #         assert_eq!(Some(2), x);
  326. #         assert_eq!(1, v.len());
  327. #         v.push(10);
  328. #         let x = v.remove(0);
  329. #         assert_eq!(5, x);
  330. #         assert_eq!(1, v.len());
  331. #     }
  332. #
  333. #     pub fn iter_test() {
  334. #         let mut v = Vec::new();
  335. #         for i in 0..10 {
  336. #             v.push(Box::new(i))
  337. #         }
  338. #         let mut iter = v.into_iter();
  339. #         let first = iter.next().unwrap();
  340. #         let last = iter.next_back().unwrap();
  341. #         drop(iter);
  342. #         assert_eq!(0, *first);
  343. #         assert_eq!(9, *last);
  344. #     }
  345. #
  346. #     pub fn test_drain() {
  347. #         let mut v = Vec::new();
  348. #         for i in 0..10 {
  349. #             v.push(Box::new(i))
  350. #         }
  351. #         {
  352. #             let mut drain = v.drain();
  353. #             let first = drain.next().unwrap();
  354. #             let last = drain.next_back().unwrap();
  355. #             assert_eq!(0, *first);
  356. #             assert_eq!(9, *last);
  357. #         }
  358. #         assert_eq!(0, v.len());
  359. #         v.push(Box::new(1));
  360. #         assert_eq!(1, *v.pop().unwrap());
  361. #     }
  362. #
  363. #     pub fn test_zst() {
  364. #         let mut v = Vec::new();
  365. #         for _i in 0..10 {
  366. #             v.push(())
  367. #         }
  368. #
  369. #         let mut count = 0;
  370. #
  371. #         for _ in v.into_iter() {
  372. #             count += 1
  373. #         }
  374. #
  375. #         assert_eq!(10, count);
  376. #     }
  377. # }