dope-core 0.11.0

use std::{io, ptr::NonNull, slice};

use io_uring::Submitter;
use io_uring::types::BufRingEntry;
use std::io::{Error, ErrorKind};

use crate::backend::uring::tail::Tail;

use super::ffi::Mmap;

struct Entries {
    mem: Mmap,
    count: u16,
}

impl Entries {
    fn new(count: u16) -> io::Result<Self> {
        let size = size_of::<BufRingEntry>()
            .checked_mul(count as usize)
            .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "pbuf_ring size overflow"))?;
        Ok(Self {
            mem: Mmap::new_zeroed(size)?,
            count,
        })
    }

    fn raw_ptr(&self) -> NonNull<BufRingEntry> {
        unsafe { NonNull::new_unchecked(self.mem.as_ptr().cast_mut().cast()) }
    }

    fn as_mut_slice(&mut self) -> &mut [BufRingEntry] {
        unsafe { slice::from_raw_parts_mut(self.mem.as_mut_ptr().cast(), self.count as usize) }
    }

    unsafe fn get_unchecked_mut(&mut self, index: usize) -> &mut BufRingEntry {
        unsafe { self.as_mut_slice().get_unchecked_mut(index) }
    }
}

struct Buffers {
    mem: Mmap,
    buf_len: usize,
}

impl Buffers {
    fn new(count: u16, buf_len: usize) -> io::Result<Self> {
        let total = (count as usize)
            .checked_mul(buf_len)
            .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "pbuf backing size overflow"))?;
        let mut mem = Mmap::new_zeroed(total)?;
        mem.prewarm();
        Ok(Self { mem, buf_len })
    }

    fn buf_len(&self) -> usize {
        self.buf_len
    }

    fn addr(&self, bid: u16) -> u64 {
        self.mem.as_ptr() as u64 + (bid as usize * self.buf_len) as u64
    }

    fn ptr_len(&self, bid: u16, len: usize) -> (*const u8, usize) {
        let len = len.min(self.buf_len);
        let ptr = unsafe { self.mem.as_ptr().add(bid as usize * self.buf_len) };
        (ptr, len)
    }
}

pub(crate) struct Ring {
    tail_pos: u16,
    last_published: u16,
    mask: u16,
    tail: Tail,
    buffers: Buffers,
    entries: Entries,
}

impl Ring {
    pub(crate) const BGID: u16 = 1;

    pub(crate) fn new(submitter: &Submitter<'_>, entries: u16, buf_len: usize) -> io::Result<Self> {
        if !entries.is_power_of_two() {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                "pbuf_ring entries must be power-of-two",
            ));
        }

        let mut entries_mem = Entries::new(entries)?;
        let buffers = Buffers::new(entries, buf_len)?;

        for (bid, e) in entries_mem.as_mut_slice().iter_mut().enumerate() {
            let bid = bid as u16;
            let addr = buffers.addr(bid);
            e.set_addr(addr);
            e.set_len(buf_len as u32);
            e.set_bid(bid);
        }

        let tail = unsafe {
            Tail::new(BufRingEntry::tail(entries_mem.raw_ptr().as_ptr() as *const _) as *mut u16)
        };
        let ring = Self {
            tail_pos: entries,
            last_published: entries,
            mask: entries.wrapping_sub(1),
            tail,
            buffers,
            entries: entries_mem,
        };
        ring.store_tail(entries);

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

        Ok(ring)
    }

    pub(crate) fn buf_len(&self) -> usize {
        self.buffers.buf_len()
    }

    pub(crate) fn ptr_len(&self, bid: u16, len: usize) -> (*const u8, usize) {
        self.buffers.ptr_len(bid & self.mask, len)
    }

    pub(crate) fn defer(&mut self, bid: u16) {
        let bid = bid & self.mask;
        let slot = (self.tail_pos & self.mask) as usize;
        let addr = self.buffers.addr(bid);
        let buf_len = self.buffers.buf_len() as u32;
        let e = unsafe { self.entries.get_unchecked_mut(slot) };
        e.set_addr(addr);
        e.set_len(buf_len);
        e.set_bid(bid);
        self.tail_pos = self.tail_pos.wrapping_add(1);
    }

    pub(crate) fn flush(&mut self) {
        if self.tail_pos == self.last_published {
            return;
        }
        self.store_tail(self.tail_pos);
        self.last_published = self.tail_pos;
    }

    fn store_tail(&self, value: u16) {
        self.tail.publish(value);
    }
}