tempest-io 0.0.1

use bytes::{Bytes, BytesMut};

/// # Safety
///
/// Implementors must ensure [`stable_ptr`](IoBuf::stable_ptr) returns a pointer
/// that remains valid and pinned for the entire duration of any I/O operation
/// the buffer is passed to.
pub unsafe trait IoBuf {
    /// Returns a stable pointer to the start of the initialized data.
    fn stable_ptr(&self) -> *const u8;

    /// Number of initialized bytes, valid to read.
    fn bytes_init(&self) -> usize;

    /// Total capacity, including uninitialized memory.
    /// This is what gets passed to the kernel for reads.
    fn bytes_total(&self) -> usize;

    /// Returns a bounded view into this buffer.
    fn slice(self, range: impl std::ops::RangeBounds<usize>) -> Slice<Self>
    where
        Self: Sized,
    {
        let begin = match range.start_bound() {
            std::ops::Bound::Included(&n) => n,
            std::ops::Bound::Unbounded => 0,
            _ => unreachable!(),
        };
        let end = match range.end_bound() {
            std::ops::Bound::Excluded(&n) => n,
            std::ops::Bound::Included(&n) => n + 1,
            std::ops::Bound::Unbounded => self.bytes_total(),
        };
        assert!(begin <= end && end <= self.bytes_total());
        Slice::new(self, begin, end)
    }
}

/// # Safety
///
/// Implementors must ensure [`stable_mut_ptr`](IoBufMut::stable_mut_ptr) returns
/// a pointer that remains valid and pinned for the entire duration of any I/O
/// operation the buffer is passed to.
pub unsafe trait IoBufMut: IoBuf {
    /// Returns a stable mutable pointer to the start of the buffer.
    fn stable_mut_ptr(&mut self) -> *mut u8;

    /// Mark the first `pos` bytes as initialized.
    ///
    /// # Safety
    ///
    /// Caller must ensure all bytes from `stable_mut_ptr()` up to `pos`
    /// are actually initialized and may be read.
    unsafe fn set_init(&mut self, pos: usize);

    /// Clear this buffer, setting the initialized bytes to 0.
    fn clear(&mut self) {
        // SAFETY: it is always safe, since we dont assume anything is initialized
        unsafe { self.set_init(0) };
    }
}

/// A bounded view into an owned [`IoBuf`] or [`IoBufMut`] buffer.
///
/// Allows passing a subset of a buffer to an I/O operation without
/// splitting or consuming it. The original buffer is reclaimed via
/// [`into_inner`](Slice::into_inner) after the operation completes.
// TODO: add example section once VirtualIo is implemented
#[derive(Debug)]
pub struct Slice<B> {
    buf: B,
    begin: usize,
    end: usize,
}

impl<B: IoBuf> Slice<B> {
    /// Creates a new slice into `buf` bounded by `[begin, end)`.
    ///
    /// # Panics
    ///
    /// Panics if `begin > end` or `end > buf.bytes_total()`.
    pub fn new(buf: B, begin: usize, end: usize) -> Self {
        assert!(begin <= end && end <= buf.bytes_total());
        Self { buf, begin, end }
    }

    /// Consumes the slice and returns the underlying buffer.
    pub fn into_inner(self) -> B {
        self.buf
    }

    /// Returns the length of initialized bytes of this slice.
    pub fn len(&self) -> usize {
        self.bytes_init()
    }

    /// Returns `true` if the length of this slice is zero.
    pub fn is_empty(&self) -> bool {
        self.bytes_init() == 0
    }
}

unsafe impl<B: IoBuf> IoBuf for Slice<B> {
    fn stable_ptr(&self) -> *const u8 {
        unsafe { self.buf.stable_ptr().add(self.begin) }
    }

    fn bytes_init(&self) -> usize {
        // slice may reach into uninitialized bytes
        self.buf
            .bytes_init()
            .saturating_sub(self.begin)
            .min(self.end - self.begin)
    }

    fn bytes_total(&self) -> usize {
        self.end - self.begin
    }
}

unsafe impl<B: IoBufMut> IoBufMut for Slice<B> {
    fn stable_mut_ptr(&mut self) -> *mut u8 {
        unsafe { self.buf.stable_mut_ptr().add(self.begin) }
    }

