Search code examples
rustundefined-behaviorunsafe

Polymorphic update on struct fields in Rust


Suppose I have a polymorphic type T<A>:

#[repr(C)]
pub struct T<A> {
    x: u32,
    y: Box<A>,
}

Below are my reasonings:

  • According to the Memory Layout Section of std::boxed:

    So long as T: Sized, a Box<T> is guaranteed to be represented as a single pointer and is also ABI-compatible with C pointers (i.e. the C type T*).

    So no matter what A is, y should have the same layout;

  • Considering additionally the #[repr(C)] attribute on T, I expect that for all As, T<A> would share the same layout;

  • Thus I should be able to modify y in-place, even to values of a different type Box<B>.

My question is whether the following code is well-formed or has undefined behaviours? (The code below has been edited.)

fn update<A, B>(t: Box<T<A>>, f: impl FnOnce(A) -> B) -> Box<T<B>> {
    unsafe {
        let p = Box::into_raw(t);
        let a = std::ptr::read(&(*p).y);
        let q = p as *mut T<B>;
        std::ptr::write(&mut (*q).y, Box::new(f(*a)));
        Box::from_raw(q)
    }
}

Note:

The above code aims to perform a polymorphic update in place, so that the x field is left as-is. Imagine x were not simply a u32, but some very large chunk of data. The whole idea is to change the type of y (along with its value) without touching field x.


As is pointed out by Frxstrem, the code below indeed causes undefined behaviour. I made a silly mistake: I forgot to re-allocate memory for the B produced by f. It seems the new code above passes Miri check.

fn update<A, B>(t: Box<T<A>>, f: impl FnOnce(A) -> B) -> Box<T<B>> {
    unsafe {
        let mut u: Box<T<std::mem::MaybeUninit<B>>> = std::mem::transmute(t);
        let a = std::ptr::read::<A>(u.y.as_ptr() as *const _);
        u.y.as_mut_ptr().write(f(a));
        std::mem::transmute(u)
    }
}

Solution

  • You're confusing yourself with the whole T stuff. This should be much easier to analyze. Specifically, read here.

    use core::mem::ManuallyDrop;
    use core::alloc::Layout;
    
    unsafe trait SharesLayout<T: Sized>: Sized {
        fn assert_same_layout() {
            assert!(Layout::new::<Self>() == Layout::new::<T>());
        }
    }
    
    /// Replaces the contents of a box, without reallocating.
    fn box_map<A, B: SharesLayout<A>>(b: Box<A>, f: impl FnOnce(A) -> B) -> Box<B> {
        unsafe {
            B::assert_same_layout();
            let p = Box::into_raw(b);
            let mut dealloc_on_panic = Box::from_raw(p as *mut ManuallyDrop<A>);
            let new_content = f(ManuallyDrop::take(&mut *dealloc_on_panic));
            std::mem::forget(dealloc_on_panic);
            std::ptr::write(p as *mut B, new_content);
            Box::from_raw(p as *mut B)
        }
    }
    

    Then simply:

    unsafe impl<A, B> SharesLayout<T<A>> for T<B> {}
    
    fn update<A, B>(bt: Box<T<A>>, f: impl FnOnce(A) -> B) -> Box<T<B>> {
        box_map(bt, |t| {
            T { x: t.x, y: Box::new(f(*t.y))}
        })
    }