// A minimal io_uring ring, on the raw syscalls. Only io.h uses it
//
// We write SQEs into the SQ and the kernel takes them when we submit. The
// kernel writes CQEs into the CQ and we take them. Each side moves only its
// own index, and reads the other's with acquire, so it sees the entries that
// came before
#pragma once

#ifdef JSRT_IMPLEMENTATION
#include <errno.h>
#include <linux/io_uring.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <unistd.h>

typedef struct {
    int fd;

    // The SQ: we own the tail, the kernel owns the head. `array` maps each
    // slot of the ring to an SQE, and we keep it the identity
    unsigned *sq_head, *sq_tail, *sq_array;
    unsigned sq_mask, sq_entries;
    struct io_uring_sqe *sqes;
    unsigned tail;                  // our copy of the SQ tail, ahead of it until submit
    unsigned to_submit;             // SQEs written that the kernel did not take yet

    // The CQ: the kernel owns the tail, we own the head
    unsigned *cq_head, *cq_tail;
    unsigned cq_mask, cq_entries;
    struct io_uring_cqe *cqes;
} jsrt_uring;

_Noreturn static void jsrt_uring_fail(const char *what) {
    fprintf(stderr, "jsrt: %s: %s\n", what, strerror(errno));
    exit(1);
}

// The features that io.h depends on, all from Linux 5.5:
// SINGLE_MMAP  the SQ and the CQ rings are one mapping
// NODROP       the kernel keeps CQEs that do not fit the CQ, never drops them
// SUBMIT_STABLE  the kernel copies what an SQE points to (a timespec, a path)
//              before submit returns, so it needs to stay only until then
static void jsrt_uring_init(jsrt_uring *r, unsigned entries) {
    struct io_uring_params p;
    memset(&p, 0, sizeof p);

    int fd = (int)syscall(__NR_io_uring_setup, entries, &p);

    if (fd < 0) {
        jsrt_uring_fail("io_uring_setup");
    }

    unsigned need = IORING_FEAT_SINGLE_MMAP | IORING_FEAT_NODROP | IORING_FEAT_SUBMIT_STABLE;

    if ((p.features & need) != need) {
        fputs("jsrt: io_uring: the kernel is older than Linux 5.5\n", stderr);
        exit(1);
    }

    size_t sq_size = p.sq_off.array + p.sq_entries * sizeof(unsigned);
    size_t cq_size = p.cq_off.cqes + p.cq_entries * sizeof(struct io_uring_cqe);
    size_t size = sq_size > cq_size ? sq_size : cq_size;
    int prot = PROT_READ | PROT_WRITE, flags = MAP_SHARED | MAP_POPULATE;

    char *ring = mmap(NULL, size, prot, flags, fd, IORING_OFF_SQ_RING);
    void *sqes = mmap(NULL, p.sq_entries * sizeof(struct io_uring_sqe), prot, flags, fd, IORING_OFF_SQES);

    if (ring == MAP_FAILED || sqes == MAP_FAILED) {
        jsrt_uring_fail("io_uring mmap");
    }

    *r = (jsrt_uring){
        .fd = fd,
        .sq_head = (unsigned *)(ring + p.sq_off.head),
        .sq_tail = (unsigned *)(ring + p.sq_off.tail),
        .sq_array = (unsigned *)(ring + p.sq_off.array),
        .sq_mask = *(unsigned *)(ring + p.sq_off.ring_mask),
        .sq_entries = p.sq_entries,
        .sqes = sqes,
        .tail = *(unsigned *)(ring + p.sq_off.tail),
        .cq_head = (unsigned *)(ring + p.cq_off.head),
        .cq_tail = (unsigned *)(ring + p.cq_off.tail),
        .cq_mask = *(unsigned *)(ring + p.cq_off.ring_mask),
        .cq_entries = p.cq_entries,
        .cqes = (struct io_uring_cqe *)(ring + p.cq_off.cqes),
    };
}

// A new SQE, all zero, or NULL if the SQ is full. The kernel sees it at the
// next submit
static struct io_uring_sqe *jsrt_uring_sqe(jsrt_uring *r) {
    unsigned head = __atomic_load_n(r->sq_head, __ATOMIC_ACQUIRE);

    if (r->tail - head == r->sq_entries) {
        return NULL;
    }

    unsigned i = r->tail & r->sq_mask;
    r->sq_array[i] = i;
    r->tail++;
    r->to_submit++;

    memset(&r->sqes[i], 0, sizeof r->sqes[i]);
    return &r->sqes[i];
}

static bool jsrt_uring_cq_empty(jsrt_uring *r) {
    return *r->cq_head == __atomic_load_n(r->cq_tail, __ATOMIC_ACQUIRE);
}

// Submits every SQE written, and if `wait`, waits until the CQ has a CQE
//
// The kernel can take fewer SQEs than we give, for example when one fails
// before it starts. That one still gets a CQE, and we give the rest again
static void jsrt_uring_submit(jsrt_uring *r, bool wait) {
    __atomic_store_n(r->sq_tail, r->tail, __ATOMIC_RELEASE);

    for (;;) {
        bool get = wait && jsrt_uring_cq_empty(r);

        if (!r->to_submit && !get) {
            return;
        }

        unsigned flags = get ? IORING_ENTER_GETEVENTS : 0;
        int n = (int)syscall(__NR_io_uring_enter, r->fd, r->to_submit, get ? 1 : 0, flags, NULL, 0);

        if (n < 0) {
            if (errno == EINTR) {
                continue;
            }

            jsrt_uring_fail("io_uring_enter");
        }

        r->to_submit -= (unsigned)n;
    }
}

// Takes the next CQE, if there is one
static bool jsrt_uring_cqe(jsrt_uring *r, uint64_t *user_data, int32_t *res) {
    if (jsrt_uring_cq_empty(r)) {
        return false;
    }

    unsigned head = *r->cq_head;
    struct io_uring_cqe *cqe = &r->cqes[head & r->cq_mask];
    *user_data = cqe->user_data;
    *res = cqe->res;

    __atomic_store_n(r->cq_head, head + 1, __ATOMIC_RELEASE);
    return true;
}

#endif // JSRT_IMPLEMENTATION