    unsafe fn set_init(&mut self, pos: usize) {
        // pos is relative to the slice start, adjust to the full buffer
        unsafe { self.buf.set_init(self.begin + pos) };
    }
}

impl<B: IoBuf> AsRef<[u8]> for Slice<B> {
    fn as_ref(&self) -> &[u8] {
        let ptr = self.stable_ptr();
        let len = self.bytes_init();
        // SAFETY: ptr is valid and initialized up until ptr + len
        unsafe { std::slice::from_raw_parts(ptr, len) }
    }
}

pub enum MaybeRegistered<R, F> {
    Registered(R),
    Fallback(F),
}

unsafe impl<R: IoBuf, F: IoBuf> IoBuf for MaybeRegistered<R, F> {
    fn stable_ptr(&self) -> *const u8 {
        match self {
            Self::Registered(r) => r.stable_ptr(),
            Self::Fallback(f) => f.stable_ptr(),
        }
    }

    fn bytes_init(&self) -> usize {
        match self {
            Self::Registered(r) => r.bytes_init(),
            Self::Fallback(f) => f.bytes_init(),
        }
    }

    fn bytes_total(&self) -> usize {
        match self {
            Self::Registered(r) => r.bytes_total(),
            Self::Fallback(f) => f.bytes_total(),
        }
    }
}

unsafe impl<R: IoBufMut, F: IoBufMut> IoBufMut for MaybeRegistered<R, F> {
    fn stable_mut_ptr(&mut self) -> *mut u8 {
        match self {
            Self::Registered(r) => r.stable_mut_ptr(),
            Self::Fallback(f) => f.stable_mut_ptr(),
        }
    }

    unsafe fn set_init(&mut self, pos: usize) {
        // SAFETY: R and F are required to implement set_init correctly
        unsafe {
            match self {
                Self::Registered(r) => r.set_init(pos),
                Self::Fallback(f) => f.set_init(pos),
            }
        }
    }
}

unsafe impl<const N: usize> IoBuf for [u8; N] {
    fn stable_ptr(&self) -> *const u8 {
        self.as_ptr()
    }

    fn bytes_init(&self) -> usize {
        N
    }

    fn bytes_total(&self) -> usize {
        N
    }
}

unsafe impl<const N: usize> IoBufMut for [u8; N] {
    fn stable_mut_ptr(&mut self) -> *mut u8 {
        self.as_mut_ptr()
    }

    unsafe fn set_init(&mut self, _pos: usize) {} // always fully initialized
}

unsafe impl IoBuf for &'static [u8] {
    fn stable_ptr(&self) -> *const u8 {
        self.as_ptr()
    }

    fn bytes_init(&self) -> usize {
        self.len()
    }

    fn bytes_total(&self) -> usize {
        self.len()
    }
}

unsafe impl IoBuf for Bytes {
    fn stable_ptr(&self) -> *const u8 {
        self.as_ptr()
    }

    fn bytes_init(&self) -> usize {
        self.len()
    }

    fn bytes_total(&self) -> usize {
        // `Bytes` does not have a capacity, since it is immutable
        self.len()
    }
}

unsafe impl IoBuf for BytesMut {
    fn stable_ptr(&self) -> *const u8 {
        self.as_ptr()
    }

    fn bytes_init(&self) -> usize {
        self.len()
    }

    fn bytes_total(&self) -> usize {
        self.capacity()
    }
}

unsafe impl IoBufMut for BytesMut {
    fn stable_mut_ptr(&mut self) -> *mut u8 {
        self.as_mut_ptr()
    }

    unsafe fn set_init(&mut self, pos: usize) {
        // SAFETY: caller must ensure all bytes up to `pos` are valid
        unsafe { self.set_len(pos) };
    }
}

unsafe impl IoBuf for Vec<u8> {
    fn stable_ptr(&self) -> *const u8 {
        self.as_ptr()
    }
    fn bytes_init(&self) -> usize {
        self.len()
    }
    fn bytes_total(&self) -> usize {
        self.capacity()
    }
}

unsafe impl IoBufMut for Vec<u8> {
    fn stable_mut_ptr(&mut self) -> *mut u8 {
        self.as_mut_ptr()
    }

    unsafe fn set_init(&mut self, pos: usize) {
        // SAFETY: caller ensures bytes up to pos are initialized
        unsafe { self.set_len(pos) };
    }
}