Get a `&mut u8` reference to part of a `&mut u32`

To convert a u32 to its bytes, I know there is already:

• u32::to_le_bytes
• u32::to_be_bytes
• u32::to_ne_bytes
• u32::from_le_bytes
• u32::from_be_bytes
• u32::from_ne_bytes

However, is there any sound and cross-platform way to convert a &mut u32 into references of its bytes, as &mut u8?

Like this:

fn as_be_bytes_mut(num: &mut u32) -> [&mut u8; 4] {
    todo!()
}

fn main() {
    let mut num: u32 = 0x12345678;

    // Prints `0x12345678 - [12, 34, 56, 78]`
    println!("{:#x} - {:x?}", num, num.to_be_bytes());

    let parts = as_be_bytes_mut(&mut num);
    *parts[2] = 0xfe;

    // Should print `0x1234fe78 - [12, 34, fe, 78]`
    println!("{:#x} - {:x?}", num, num.to_be_bytes());
}

The rationale is that it should be theoretically possible, because there are no invalid states of a u32, no matter how you modify its underlying bytes.

Attempt:

fn as_ne_bytes_mut(num: &mut u32) -> [&mut u8; 4] {
    unsafe {
        let ptr: *mut u8 = (num as *mut u32).cast();
        [
            &mut *ptr.add(0),
            &mut *ptr.add(1),
            &mut *ptr.add(2),
            &mut *ptr.add(3),
        ]
    }
}

fn main() {
    let mut num: u32 = 0x12345678;

    println!("{:#x} - {:x?}", num, num.to_be_bytes());

    let parts = as_ne_bytes_mut(&mut num);
    *parts[1] = 0xfe;

    println!("{:#x} - {:x?}", num, num.to_be_bytes());
}

0x12345678 - [12, 34, 56, 78]
0x1234fe78 - [12, 34, fe, 78]

I think (tm) this is sound, because a u32 is a packed array of u8's and u8's are always correctly aligned, and the lifetimes should also match. I didn't find a way yet to implement as_be/le_bytes_mut yet. Also I'm not 100% sure this is sound, so some feedback would help.

It should be sound to get a &mut [u8; 4] out of it.

pub fn as_ne_bytes_mut(num: &mut u32) -> &mut [u8; 4] {
    unsafe {
        let arr: *mut [u8; 4] = (num as *mut u32).cast();
        &mut *arr
    }
}

This is much better than [&mut u8; 4] since it's a no-op. However, you aren't going to be able to create le and be versions unless you swap the bits around and return an impl DerefMut<Target = &mut [u8; 4]> guard that swaps the bits back when it's dropped.

I would probably go with four functions that return one byte each instead of le and be functions that return [&mut u8; 4]. Although if you use them the same, they should optimize to the same thing.

pub fn most_sig_byte_0(num: &mut u32) -> &mut u8 {
    if cfg!(target_endian = "big") {
        &mut as_ne_bytes_mut(num)[0]
    } else {
        &mut as_ne_bytes_mut(num)[3]
    }
}

pub fn most_sig_byte_1(num: &mut u32) -> &mut u8 {
    if cfg!(target_endian = "big") {
        &mut as_ne_bytes_mut(num)[1]
    } else {
        &mut as_ne_bytes_mut(num)[2]
    }
}

pub fn most_sig_byte_2(num: &mut u32) -> &mut u8 {
    if cfg!(target_endian = "big") {
        &mut as_ne_bytes_mut(num)[2]
    } else {
        &mut as_ne_bytes_mut(num)[1]
    }
}

pub fn most_sig_byte_3(num: &mut u32) -> &mut u8 {
    if cfg!(target_endian = "big") {
        &mut as_ne_bytes_mut(num)[3]
    } else {
        &mut as_ne_bytes_mut(num)[0]
    }
}

Miri thinks these are fine, at least. You could make them a single const generic function, but it would be less ergonomic to give it the 0..4 bound. (playground)

This can be achieved with the help of cfg(target_endian = "..."):

pub fn as_ne_bytes_mut(num: &mut u32) -> [&mut u8; 4] {
    unsafe {
        let ptr: *mut u8 = (num as *mut u32).cast();
        [
            &mut *ptr.add(0),
            &mut *ptr.add(1),
            &mut *ptr.add(2),
            &mut *ptr.add(3),
        ]
    }
}

pub fn as_be_bytes_mut(num: &mut u32) -> [&mut u8; 4] {
    let mut b = as_ne_bytes_mut(num);

    #[cfg(target_endian = "little")]
    b.reverse();

    b
}

pub fn as_le_bytes_mut(num: &mut u32) -> [&mut u8; 4] {
    let mut b = as_be_bytes_mut(num);
    b.reverse();
    b
}

fn main() {
    let mut num: u32 = 0x12345678;

    // Prints `0x12345678 - [12, 34, 56, 78]`
    println!("{:#x} - {:x?}", num, num.to_be_bytes());

    let parts = as_be_bytes_mut(&mut num);
    *parts[2] = 0xfe;

    // Should print `0x1234fe78 - [12, 34, fe, 78]`
    println!("{:#x} - {:x?}", num, num.to_be_bytes());
}

0x12345678 - [12, 34, 56, 78]
0x1234fe78 - [12, 34, fe, 78]

While it does look a little convoluted, the compiler manages to optimize it perfectly:

example::as_ne_bytes_mut:
        mov     rax, rdi
        lea     rcx, [rsi + 1]
        lea     rdx, [rsi + 2]
        mov     qword ptr [rdi], rsi
        add     rsi, 3
        mov     qword ptr [rdi + 8], rcx
        mov     qword ptr [rdi + 16], rdx
        mov     qword ptr [rdi + 24], rsi
        ret

example::as_be_bytes_mut:
        mov     rax, rdi
        lea     rcx, [rsi + 1]
        lea     rdx, [rsi + 2]
        lea     rdi, [rsi + 3]
        mov     qword ptr [rax], rdi
        mov     qword ptr [rax + 24], rsi
        mov     qword ptr [rax + 8], rdx
        mov     qword ptr [rax + 16], rcx
        ret

example::as_le_bytes_mut:
        mov     rax, rdi
        lea     rcx, [rsi + 1]
        lea     rdx, [rsi + 2]
        mov     qword ptr [rdi], rsi
        add     rsi, 3
        mov     qword ptr [rdi + 24], rsi
        mov     qword ptr [rdi + 8], rcx
        mov     qword ptr [rdi + 16], rdx
        ret

It is possible to do this without pointer arithmetics.

Code:

pub fn as_le_bytes_mut(num: &mut u32)->[&mut u8; 4]{
    let num_slice = std::slice::from_mut(num);
    let (pref, middle, suff): (_,&mut [u8],_) = unsafe{
        num_slice.align_to_mut()
    };
    // This would be always true when we cast to u8
    assert!(pref.is_empty() && suff.is_empty());
    
    match middle {
        #[cfg(target_endian = "little")]
        [a, b, c, d] => [a, b, c, d],
        #[cfg(target_endian = "big")]
        [a, b, c, d] => [d, c, b, a],
        _ => unreachable!()
    }
}

pub fn as_be_bytes_mut(num: &mut u32)->[&mut u8; 4]{
    let mut r = as_le_bytes_mut(num);
    r.reverse();
    r
}

And it compiles to nice assembly too: godbolt link.