Skip to content

Commit

Permalink
Implement BufMut for &mut [MaybeUninit<u8>]
Browse files Browse the repository at this point in the history
This will allow using BufMut without having to initialise memory
first.  In particular, this will allow simple integration with tokio:

    #![feature(maybe_uninit_slice)]

    use tokio::io::AsyncReadExt;

    let mut tokio_file = /* ... */;
    let mut buf = [MaybeUninit::uninit(); 4096];
    let len = tokio_file.read_buf(&mut buf);
    // SAFETY: read above initialised first len bytes of the buffer.
    let buf = unsafe {
        MaybeUninit::slice_assume_init_mut(&buf[..len])
    };
  • Loading branch information
mina86 committed Feb 9, 2023
1 parent 74b04c7 commit 72ef20d
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 0 deletions.
35 changes: 35 additions & 0 deletions src/buf/buf_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1419,6 +1419,41 @@ unsafe impl BufMut for &mut [u8] {
}
}

unsafe impl BufMut for &mut [core::mem::MaybeUninit<u8>] {
#[inline]
fn remaining_mut(&self) -> usize {
self.len()
}

#[inline]
fn chunk_mut(&mut self) -> &mut UninitSlice {
UninitSlice::from_slice(self)
}

#[inline]
unsafe fn advance_mut(&mut self, cnt: usize) {
// Lifetime dance taken from `impl Write for &mut [u8]`.
let (_, b) = core::mem::replace(self, &mut []).split_at_mut(cnt);
*self = b;
}

#[inline]
fn put_slice(&mut self, src: &[u8]) {
self.chunk_mut()[..src.len()].copy_from_slice(src);
unsafe {
self.advance_mut(src.len());
}
}

fn put_bytes(&mut self, val: u8, cnt: usize) {
assert!(self.remaining_mut() >= cnt);
unsafe {
ptr::write_bytes(self.as_mut_ptr() as *mut u8, val, cnt);
self.advance_mut(cnt);
}
}
}

unsafe impl BufMut for Vec<u8> {
#[inline]
fn remaining_mut(&self) -> usize {
Expand Down
36 changes: 36 additions & 0 deletions tests/test_buf_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
use bytes::buf::UninitSlice;
use bytes::{BufMut, BytesMut};
use core::fmt::Write;
use core::mem::{transmute, MaybeUninit};
use core::usize;

#[test]
Expand Down Expand Up @@ -121,6 +122,41 @@ fn test_slice_put_bytes() {
assert_eq!(&[17, 19, 19, 0], &v[..]);
}

#[test]
fn test_uninit_buf_mut() {
let mut buf = [b'X'; 100];
let fill = [b'Z'; 100];

fn make(slice: &mut [u8]) -> &mut [MaybeUninit<u8>] {
// SAFETY: [T] and [MaybeUninit<T>] have the same layout.
unsafe { transmute(slice) }
}

{
let mut slice = make(&mut buf[..]);
slice.put_bytes(b'A', 2);
slice.put_slice(b"BBCC");
assert_eq!(b"AABBCCXX", &buf[..8]);
}

for buf_len in 0..100 {
for fill_len in 0..=buf_len {
let mut slice = make(&mut buf[..buf_len]);
slice.put_slice(&fill[..fill_len]);
assert_eq!(buf_len - fill_len, slice.remaining_mut());
}
}
}

#[test]
#[should_panic]
fn test_uninit_buf_mut_overflow() {
let mut buf = [b'X'; 4];
// SAFETY: [T] and [MaybeUninit<T>] have the same layout.
let mut slice: &mut [MaybeUninit<u8>] = unsafe { transmute(&mut buf[..]) };
slice.put_slice(b"AABBCC");
}

#[test]
fn test_deref_bufmut_forwards() {
struct Special;
Expand Down

0 comments on commit 72ef20d

Please sign in to comment.