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

code review
1

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

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

1 Like
2

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

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
4

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

5

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
6

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
7

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 t = pt.read();
            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!

8

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.

9

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

1 Like
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
11

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
12

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

I'll update my code accordingly.

13

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
14

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

15

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

16

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
17

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(
                  start.add(self.u_len).cast(),
                  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 :smile:

18

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?

19

Yes.

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

1 Like
20

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