Line data    Source code

       1              : use tokio_epoll_uring::BoundedBuf;
       2              : use tokio_epoll_uring::BoundedBufMut;
       3              : use tokio_epoll_uring::IoBufMut;
       4              : use tokio_epoll_uring::Slice;
       5              : 
       6              : pub(crate) trait SliceMutExt {
       7              :     /// Get a `&mut[0..self.bytes_total()`] slice, for when you need to do borrow-based IO.
       8              :     ///
       9              :     /// See the test case `test_slice_full_zeroed` for the difference to just doing `&slice[..]`
      10              :     fn as_mut_rust_slice_full_zeroed(&mut self) -> &mut [u8];
      11              : }
      12              : 
      13              : impl<B> SliceMutExt for Slice<B>
      14              : where
      15              :     B: IoBufMut,
      16              : {
      17              :     #[inline(always)]
      18       778213 :     fn as_mut_rust_slice_full_zeroed(&mut self) -> &mut [u8] {
      19       778213 :         // zero-initialize the uninitialized parts of the buffer so we can create a Rust slice
      20       778213 :         //
      21       778213 :         // SAFETY: we own `slice`, don't write outside the bounds
      22       778213 :         unsafe {
      23       778213 :             let to_init = self.bytes_total() - self.bytes_init();
      24       778213 :             self.stable_mut_ptr()
      25       778213 :                 .add(self.bytes_init())
      26       778213 :                 .write_bytes(0, to_init);
      27       778213 :             self.set_init(self.bytes_total());
      28       778213 :         };
      29       778213 :         let bytes_total = self.bytes_total();
      30       778213 :         &mut self[0..bytes_total]
      31       778213 :     }
      32              : }
      33              : 
      34              : #[cfg(test)]
      35              : mod tests {
      36              :     use std::io::Read;
      37              : 
      38              :     use super::*;
      39              :     use bytes::Buf;
      40              :     use tokio_epoll_uring::Slice;
      41              : 
      42              :     #[test]
      43            2 :     fn test_slice_full_zeroed() {
      44            6 :         let make_fake_file = || bytes::BytesMut::from(&b"12345"[..]).reader();
      45              : 
      46              :         // before we start the test, let's make sure we have a shared understanding of what slice_full does
      47              :         {
      48            2 :             let buf = Vec::with_capacity(3);
      49            2 :             let slice: Slice<_> = buf.slice_full();
      50            2 :             assert_eq!(slice.bytes_init(), 0);
      51            2 :             assert_eq!(slice.bytes_total(), 3);
      52            2 :             let rust_slice = &slice[..];
      53            2 :             assert_eq!(
      54            2 :                 rust_slice.len(),
      55              :                 0,
      56            0 :                 "Slice only derefs to a &[u8] of the initialized part"
      57              :             );
      58              :         }
      59              : 
      60              :         // and also let's establish a shared understanding of .slice()
      61              :         {
      62            2 :             let buf = Vec::with_capacity(3);
      63            2 :             let slice: Slice<_> = buf.slice(0..2);
      64            2 :             assert_eq!(slice.bytes_init(), 0);
      65            2 :             assert_eq!(slice.bytes_total(), 2);
      66            2 :             let rust_slice = &slice[..];
      67            2 :             assert_eq!(
      68            2 :                 rust_slice.len(),
      69              :                 0,
      70            0 :                 "Slice only derefs to a &[u8] of the initialized part"
      71              :             );
      72              :         }
      73              : 
      74              :         // the above leads to the easy mistake of using slice[..] for borrow-based IO like so:
      75              :         {
      76            2 :             let buf = Vec::with_capacity(3);
      77            2 :             let mut slice: Slice<_> = buf.slice_full();
      78            2 :             assert_eq!(slice[..].len(), 0);
      79            2 :             let mut file = make_fake_file();
      80            2 :             file.read_exact(&mut slice[..]).unwrap(); // one might think this reads 3 bytes but it reads 0
      81            2 :             assert_eq!(&slice[..] as &[u8], &[][..] as &[u8]);
      82              :         }
      83              : 
      84              :         // With owned buffers IO like with VirtualFilem, you could totally
      85              :         // pass in a `Slice` with bytes_init()=0 but bytes_total()=5
      86              :         // and it will read 5 bytes into the slice, and return a slice that has bytes_init()=5.
      87            2 :         {
      88            2 :             // TODO: demo
      89            2 :         }
      90            2 : 
      91            2 :         //
      92            2 :         // Ok, now that we have a shared understanding let's demo how to use the extension trait.
      93            2 :         //
      94            2 : 
      95            2 :         // slice_full()
      96            2 :         {
      97            2 :             let buf = Vec::with_capacity(3);
      98            2 :             let mut slice: Slice<_> = buf.slice_full();
      99            2 :             let rust_slice = slice.as_mut_rust_slice_full_zeroed();
     100            2 :             assert_eq!(rust_slice.len(), 3);
     101            2 :             assert_eq!(rust_slice, &[0, 0, 0]);
     102            2 :             let mut file = make_fake_file();
     103            2 :             file.read_exact(rust_slice).unwrap();
     104            2 :             assert_eq!(rust_slice, b"123");
     105            2 :             assert_eq!(&slice[..], b"123");
     106              :         }
     107              : 
     108              :         // .slice(..)
     109              :         {
     110            2 :             let buf = Vec::with_capacity(3);
     111            2 :             let mut slice: Slice<_> = buf.slice(0..2);
     112            2 :             let rust_slice = slice.as_mut_rust_slice_full_zeroed();
     113            2 :             assert_eq!(rust_slice.len(), 2);
     114            2 :             assert_eq!(rust_slice, &[0, 0]);
     115            2 :             let mut file = make_fake_file();
     116            2 :             file.read_exact(rust_slice).unwrap();
     117            2 :             assert_eq!(rust_slice, b"12");
     118            2 :             assert_eq!(&slice[..], b"12");
     119              :         }
     120            2 :     }
     121              : }