We've actually reached an interesting situation here: we've duplicated the logic for specifying a buffer and freeing its memory in Vec and IntoIter. Now that we've implemented it and identified actual logic duplication, this is a good time to perform some logic compression.
We're going to abstract out the (ptr, cap)
pair and give them the logic for
allocating, growing, and freeing:
struct RawVec<T> { ptr: Unique<T>, cap: usize, } impl<T> RawVec<T> { fn new() -> Self { assert!(mem::size_of::<T>() != 0, "TODO: implement ZST support"); unsafe { RawVec { ptr: Unique::new(heap::EMPTY as *mut T), cap: 0 } } } // unchanged from Vec fn grow(&mut self) { unsafe { let align = mem::align_of::<T>(); let elem_size = mem::size_of::<T>(); let (new_cap, ptr) = if self.cap == 0 { let ptr = heap::allocate(elem_size, align); (1, ptr) } else { let new_cap = 2 * self.cap; let ptr = heap::reallocate(*self.ptr as *mut _, self.cap * elem_size, new_cap * elem_size, align); (new_cap, ptr) }; // If allocate or reallocate fail, we'll get `null` back if ptr.is_null() { oom() } self.ptr = Unique::new(ptr as *mut _); self.cap = new_cap; } } } impl<T> Drop for RawVec<T> { fn drop(&mut self) { if self.cap != 0 { let align = mem::align_of::<T>(); let elem_size = mem::size_of::<T>(); let num_bytes = elem_size * self.cap; unsafe { heap::deallocate(*self.ptr as *mut _, num_bytes, align); } } } }
And change Vec as follows:
fn main() { pub struct Vec<T> { buf: RawVec<T>, len: usize, } impl<T> Vec<T> { fn ptr(&self) -> *mut T { *self.buf.ptr } fn cap(&self) -> usize { self.buf.cap } pub fn new() -> Self { Vec { buf: RawVec::new(), len: 0 } } // push/pop/insert/remove largely unchanged: // * `self.ptr -> self.ptr()` // * `self.cap -> self.cap()` // * `self.grow -> self.buf.grow()` } impl<T> Drop for Vec<T> { fn drop(&mut self) { while let Some(_) = self.pop() {} // deallocation is handled by RawVec } } }pub struct Vec<T> { buf: RawVec<T>, len: usize, } impl<T> Vec<T> { fn ptr(&self) -> *mut T { *self.buf.ptr } fn cap(&self) -> usize { self.buf.cap } pub fn new() -> Self { Vec { buf: RawVec::new(), len: 0 } } // push/pop/insert/remove largely unchanged: // * `self.ptr -> self.ptr()` // * `self.cap -> self.cap()` // * `self.grow -> self.buf.grow()` } impl<T> Drop for Vec<T> { fn drop(&mut self) { while let Some(_) = self.pop() {} // deallocation is handled by RawVec } }
And finally we can really simplify IntoIter:
fn main() { struct IntoIter<T> { _buf: RawVec<T>, // we don't actually care about this. Just need it to live. start: *const T, end: *const T, } // next and next_back literally unchanged since they never referred to the buf impl<T> Drop for IntoIter<T> { fn drop(&mut self) { // only need to ensure all our elements are read; // buffer will clean itself up afterwards. for _ in &mut *self {} } } impl<T> Vec<T> { pub fn into_iter(self) -> IntoIter<T> { unsafe { // need to use ptr::read to unsafely move the buf out since it's // not Copy, and Vec implements Drop (so we can't destructure it). let buf = ptr::read(&self.buf); let len = self.len; mem::forget(self); IntoIter { start: *buf.ptr, end: buf.ptr.offset(len as isize), _buf: buf, } } } } }struct IntoIter<T> { _buf: RawVec<T>, // we don't actually care about this. Just need it to live. start: *const T, end: *const T, } // next and next_back literally unchanged since they never referred to the buf impl<T> Drop for IntoIter<T> { fn drop(&mut self) { // only need to ensure all our elements are read; // buffer will clean itself up afterwards. for _ in &mut *self {} } } impl<T> Vec<T> { pub fn into_iter(self) -> IntoIter<T> { unsafe { // need to use ptr::read to unsafely move the buf out since it's // not Copy, and Vec implements Drop (so we can't destructure it). let buf = ptr::read(&self.buf); let len = self.len; mem::forget(self); IntoIter { start: *buf.ptr, end: buf.ptr.offset(len as isize), _buf: buf, } } } }
Much better.