I'm currently playing around with implementing a stack based vec with limited (const) size and wanted to implement something like Drain for it.
So currently my problem looks like this (simplified):
pub struct Drain<'a, T> {
small_vec: &'a mut [MaybeUninit<T>],
}
impl<'a, T> Iterator for Drain<'a, T> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
if self.small_vec.len() == 0 {
None
} else {
let ret = self.small_vec.last_mut().unwrap().as_mut_ptr();
let ret = unsafe { ret.read() };
let new_len = self.small_vec.len() - 1;
let new_vec: &mut [MaybeUninit<T>] = &mut self.small_vec[..new_len];
//transmute is used to extend the lifetime. Is this safe??
self.small_vec = unsafe { core::mem::transmute(new_vec) };
Some(ret)
}
}
}
What I'm wondering is: Is there a better way to shorten the slice and work on the last element?
I found methods like split_last_mut, but wasn't able to convince rustc of what I was doing. The problem always was when trying to assign self.small_vec, a short example:
pub struct Drain<'a, T> {
small_vec: &'a mut [MaybeUninit<T>],
}
impl<'a, T> Iterator for Drain<'a, T> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
if self.small_vec.len() == 0 {
None
} else {
let (ret, front) = self.small_vec.split_last_mut().unwrap();
self.small_vec = front;
let ret = ret.as_mut_ptr();
Some(unsafe { ret.read() })
}
}
}
Maybe it would be useful to have a function similar to split_last_mut but that modified self in place and additionally returned the last element?
You can take the mutable reference to the slice with mem::take. The type &'a mut [T]has a Default implementation, yielding a reference to the empty slice.
use std::mem::{self, MaybeUninit};
pub struct Drain<'a, T> {
small_vec: &'a mut [MaybeUninit<T>],
}
impl<'a, T> Iterator for Drain<'a, T> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
if self.small_vec.len() == 0 {
None
} else {
let (ret, front) = mem::take(&mut self.small_vec).split_last_mut().unwrap();
self.small_vec = front;
let ret = ret.as_mut_ptr();
Some(unsafe { ret.read() })
}
}
}
It is always good to see that there are safe ways to do stuff, but in this case I think I will keep the unsafe one. My main reasons for this are:
SmallVec does a bit of work to track initialized and uninitialized values which will always require some unsafe
The surface area of the exposed api is still quite manageable and as long as that's safe I don't really bother
Most use of unsafe in my implementation was more or less copied from std::Vec (the only exception being the drain function)
I want to avoid unnecessary writes (or copies), since this will be a rather hot part of the code (I now I should benchmark stuff before optimizing it especially before using unsafe, but this is just a small side project of mine and benchmarking everything just takes too much time for a small side project)
Also, I too should go to sleep
For the interested, my drain function looks as follows:
pub fn drain<R>(&mut self, range: R) -> Drain<'_, T>
where
R: RangeBounds<usize>,
{
let Range { start, end } = check_range(self.len, range);
let len = self.len;
let slice = self.as_mut_slice();
slice[start..].reverse();
let reverse_elements = len - end;
slice[start..start + reverse_elements].reverse();
//the elements to drain are now in reverse order at the end of the (initialized) array
//basically mem::forget the elements that get drained
self.len -= end - start;
Drain {
small_vec: &mut self.data[self.len..len],
}
}
The drain function on Vec is certainly more efficient, but I think my function allows much simpler behavior for stuff like mem::forget. Basically I use the method that appeared on Quote of the Week some days ago: pre poop your pants