1 use std::ptr;
2 
3 /// A typesafe helper that separates new value construction from
4 /// vector growing, allowing LLVM to ideally construct the element in place.
5 pub struct VecAllocation<'a, T: 'a> {
6     vec: &'a mut Vec<T>,
7     index: usize,
8 }
9 
10 impl<'a, T> VecAllocation<'a, T> {
11     /// Consumes self and writes the given value into the allocation.
12     // writing is safe because alloc() ensured enough capacity
13     // and `Allocation` holds a mutable borrow to prevent anyone else
14     // from breaking this invariant.
15     #[inline(always)]
init(self, value: T) -> usize16     pub fn init(self, value: T) -> usize {
17         unsafe {
18             ptr::write(self.vec.as_mut_ptr().add(self.index), value);
19             self.vec.set_len(self.index + 1);
20         }
21         self.index
22     }
23 }
24 
25 /// An entry into a vector, similar to `std::collections::hash_map::Entry`.
26 pub enum VecEntry<'a, T: 'a> {
27     /// Entry has just been freshly allocated.
28     Vacant(VecAllocation<'a, T>),
29     /// Existing entry.
30     Occupied(&'a mut T),
31 }
32 
33 impl<'a, T> VecEntry<'a, T> {
34     /// Sets the value for this entry.
35     #[inline(always)]
set(self, value: T)36     pub fn set(self, value: T) {
37         match self {
38             VecEntry::Vacant(alloc) => { alloc.init(value); }
39             VecEntry::Occupied(slot) => { *slot = value; }
40         }
41     }
42 }
43 
44 /// Helper trait for a `Vec` type that allocates up-front.
45 pub trait VecHelper<T> {
46     /// Grows the vector by a single entry, returning the allocation.
alloc(&mut self) -> VecAllocation<T>47     fn alloc(&mut self) -> VecAllocation<T>;
48     /// Either returns an existing element, or grows the vector by one.
49     /// Doesn't expect indices to be higher than the current length.
entry(&mut self, index: usize) -> VecEntry<T>50     fn entry(&mut self, index: usize) -> VecEntry<T>;
51 }
52 
53 impl<T> VecHelper<T> for Vec<T> {
alloc(&mut self) -> VecAllocation<T>54     fn alloc(&mut self) -> VecAllocation<T> {
55         let index = self.len();
56         if self.capacity() == index {
57             self.reserve(1);
58         }
59         VecAllocation {
60             vec: self,
61             index,
62         }
63     }
64 
entry(&mut self, index: usize) -> VecEntry<T>65     fn entry(&mut self, index: usize) -> VecEntry<T> {
66         if index < self.len() {
67             VecEntry::Occupied(unsafe {
68                 self.get_unchecked_mut(index)
69             })
70         } else {
71             assert_eq!(index, self.len());
72             VecEntry::Vacant(self.alloc())
73         }
74     }
75 }
76