Is My Highly Unsafe Code Correct? In Place Mapping a Vector

code review
10

You now have a memory leak if f panics. That's better than being UB, of course, but probably still not what you want.

The usual way to fix this is a guard that cleans up on drop which you forget once you've done all the work you need to do. Something like this:

#![feature(vec_into_raw_parts)]
pub fn transform_in_place<T, U, F>(v: Vec<T>, mut f: F) -> Vec<U>
where
    F: FnMut(T) -> U,
{
    use std::marker::PhantomData;
    use std::ptr;

    assert!(std::mem::align_of::<T>() == std::mem::align_of::<U>());
    assert!(std::mem::size_of::<T>() == std::mem::size_of::<U>());

    struct Dropper<T, U> {
        pstart: *mut T,
        j: usize,
        len: usize,
        phantom_t: PhantomData<[T]>,
        phantom_u: PhantomData<[U]>,
    }
    impl<T, U> Drop for Dropper<T, U> {
        fn drop(&mut self) {
            unsafe {
                ptr::drop_in_place(ptr::slice_from_raw_parts_mut::<T>(self.pstart, self.j));
                ptr::drop_in_place(ptr::slice_from_raw_parts_mut::<U>(self.pstart.add(self.j).cast(), self.len - self.j));
            }   
        }
    }

    unsafe {
        let (pstart, len, cap) = v.into_raw_parts();
        let mut dropper = Dropper::<T, U> { pstart, j: 0, len, phantom_t: PhantomData, phantom_u: PhantomData };
        for j in 0..len {
            let pt = pstart.add(j);
            let t = pt.read();
            let u = f(t);
            let pu: *mut U = pt as _;
            pu.write(u);

            dropper.j = j;
        }
        std::mem::forget(dropper);
        Vec::from_raw_parts(pstart as _, len, cap)
    }
}

https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=b91913def5a8c1fd0182643703d19af1

Completely untested; be sure to fix the mistakes I probably made before using it

EDIT: Wait, that's not enough either, since it's still leaking the memory from the Vec itself.

Maybe you want to turn your Vec<T> into a Vec<Union<T, U>> first, and work on that, so that Vec's destructor will still clean up its memory. It'll still need the guard, though, to drop the elements, as a union can't.

5 Likes