SHARE
TWEET

Untitled

a guest Sep 20th, 2019 89 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. use core::marker::{PhantomPinned, Unpin};
  2. use core::mem::MaybeUninit;
  3. use core::pin::Pin;
  4. extern crate alloc;
  5. use alloc::alloc as ma;
  6.  
  7. #[repr(C)]
  8. pub struct PinSmallVecData<A>
  9. where
  10.     A: Array + Unpin,
  11.     A::Item: Unpin,
  12. {
  13.     ptr: *mut A::Item,
  14.     len: usize,
  15.     cap: usize,
  16.     inline_items: MaybeUninit<A>,
  17.     _boo: PhantomPinned,
  18. }
  19.  
  20. impl<A> PinSmallVecData<A>
  21. where
  22.     A: Array + Unpin,
  23.     A::Item: Unpin,
  24. {
  25.     pub fn new() -> Self {
  26.         Self {
  27.             ptr: core::ptr::null_mut(),
  28.             len: 0,
  29.             cap: 0,
  30.             inline_items: MaybeUninit::uninit(),
  31.             _boo: PhantomPinned,
  32.         }
  33.     }
  34. }
  35.  
  36. pub struct PinSmallVec<'a, A>
  37. where
  38.     A: Array + Unpin,
  39.     A::Item: Unpin,
  40. {
  41.     data: Pin<&'a mut PinSmallVecData<A>>,
  42. }
  43.  
  44. impl<'a, A> PinSmallVec<'a, A>
  45. where
  46.     A: Array + Unpin,
  47.     A::Item: Unpin,
  48. {
  49.     #[doc(hidden)]
  50.     pub unsafe fn new(v: &'a mut PinSmallVecData<A>) -> Self {
  51.         v.ptr = v.inline_items.as_mut_ptr() as *mut _;
  52.         v.cap = A::size();
  53.         v.len = 0;
  54.         Self {
  55.             data: Pin::new_unchecked(v),
  56.         }
  57.     }
  58.  
  59.     #[inline]
  60.     pub fn len(&self) -> usize {
  61.         self.data.len
  62.     }
  63.  
  64.     #[inline]
  65.     pub fn capacity(&self) -> usize {
  66.         self.data.cap
  67.     }
  68.  
  69.     #[inline]
  70.     pub fn is_small(&self) -> bool {
  71.         (self.data.ptr as usize) == (self.data.inline_items.as_ptr() as usize)
  72.     }
  73.  
  74.     #[inline]
  75.     pub fn as_slice(&self) -> &[A::Item] {
  76.         unsafe { core::slice::from_raw_parts(self.data.ptr, self.data.len) }
  77.     }
  78.  
  79.     #[inline]
  80.     pub fn as_slice_mut(&mut self) -> &mut [A::Item] {
  81.         unsafe { core::slice::from_raw_parts_mut(self.data.ptr, self.data.len) }
  82.     }
  83.  
  84.     #[inline]
  85.     pub fn push(&mut self, v: A::Item) {
  86.         if self.capacity() == self.len() {
  87.             self.grow(self.capacity() + 1);
  88.         }
  89.         // debug_assert!(self.capacity() > self.len());
  90.         unsafe {
  91.             let mut data = self.data.as_mut().get_unchecked_mut();
  92.             data.ptr.add(data.len).write(v);
  93.             data.len += 1;
  94.         }
  95.     }
  96.  
  97.     #[inline]
  98.     pub fn pop(&mut self) -> Option<A::Item> {
  99.         if self.len() == 0 {
  100.             return None;
  101.         }
  102.         unsafe {
  103.             let mut data = self.data.as_mut().get_unchecked_mut();
  104.             data.len -= 1;
  105.             let len = data.len;
  106.             Some(data.ptr.add(len).read())
  107.         }
  108.     }
  109.  
  110.     #[cold]
  111.     fn grow(&mut self, min_c: usize) {
  112.         // note: this is not intended to be fully production ready.
  113.         unsafe {
  114.             let mut data = self.data.as_mut().get_unchecked_mut();
  115.             let new_c = min_c.max(2 * data.cap + 1);
  116.             let small = (data.ptr as usize) == (data.inline_items.as_ptr() as usize);
  117.  
  118.             if !small {
  119.                 let layout = ma::Layout::from_size_align(
  120.                     core::mem::size_of::<A::Item>() * data.cap,
  121.                     core::mem::align_of::<A::Item>(),
  122.                 ).unwrap();
  123.  
  124.                 let np = ma::realloc(data.ptr as *mut u8, layout, new_c * core::mem::size_of::<A::Item>());
  125.                 if np.is_null() {
  126.                     ma::handle_alloc_error(layout);
  127.                 }
  128.                 data.ptr = np as *mut A::Item;
  129.                 data.cap = new_c;
  130.             } else {
  131.                 let layout = ma::Layout::from_size_align(
  132.                     core::mem::size_of::<A::Item>() * new_c,
  133.                     core::mem::align_of::<A::Item>(),
  134.                 )
  135.                 .unwrap();
  136.  
  137.                 let np = ma::alloc(layout);
  138.                 if np.is_null() {
  139.                     ma::handle_alloc_error(layout);
  140.                 }
  141.  
  142.                 let np = np as *mut A::Item;
  143.                 core::ptr::copy_nonoverlapping(data.ptr, np, data.len);
  144.                 data.ptr = np;
  145.                 data.cap = new_c;
  146.             }
  147.         }
  148.     }
  149. }
  150.  
  151. impl<'a, A> Drop for PinSmallVec<'a, A>
  152. where
  153.     A: Array + Unpin,
  154.     A::Item: Unpin,
  155. {
  156.     fn drop(&mut self) {
  157.         while let Some(_) = self.pop() {};
  158.         if !self.is_small() {
  159.             let c = self.capacity();
  160.             let layout = ma::Layout::from_size_align(
  161.                 core::mem::size_of::<A::Item>() * c,
  162.                 core::mem::align_of::<A::Item>(),
  163.             ).unwrap();
  164.             unsafe {
  165.                 ma::dealloc(self.data.ptr as *mut _, layout);
  166.             }
  167.         }
  168.     }
  169. }
  170.  
  171.  
  172. macro_rules! local_smallvec {
  173.     ($name:ident : [$T:ty; $N:expr]) => {
  174.         let mut $name: PinSmallVecData<[$T; $N]> = PinSmallVecData::new();
  175.         #[allow(unused_mut)]
  176.         let mut $name: PinSmallVec<'_, [$T; $N]> = unsafe { PinSmallVec::new(&mut $name) };
  177.     };
  178. }
  179.  
  180.  
  181. fn main() {
  182.     local_smallvec!(foobars: [i32; 0x40]);
  183.     for x in 0..100i32 {
  184.         foobars.push(x);
  185.     }
  186.     for x in (0..100i32).rev() {
  187.         assert_eq!(foobars.pop(), Some(x));
  188.     }
  189.     assert_eq!(foobars.pop(), None);
  190. }
  191.  
  192.  
  193. /// rest off this is from smallvec crate once the maybeuninit patch lands
  194.  
  195. pub unsafe trait Array {
  196.     type Item;
  197.     fn size() -> usize;
  198. }
  199.  
  200. macro_rules! impl_array {
  201.     ($($N:expr),*) => {$(
  202.         unsafe impl<T> Array for [T; $N] {
  203.             #[inline] fn size() -> usize { $N }
  204.             type Item = T;
  205.         }
  206.     )*}
  207. }
  208.  
  209. impl_array!(
  210.     0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 20, 24, 32, 36, 0x40, 0x80, 0x100, 0x200, 0x400, 0x800,
  211.     0x1000, 0x2000, 0x4000, 0x8000, 0x10000, 0x20000, 0x40000, 0x80000, 0x100000
  212. );
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!
 
Top