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      2329021 :     fn as_mut_rust_slice_full_zeroed(&mut self) -> &mut [u8] {
      19      2329021 :         // zero-initialize the uninitialized parts of the buffer so we can create a Rust slice
      20      2329021 :         //
      21      2329021 :         // SAFETY: we own `slice`, don't write outside the bounds
      22      2329021 :         unsafe {
      23      2329021 :             let to_init = self.bytes_total() - self.bytes_init();
      24      2329021 :             self.stable_mut_ptr()
      25      2329021 :                 .add(self.bytes_init())
      26      2329021 :                 .write_bytes(0, to_init);
      27      2329021 :             self.set_init(self.bytes_total());
      28      2329021 :         };
      29      2329021 :         let bytes_total = self.bytes_total();
      30      2329021 :         &mut self[0..bytes_total]
      31      2329021 :     }
      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            6 :     fn test_slice_full_zeroed() {
      44           18 :         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            6 :             let buf = Vec::with_capacity(3);
      49            6 :             let slice: Slice<_> = buf.slice_full();
      50            6 :             assert_eq!(slice.bytes_init(), 0);
      51            6 :             assert_eq!(slice.bytes_total(), 3);
      52            6 :             let rust_slice = &slice[..];
      53            6 :             assert_eq!(
      54            6 :                 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            6 :             let buf = Vec::with_capacity(3);
      63            6 :             let slice: Slice<_> = buf.slice(0..2);
      64            6 :             assert_eq!(slice.bytes_init(), 0);
      65            6 :             assert_eq!(slice.bytes_total(), 2);
      66            6 :             let rust_slice = &slice[..];
      67            6 :             assert_eq!(
      68            6 :                 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            6 :             let buf = Vec::with_capacity(3);
      77            6 :             let mut slice: Slice<_> = buf.slice_full();
      78            6 :             assert_eq!(slice[..].len(), 0);
      79            6 :             let mut file = make_fake_file();
      80            6 :             file.read_exact(&mut slice[..]).unwrap(); // one might think this reads 3 bytes but it reads 0
      81            6 :             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            6 :         {
      88            6 :             // TODO: demo
      89            6 :         }
      90            6 : 
      91            6 :         //
      92            6 :         // Ok, now that we have a shared understanding let's demo how to use the extension trait.
      93            6 :         //
      94            6 : 
      95            6 :         // slice_full()
      96            6 :         {
      97            6 :             let buf = Vec::with_capacity(3);
      98            6 :             let mut slice: Slice<_> = buf.slice_full();
      99            6 :             let rust_slice = slice.as_mut_rust_slice_full_zeroed();
     100            6 :             assert_eq!(rust_slice.len(), 3);
     101            6 :             assert_eq!(rust_slice, &[0, 0, 0]);
     102            6 :             let mut file = make_fake_file();
     103            6 :             file.read_exact(rust_slice).unwrap();
     104            6 :             assert_eq!(rust_slice, b"123");
     105            6 :             assert_eq!(&slice[..], b"123");
     106              :         }
     107              : 
     108              :         // .slice(..)
     109              :         {
     110            6 :             let buf = Vec::with_capacity(3);
     111            6 :             let mut slice: Slice<_> = buf.slice(0..2);
     112            6 :             let rust_slice = slice.as_mut_rust_slice_full_zeroed();
     113            6 :             assert_eq!(rust_slice.len(), 2);
     114            6 :             assert_eq!(rust_slice, &[0, 0]);
     115            6 :             let mut file = make_fake_file();
     116            6 :             file.read_exact(rust_slice).unwrap();
     117            6 :             assert_eq!(rust_slice, b"12");
     118            6 :             assert_eq!(&slice[..], b"12");
     119              :         }
     120            6 :     }
     121              : }