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

Hey all, I believe I recently had a breakthrough concerning how to think about unsafe code and I tried to prove it to myself by implementing a function that performs a mapping of `Vec<T>` to `Vec<U>` without allocating (given equal size and alignment of T and U). So here is my code

``````#![feature(vec_into_raw_parts)]

fn transform_in_place<T,U,F>(v: Vec<T>, mut f: F) -> Vec<U>
where F : FnMut(T) -> U {
assert!(std::mem::align_of::<T>() == std::mem::align_of::<U>());
assert!(std::mem::size_of::<T>() == std::mem::size_of::<U>());

unsafe {
let (pstart,len,cap) = v.into_raw_parts();
for pelem in (0..len).map(|j| pstart.offset(j.try_into().unwrap())) {
let mut t = MaybeUninit::uninit();
core::intrinsics::copy_nonoverlapping(pelem, t.as_mut_ptr(), 1);
let u = f(t.assume_init());
let pu :*mut U = pelem as _;
*pu = u;
}
Vec::from_raw_parts(pstart as _, len, cap)
}
}
``````

So what I do is

1. get the raw parts of a vector (pointer to mem, len, capacity)
2. iterate over the addresses of the elements in the vector
3. create a temporary element of type t into which I copy the vector element without actually dereferencing the element (so T does not have to implement Copy)
4. I perform the transformation and put the result into the memory previously occupied by the element
5. finally I create a new vector from the raw parts

I believe I understand what I am doing and I did create a test case and run it with miri without complaints. However, I really appreciate a code review that points out what I could do better and / or where I am wrong.

If there is something I can do better in terms of performance, I also appreciate hints

and yes, I know that the standard library is almost guarantee to do what I want if I take the idiomatic approach. But as I said, this is mostly about testing my understanding of unsafe.

1 Like

`*pu = u;` will drop the data in `*pelem` as if it were a `U`, but it is actually an already-moved `T`. You need to use `ptr::write()` to put the `U` in its place.

2 Likes

The `*tu = u;` line is likely wrong because it drops the old value (again).

The whole dance with `MaybeUninit` seems unnecessary too. You could just do `ptr::read()`, call the function, `ptr::write()`.

2 Likes

amazing! thank you both. This is somehting I completely missed, but I now looking at it I don't understand why my code is wrong (not saying it isn't!!). To my understanding I want to call the destructor of every element t in the Vec once. But after splitting the vector into raw parts the destructors won't be run. Then I copy the element into the temporary element of type t and transform it. This t is dropped as it should be. But then I can see that there is a problem with `*pu = u` because OH MY WORD now I see it...

so what you are saying is that the `*pu = u` line drops the thing at element t because it actually points there... wow

An ordinary safe assignment to `*some_mutable_reference` must drop the existing value or it would be a memory leak. So, pointers work the same way.

1 Like

Note that this behavior used to be in `collections`. It was dropped for being impractical and insufficiently general to be in the standard library, but not because the implementation was unsound.

3 Likes

Thanks, my test cases were too simple. I only tried it for Vec -> Vec which does not invoke the destructor. When I redid it with Vec<Box> -> Vec<Box>, miri did catch the double free.

I've boiled my code down to this now:

``````fn transform_in_place<T, U, F>(v: Vec<T>, mut f: F) -> Vec<U>
where
F: FnMut(T) -> U,
{
assert!(std::mem::align_of::<T>() == std::mem::align_of::<U>());
assert!(std::mem::size_of::<T>() == std::mem::size_of::<U>());

unsafe {
let (pstart, len, cap) = v.into_raw_parts();
for pt in (0..len).map(|j| pstart.offset(j.try_into().unwrap())) {
let u = f(t);
let pu: *mut U = pt as _;
pu.write(u);
}
Vec::from_raw_parts(pstart as _, len, cap)
}
}
``````

I know I could boil down the code inside the for loop to one line but I like the step by step more.

Did I miss something else? Thanks again to both of you!

oh cool, did not know that. Rust 1.3. was way before my time with Rust. From what I understand iterators and specializations inside the standard library can achieve the same effect so that the idiomatic `v.iter().map(...).collect()` pattern can be used instead but I am not sure if that is guaranteed.

