dope-uring 0.2.2

use std::{
    alloc::{self, Layout},
    io,
    sync::atomic::{AtomicU16, Ordering},
};

use dope_core::driver::ProvidedRingOps;
use io_uring::types::BufRingEntry;

use crate::driver::{ProvidedBuf, Ring};

pub struct ProvidedRing {
    tail_pos: u16,
    last_published: u16,
    mask: u16,
    buf_len: u32,
    buf_base: *mut u8,
    entries: *mut BufRingEntry,
    tail: *mut AtomicU16,
    _buf_mem: Box<[u8]>,
    entries_layout: Layout,
}

impl ProvidedRing {
    const BGID: u16 = 1;
    const RING_ALIGN: usize = 16384;

    pub(crate) fn new(ring: &Ring, entries: u16, buf_len: usize) -> io::Result<Self> {
        if !entries.is_power_of_two() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "pbuf_ring entries must be power-of-two",
            ));
        }
        let entry_size = std::mem::size_of::<BufRingEntry>();
        let alloc_len = entry_size.checked_mul(entries as usize).ok_or_else(|| {
            io::Error::new(io::ErrorKind::InvalidInput, "pbuf_ring size overflow")
        })?;
        let entries_layout = Layout::from_size_align(alloc_len, Self::RING_ALIGN)
            .map_err(|_| io::Error::other("pbuf_ring layout"))?;
        let entries_ptr = unsafe { alloc::alloc_zeroed(entries_layout) } as *mut BufRingEntry;
        if entries_ptr.is_null() {
            return Err(io::Error::other("pbuf_ring alloc"));
        }

        let total_buf_len = (entries as usize).checked_mul(buf_len).ok_or_else(|| {
            io::Error::new(io::ErrorKind::InvalidInput, "pbuf backing size overflow")
        })?;
        let mut buf_mem = vec![0u8; total_buf_len].into_boxed_slice();
        let buf_base = buf_mem.as_mut_ptr();

        for bid in 0..entries {
            let e = unsafe { &mut *entries_ptr.add(bid as usize) };
            let addr = unsafe { buf_base.add(bid as usize * buf_len) };
            e.set_addr(addr as u64);
            e.set_len(buf_len as u32);
            e.set_bid(bid);
        }

        let tail =
            unsafe { BufRingEntry::tail(entries_ptr as *const _) as *mut u16 } as *mut AtomicU16;
        unsafe { (*tail).store(entries, Ordering::Release) };

        unsafe {
            ring.submitter().register_buf_ring_with_flags(
                entries_ptr as u64,
                entries,
                Self::BGID,
                0,
            )?;
        }

        Ok(Self {
            tail_pos: entries,
            last_published: entries,
            mask: entries.wrapping_sub(1),
            buf_len: buf_len as u32,
            buf_base,
            entries: entries_ptr,
            tail,
            _buf_mem: buf_mem,
            entries_layout,
        })
    }

    pub(crate) fn group(&self) -> (u16, u32) {
        (Self::BGID, self.buf_len)
    }

    pub(crate) fn make_buf(&mut self, bid: u16, len: usize) -> ProvidedBuf {
        let (ptr, cap) = self.ptr_len(bid);
        ProvidedBuf::new(bid, ptr, cap, len)
    }

    pub(crate) fn release(&mut self, bid: u16) {
        self.defer(bid);
        unsafe { (*self.tail).store(self.tail_pos, Ordering::Release) };
        self.last_published = self.tail_pos;
    }

    pub(crate) fn flush(&mut self) -> u64 {
        let n = self.tail_pos.wrapping_sub(self.last_published);
        if n == 0 {
            return 0;
        }
        unsafe { (*self.tail).store(self.tail_pos, Ordering::Release) };
        self.last_published = self.tail_pos;
        n as u64
    }
}

impl ProvidedRingOps for ProvidedRing {
    #[inline(always)]
    fn ptr_len(&self, bid: u16) -> (*mut u8, usize) {
        let off = bid as usize * self.buf_len as usize;
        let ptr = unsafe { self.buf_base.add(off) };
        (ptr, self.buf_len as usize)
    }

    #[inline(always)]
    fn defer(&mut self, bid: u16) {
        let slot = (self.tail_pos & self.mask) as usize;
        let e = unsafe { &mut *self.entries.add(slot) };
        let addr = unsafe { self.buf_base.add(bid as usize * self.buf_len as usize) };
        e.set_addr(addr as u64);
        e.set_len(self.buf_len);
        e.set_bid(bid);
        self.tail_pos = self.tail_pos.wrapping_add(1);
    }
}

impl Drop for ProvidedRing {
    fn drop(&mut self) {
        unsafe { alloc::dealloc(self.entries as *mut u8, self.entries_layout) };
    }
}