// SPDX-License-Identifier: ISC

#include <zinnia/alloc.h>
#include <zinnia/assert.h>
#include <zinnia/epoch.h>
#include <zinnia/initgraph.h>
#include <zinnia/lfqueue.h>
#include <zinnia/slab.h>

struct lfqueue_node {
    _Atomic(struct lfqueue_node*) next;
    void* value;
};

static struct slab_cache* node_cache;

static void lfqueue_stage_init(struct boot_info*) {
    node_cache =
        slab_cache_new("lfqueue_node", sizeof(struct lfqueue_node), alignof(struct lfqueue_node), SLAB_EPOCH_SAFE);
    assert(node_cache, "Unable to create the lfqueue node cache");
}
INIT_TASK(lfqueue_stage, "core.lfqueue", lfqueue_stage_init);
INIT_DEPENDS(lfqueue_stage, mem_stage);

static struct lfqueue_node* node_alloc(void* value) {
    struct lfqueue_node* node = slab_cache_alloc(node_cache, 0);
    if (node)
        node->value = value;
    return node;
}

errno_t lfqueue_init(struct lfqueue* queue) {
    struct lfqueue_node* dummy = node_alloc(nullptr);
    if (!dummy)
        return ENOMEM;
    atomic_init(&queue->head, dummy);
    atomic_init(&queue->tail, dummy);
    return 0;
}

void lfqueue_destroy(struct lfqueue* queue) {
    struct lfqueue_node* node = atomic_load(&queue->head, ATOMIC_ACQUIRE);
    while (node) {
        struct lfqueue_node* next = atomic_load(&node->next, ATOMIC_RELAXED);
        slab_cache_free(node_cache, node);
        node = next;
    }
    atomic_store(&queue->head, nullptr, ATOMIC_RELAXED);
    atomic_store(&queue->tail, nullptr, ATOMIC_RELAXED);
}

bool lfqueue_empty(struct lfqueue* queue) {
    epoch_enter();
    struct lfqueue_node* head = epoch_deref(&queue->head);
    const bool empty = atomic_load(&head->next, ATOMIC_RELAXED) == nullptr;
    epoch_exit();
    return empty;
}

errno_t lfqueue_push(struct lfqueue* queue, void* value) {
    struct lfqueue_node* node = node_alloc(value);
    if (!node)
        return ENOMEM;

    epoch_enter();
    while (true) {
        struct lfqueue_node* tail = epoch_deref(&queue->tail);
        struct lfqueue_node* next = epoch_deref(&tail->next);
        if (tail != atomic_load(&queue->tail, ATOMIC_RELAXED))
            continue;
        if (next) {
            atomic_compare_exchange(&queue->tail, &tail, next, ATOMIC_RELEASE, ATOMIC_RELAXED);
            continue;
        }
        if (atomic_compare_exchange(&tail->next, &next, node, ATOMIC_RELEASE, ATOMIC_RELAXED)) {
            atomic_compare_exchange(&queue->tail, &tail, node, ATOMIC_RELEASE, ATOMIC_RELAXED);
            break;
        }
    }
    epoch_exit();
    return 0;
}

bool lfqueue_pop(struct lfqueue* queue, void** out) {
    epoch_enter();
    while (true) {
        struct lfqueue_node* head = epoch_deref(&queue->head);
        struct lfqueue_node* tail = atomic_load(&queue->tail, ATOMIC_RELAXED);
        struct lfqueue_node* next = epoch_deref(&head->next);
        if (head != atomic_load(&queue->head, ATOMIC_RELAXED))
            continue;
        if (!next) {
            epoch_exit();
            return false;
        }
        if (head == tail) {
            atomic_compare_exchange(&queue->tail, &tail, next, ATOMIC_RELEASE, ATOMIC_RELAXED);
            continue;
        }
        void* value = next->value;
        if (atomic_compare_exchange(&queue->head, &head, next, ATOMIC_RELEASE, ATOMIC_RELAXED)) {
            slab_cache_free(node_cache, head);
            epoch_exit();
            *out = value;
            return true;
        }
    }
}