Better way to modify a slice in place

#1

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?

1 Like
#2

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() })
        }
    }
}
2 Likes
#3

That makes sense :slight_smile: Somehow I didn't think of using mem::take in this case.
For completeness, my final, untested solution:

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> {
        let small_vec = core::mem::take(&mut self.small_vec);
        small_vec.split_last_mut().and_then(|(last, front)| {
            let ret = last.as_mut_ptr();
            let ret = unsafe { ret.read() };

            self.small_vec = front;
            Some(ret)
        })
    }
}
#4 
fn next(&mut self) -> Option<Self::Item> {
    mem::take(&mut self.small_vec)
        .split_last_mut()
        .map(|(last, front)| {
            self.small_vec = front;
            unsafe { last.as_ptr().read() }
        })
}
2 Likes
#5

And to get rid of that last bit of unsafe:

- unsafe { last.as_ptr().read() }
+ mem::replace(last, MaybeUninit::uninit())
  • EDIT: this yields a MaybeUninit<T> which isn't what the method had to yield
#6

This won’t work.

  • last[0] is nonsensical, since last already is a reference to a single element
  • you cannot get rid of unsafe if you want to retrieve the value out of a MaybeUninit<T>
#7

Im too used to using .split_at_mut() that I didn't notice the last version already yields a single ref.

And yes, the assume init will need unsafe anyways, I guess I should go to sleep :sweat_smile:

#8

@steffahn was a bit faster than me :slight_smile:

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 :stuck_out_tongue:

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 :smiley: