dope-uring 0.2.2

use std::{io, os::fd::AsRawFd, ptr::NonNull};

use crate::DriverConfig;
use crate::driver::{FixedBufGuard, Ring};

pub(crate) struct Resources {
    fixed_buf: FixedBufPool,
    fixed_file: FixedFilePool,
}

pub(crate) struct ResourceToken {
    pool: NonNull<Resources>,
    idx: u16,
}

impl ResourceToken {
    pub(crate) const fn fixed(pool: NonNull<Resources>, idx: u16) -> Self {
        Self { pool, idx }
    }

    pub(crate) fn release(mut self) {
        unsafe { self.pool.as_mut() }.release_fixed(self.idx);
    }
}

fn capability_error(label: &'static str, err: io::Error) -> io::Error {
    io::Error::other(format!(
        "dope: required capability setup failed ({label}): {err}"
    ))
}

impl Resources {
    pub(crate) fn new(ring: &Ring, cfg: &DriverConfig) -> io::Result<Self> {
        Ok(Self {
            fixed_file: FixedFilePool::new(ring, cfg.fixed_file_slots)
                .map_err(|e| capability_error("fixed file registry", e))?,
            fixed_buf: FixedBufPool::new(ring, cfg.fixed_buf_count, cfg.fixed_buf_len)
                .map_err(|e| capability_error("fixed buffer registry", e))?,
        })
    }

    pub(crate) fn release_fixed(&mut self, idx: u16) {
        self.fixed_buf.release(idx);
    }

    pub(in crate::driver) fn fixed_guard(&mut self) -> Option<FixedBufGuard> {
        let buf_index = self.fixed_buf.lease()?;
        let (ptr, cap) = self.fixed_buf.ptr_len(buf_index);
        let pool = NonNull::from(&mut *self);
        Some(FixedBufGuard::new(
            ResourceToken::fixed(pool, buf_index),
            buf_index,
            ptr,
            cap,
        ))
    }

    pub(in crate::driver) fn register_fixed(
        &mut self,
        uring: &Ring,
        io: crate::Fd,
    ) -> io::Result<Option<u32>> {
        self.fixed_file.register(uring, io).map(Some)
    }

    pub(in crate::driver) fn unregister_fixed(&mut self, ring: &Ring, idx: u32) {
        self.fixed_file.unregister(ring, idx);
    }
}

struct FixedBufPool {
    backing: Box<[u8]>,
    buf_len: usize,
    free: Vec<u16>,
}

impl FixedBufPool {
    fn new(ring: &Ring, buf_count: usize, buf_len: usize) -> io::Result<Self> {
        let total = buf_count.checked_mul(buf_len).ok_or_else(|| {
            io::Error::new(io::ErrorKind::InvalidInput, "fixed buf pool size overflow")
        })?;
        let mut backing = vec![0u8; total].into_boxed_slice();
        let base = backing.as_mut_ptr();
        let mut iovecs = Vec::with_capacity(buf_count);
        for i in 0..buf_count {
            let off = i * buf_len;

            iovecs.push(libc::iovec {
                iov_base: unsafe { base.add(off) } as *mut libc::c_void,
                iov_len: buf_len,
            });
        }
        if !iovecs.is_empty() {
            unsafe {
                ring.submitter().register_buffers(&iovecs)?;
            }
        }
        let mut free = Vec::with_capacity(buf_count);
        for i in (0..buf_count as u16).rev() {
            free.push(i);
        }
        Ok(Self {
            backing,
            buf_len,
            free,
        })
    }

    fn lease(&mut self) -> Option<u16> {
        self.free.pop()
    }

    fn release(&mut self, idx: u16) {
        self.free.push(idx);
    }

    fn ptr_len(&self, idx: u16) -> (*mut u8, usize) {
        let off = idx as usize * self.buf_len;

        let ptr = unsafe { self.backing.as_ptr().add(off) } as *mut u8;
        (ptr, self.buf_len)
    }
}

struct FixedFilePool {
    free: Vec<u32>,
}

impl FixedFilePool {
    fn new(uring: &Ring, files: u32) -> io::Result<Self> {
        uring.submitter().register_files_sparse(files)?;
        let mut free = Vec::with_capacity(files as usize);
        for i in 0..files {
            free.push(i);
        }
        Ok(Self { free })
    }

    fn register(&mut self, uring: &Ring, fd: crate::Fd) -> io::Result<u32> {
        let Some(idx) = self.free.pop() else {
            return Err(io::Error::other("fixed file table exhausted"));
        };
        match uring
            .submitter()
            .register_files_update(idx, &[fd.as_raw_fd()])
        {
            Ok(1) => Ok(idx),
            Ok(_) => {
                self.free.push(idx);
                Err(io::Error::other("register_files_update: unexpected count"))
            }
            Err(e) => {
                self.free.push(idx);
                Err(e)
            }
        }
    }

    fn unregister(&mut self, ring: &Ring, idx: u32) {
        let _ = ring.submitter().register_files_update(idx, &[-1]);
        self.free.push(idx);
    }
}