It's better practice to call `.cast()` on pointers instead of `as _` because that catches changes in mutability.

1 Like

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));
}
}
}

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

Itâ€™s only for `.into_iter()`, but otherwise, yes, as far as Iâ€™m aware that will predictably always happen with the current version of the compiler. I donâ€™t believe this is documented, so thereâ€™s no 100% guarantee that this optimization will never be removed. On the other hand, that would be somewhat of a performance regression, so Iâ€™d think itâ€™s rather unlikely that will ever happen without good reason.

4 Likes

oh wow, great point. I always forget panic safety... and I had to think about what you meant by "leaking memory from the vec itself"... wow this is all really complicated

I'll update my code accordingly.

If you haven't been through it yet, be sure to read through the nomicon chapter about implementing `Vec`. It has lots of techniques in it for using type to get the compiler to help you remember to do things like this, such as the one in RawVec - The Rustonomicon

1 Like

thank you kindly, I haven't had time to tackle it yet. Will give it a read

Okay so I've taken another stab at this... I did take the idea of the Vec<Union<T,U>> type like so:

``````  union Union<T, U> {
pub first: ManuallyDrop<T>,
pub second: ManuallyDrop<U>,
}

struct TransitioningVec<T, U> {
vector: Vec<Union<T, U>>,
// The number of elements that have transitioned
// to type U.
u_len: usize,
}

impl<T, U> TransitioningVec<T, U> {
pub fn new(v: Vec<T>) -> Self {
assert_eq!(std::mem::size_of::<T>(), std::mem::size_of::<U>());
assert_eq!(std::mem::align_of::<T>(), std::mem::align_of::<U>());

let mut v = ManuallyDrop::new(v); => ManuallyDrop<Vec<T>>
let data = unsafe { Vec::from_raw_parts(v.as_mut_ptr() as _, v.len(), v.capacity()) }; <- (length) => Vec<Union<T, U>>
Self {
vector: data,
u_len: 0,
}
}
}

impl<T, U> TransitioningVec<T, U> {
pub fn new(v: Vec<T>) -> Self {
assert_eq!(std::mem::size_of::<T>(), std::mem::size_of::<U>());
assert_eq!(std::mem::align_of::<T>(), std::mem::align_of::<U>());

let mut v = ManuallyDrop::new(v); => ManuallyDrop<Vec<T>>
let data = unsafe { Vec::from_raw_parts(v.as_mut_ptr() as _, v.len(), v.capacity()) }; <- (length) => Vec<Union<T, U>>
Self {
vector: data,
u_len: 0,
}
}
}

impl<T, U> TransitioningVec<T, U> {
#[inline]
pub fn map_in_place<F: FnMut(T) -> U>(mut self, mut f: F) -> Vec<U> {
let start_ptr: *mut T = self.vector.as_mut_ptr() as _;
while self.u_len < self.vector.len() {
unsafe {
let t_ptr = start_ptr.offset(self.u_len as isize); <- (count) => *mut T
let u_ptr : *mut U = t_ptr as _;
let t = t_ptr.read(); => T
u_ptr.write(f(t)); <- (val)
}
self.u_len += 1;
}
debug_assert_eq!(self.u_len, self.vector.len());
let mut me = ManuallyDrop::new(self); => ManuallyDrop<TransitioningVec<â€¦, â€¦>>
unsafe {
Vec::from_raw_parts(
me.vector.as_mut_ptr() as _,
me.vector.len(), <- (length)
me.vector.capacity(),
)
}
}
}

impl<T, U> Drop for TransitioningVec<T, U> {
fn drop(&mut self) {
unsafe {
let t_ptr: *mut T = self.vector.as_mut_ptr() as _;
let u_ptr: *mut U = self.vector.as_mut_ptr() as _;
for offset in 0..self.u_len { => usize
u_ptr.offset(offset as isize).drop_in_place(); <- (count)
}
for offset in self.u_len..self.vector.len() { => usize
t_ptr.offset(offset as isize).drop_in_place(); <- (count)
}
}
}
}

fn map_in_place<T, U, F>(v: Vec<T>, f: F) -> Vec<U>
where
F: FnMut(T) -> U,
{
TransitioningVec::new(v).map_in_place(f)
}
``````

I've created a helper structure that keeps track of the vector when it transitions from element by element from type T to U. When the function `f` panics the structure is droppend and (so I think) will correctly drop the elements as U or T depending on how many elements have made the transform T->U. It should also deallocate the associated heap memory.

I'll gladly take any additional input and corrections because so far I thought many times I was done when I wasn't

Overall this looks pretty good, though something weird happened when you copy-pasted it, since the formatting is weird and it has a bunch of extra stuff after some of the semicolons.

Two things here:

1. Never use `.offset(x as isize)`. Use `.add(x)` instead. (In general, you typically don't want to use `offset` ever. Use `add` and `sub` instead, and your life will be much simpler.)

2. If you want to drop a bunch of things that are contiguous in memory, don't write a loop, `drop_in_place` a slice instead. (See my `Dropper` above.)

Up to you, but this could just be

``````assert_eq!(std::alloc::Layout::<T>::new(), std::alloc::Layout::<U>::new());
``````

rather than checking them the parts separately.

3 Likes

argh sorry about the weird formatting. I haven't gotten my android terminal emulator to play nice with the clipboard and the code above still contains inlay hints from my IDE. I took all your suggestions (replacing offset with add, also the more elegant assertion). I also changed the dropping logic to this:

``````impl<T, U> Drop for TransitioningVec<T, U> {
fn drop(&mut self) {
let start = self.vector.as_mut_ptr();
unsafe {
let u_slice :&mut [U] = std::slice::from_raw_parts_mut(start.cast(), self.u_len);
let t_slice :&mut [T] = std::slice::from_raw_parts_mut(
self.vector.len() - self.u_len,
);
std::ptr::drop_in_place(u_slice);
std::ptr::drop_in_place(t_slice);
}
}
}
``````

I really hope this is sound now and thank you and everyone for their suggestions. If something is still wrong please let me know

hey, sorry for the necromancy, but I have revisited the code again and I think I do have to use a loop to correctly drop the slices. What you are doing is creating slices and passing them to drop_in_place. But ptr::drop_in_place takes a pointer argument. To my mind this will coerce the slice into a pointer to its first element, effectively only dropping the first element. Am I wrong here?

Yes.

No. `&mut [T]` coerces to `*mut [T]`. Rust is not C. Types matter.

1 Like

In Rust, if weâ€™re technically correct (and pedantic), what youâ€™d call a â€śsliceâ€ť is actually just the type `[T]`. Types like `Box<[T]>` or `&mut [T]` or `&[T]` then are slice pointers (or references, for the latter two). But slices support raw pointers just as well, like `*const [T]` or `*mut [T]`. All of these will be fat pointers, keeping the length information intact.

`ptr::drop_in_place` does support unsized types, as you can see in its signature

``````pub unsafe fn drop_in_place<T>(to_drop: *mut T)
where
T: ?Sized,
``````

having a `?Sized` bound. This means that it does support dropping trait objects and slices, too.

In fact, itâ€™s the primary way of dropping slices (or trait objects) as you cannot drop them by value directly (since theyâ€™re unsized). E.g. if you look at [the code for `impl<T> Drop for Vec<T>` it uses the same pattern. Well almost. It uses `ptr::slice_from_raw_parts_mut` instead of `slice::from_raw_parts_mut`, and the former should be preferred as it doesnâ€™t involve having a `&mut` reference to a slice of dropped values (even if itâ€™s only for a very short time, after the call to `drop_in_place`), and it avoids the need for any coercions, too.

Other examples in `std` include e.g. the code thatâ€™s called in the drop implementaiton of `Arc` (if itâ€™s the last copy of the `Arc`), which uses `ptr::drop_in_place` to power dropping anything from simple sized types like `Arc<u8>` (well, in that case itâ€™s a no-op) to unsized targets like `Arc<[String]>` or `Arc<dyn Trait>`.

3 Likes