// SPDX-License-Identifier: ISC

#include <zinnia/alloc.h>
#include <zinnia/assert.h>
#include <zinnia/atomic.h>
#include <zinnia/percpu.h>
#include <zinnia/sched.h>
#include <zinnia/spin.h>
#include <zinnia/turnstile.h>
#include <arch/cpu.h>
#include <sched/internal.h>

constexpr size_t TURNSTILE_CHAIN_BITS = 7;
constexpr size_t TURNSTILE_CHAINS = (1 << TURNSTILE_CHAIN_BITS);

struct turnstile_chain {
    // Aligned to cache line to avoid false sharing.
    alignas(CACHE_LINE_SIZE) struct spinlock lock;
    struct turnstile* head;
};

static struct turnstile_chain chains[TURNSTILE_CHAINS];

// Gets the chain from a lock object address.
static struct turnstile_chain* chain_for(const void* obj) {
    uint64_t hash = ((uintptr_t)obj >> 3) * 0x9E3779B97F4A7C15ull;
    return &chains[hash >> (64 - TURNSTILE_CHAIN_BITS)];
}

static struct turnstile* chain_lookup(struct turnstile_chain* chain, const void* obj) {
    struct turnstile* ts = chain->head;
    while (ts && ts->obj != obj)
        ts = ts->chain_next;
    return ts;
}

static void chain_remove(struct turnstile_chain* chain, struct turnstile* ts) {
    struct turnstile** link = &chain->head;
    while (*link != ts)
        link = &(*link)->chain_next;
    *link = ts->chain_next;
}

struct turnstile* turnstile_alloc() {
    struct turnstile* ts = mem_alloc(sizeof(struct turnstile), 0);
    if (ts) {
        for (size_t q = 0; q < TURNSTILE_NUM_QUEUES; q++)
            tailq_init(&ts->waiters[q]);
    }
    return ts;
}

static bool has_waiters(struct turnstile* ts) {
    for (size_t q = 0; q < TURNSTILE_NUM_QUEUES; q++) {
        if (!tailq_empty(&ts->waiters[q]))
            return true;
    }
    return false;
}

void turnstile_free(struct turnstile* ts) {
    assert(ts->obj == nullptr, "Freeing turnstile %p that is still in use", ts);
    mem_free(ts);
}

struct turnstile_chain* turnstile_lock(const void* obj) {
    struct turnstile_chain* chain = chain_for(obj);
    spin_lock(&chain->lock);
    return chain;
}

void turnstile_unlock(struct turnstile_chain* chain) {
    spin_unlock(&chain->lock);
}

static uint8_t task_priority(struct task* task) {
    return atomic_load(&task->priority, ATOMIC_RELAXED);
}

// Inserts `task` into the waiters of `ts`, behind all waiters that are at least as urgent.
static void waiters_insert(struct turnstile* ts, struct task* task) {
    uint8_t prio = task_priority(task);
    struct task* pos = tailq_first(&ts->waiters[task->ts_queue]);
    while (pos && task_priority(pos) >= prio)
        pos = tailq_next(pos, ts_link);

    if (pos)
        tailq_insert_before(pos, task, ts_link);
    else
        tailq_insert_tail(&ts->waiters[task->ts_queue], task, ts_link);
}

// Updates the cached priority of the first waiter. Returns true if it changed.
static bool update_top(struct turnstile* ts) {
    uint8_t top = SCHED_PRIO_MIN;
    for (size_t q = 0; q < TURNSTILE_NUM_QUEUES; q++) {
        struct task* first = tailq_first(&ts->waiters[q]);
        if (first && task_priority(first) > top)
            top = task_priority(first);
    }
    return atomic_exchange(&ts->top_priority, top, ATOMIC_RELAXED) != top;
}

// Recomputes the inherited priority of `task`. Returns true if its effective priority changed.
static bool recompute_locked(struct task* task) {
    uint8_t inherited = SCHED_PRIO_MIN;
    struct turnstile* ts;
    for (ts = tailq_first(&task->pi_lenders); ts; ts = tailq_next(ts, lender_link)) {
        uint8_t top = atomic_load(&ts->top_priority, ATOMIC_RELAXED);
        if (top > inherited)
            inherited = top;
    }
    task->lent_prio = inherited;
    return sched_update_priority_locked(task);
}

// Makes the waiters of `ts` lend their priority to `owner`, and recomputes the old owner's priority.
static void set_owner(struct turnstile* ts, struct task* owner) {
    struct task* old = ts->owner;
    if (old == owner)
        return;

    if (old) {
        spin_lock(&old->lock);
        tailq_remove(&old->pi_lenders, ts, lender_link);
        recompute_locked(old);
        spin_unlock(&old->lock);
    }

    ts->owner = owner;
    if (owner) {
        spin_lock(&owner->lock);
        tailq_insert_tail(&owner->pi_lenders, ts, lender_link);
        spin_unlock(&owner->lock);
    }
}

static void requeue(struct turnstile* ts, struct task* task) {
    tailq_remove(&ts->waiters[task->ts_queue], task, ts_link);
    waiters_insert(ts, task);
}

// Recomputes the priority of the owner of `ts` along the owner chain.
static void propagate(struct turnstile_chain* chain, struct turnstile* ts) {
    bool retried = false;
    while (ts->owner) {
        struct task* owner = ts->owner;
        spin_lock(&owner->lock);
        bool changed = recompute_locked(owner);
        const void* next_obj = atomic_load(&owner->blocked_obj, ATOMIC_RELAXED);
        spin_unlock(&owner->lock);
        if ((!changed && !retried) || next_obj == nullptr)
            break;

        struct turnstile_chain* next_chain = chain_for(next_obj);
        if (next_chain > chain) {
            spin_lock(&next_chain->lock);
        } else if (next_chain < chain && !spin_trylock(&next_chain->lock)) {
            const void* obj = ts->obj;
            turnstile_unlock(chain);
            chain = turnstile_lock(obj);
            ts = chain_lookup(chain, obj);
            if (ts == nullptr)
                break;
            retried = true;
            continue;
        }

        struct turnstile* next = chain_lookup(next_chain, next_obj);
        bool blocked = next && atomic_load(&owner->blocked_ts, ATOMIC_RELAXED) == next;
        if (next_chain != chain)
            turnstile_unlock(chain);
        chain = next_chain;
        if (!blocked)
            break;

        requeue(next, owner);
        if (!update_top(next))
            break;

        ts = next;
        retried = false;
    }
    turnstile_unlock(chain);
}

void turnstile_block(struct turnstile_chain* chain, const void* obj, struct task* owner, enum turnstile_queue queue) {
    struct task* self = percpu_get()->sched.current;
    assert(self->turnstile != nullptr, "Task is already blocked on a turnstile");
    assert(owner != self, "Task tried to block on a lock it owns");

    // Lend our turnstile to the lock if it doesn't have one yet, otherwise donate it.
    struct turnstile* ts = chain_lookup(chain, obj);
    if (ts == nullptr) {
        ts = self->turnstile;
        ts->obj = obj;
        ts->owner = nullptr;
        ts->spare = nullptr;
        ts->chain_next = chain->head;
        chain->head = ts;
    } else {
        self->turnstile->spare = ts->spare;
        ts->spare = self->turnstile;
    }
    self->turnstile = nullptr;

    spin_lock(&self->lock);
    atomic_store(&self->ts_woken, false, ATOMIC_RELAXED);
    atomic_store(&self->blocked_ts, ts, ATOMIC_RELAXED);
    atomic_store(&self->blocked_obj, obj, ATOMIC_RELAXED);
    self->ts_queue = queue;
    waiters_insert(ts, self);
    update_top(ts);
    spin_unlock(&self->lock);

    set_owner(ts, owner);
    propagate(chain, ts);

    // Check if we were actually woken.
    while (!atomic_load(&self->ts_woken, ATOMIC_ACQUIRE))
        sched_block(&percpu_get()->sched);

    spin_lock(&self->lock);
    spin_unlock(&self->lock);
}

struct task* turnstile_first(
    struct turnstile_chain* chain,
    const void* obj,
    enum turnstile_queue queue,
    bool* have_more
) {
    struct turnstile* ts = chain_lookup(chain, obj);
    struct task* first = ts ? tailq_first(&ts->waiters[queue]) : nullptr;
    *have_more = first && tailq_next(first, ts_link) != nullptr;
    return first;
}

void turnstile_wake(struct turnstile_chain* chain, const void* obj, struct task* task, struct task* new_owner) {
    struct turnstile* ts = atomic_load(&task->blocked_ts, ATOMIC_RELAXED);
    assert(ts && ts->obj == obj, "Task %p isn't blocked on %p", task, obj);

    tailq_remove(&ts->waiters[task->ts_queue], task, ts_link);
    atomic_store(&task->blocked_ts, nullptr, ATOMIC_RELAXED);
    atomic_store(&task->blocked_obj, nullptr, ATOMIC_RELAXED);

    if (!has_waiters(ts)) {
        // The last waiter takes the lock's turnstile with it.
        set_owner(ts, nullptr);
        chain_remove(chain, ts);
        ts->obj = nullptr;
        task->turnstile = ts;
    } else {
        task->turnstile = ts->spare;
        ts->spare = ts->spare->spare;
        update_top(ts);
        set_owner(ts, new_owner);
        if (new_owner) {
            spin_lock(&new_owner->lock);
            recompute_locked(new_owner);
            spin_unlock(&new_owner->lock);
        }
    }

    spin_lock(&task->lock);
    atomic_store(&task->ts_woken, true, ATOMIC_RELEASE);
    sched_wake_locked(task);
    spin_unlock(&task->lock);
}

void turnstile_reprioritize(struct task* task) {
    const void* obj = atomic_load(&task->blocked_obj, ATOMIC_RELAXED);
    if (obj == nullptr)
        return;

    struct turnstile_chain* chain = turnstile_lock(obj);
    struct turnstile* ts = chain_lookup(chain, obj);
    if (ts == nullptr || atomic_load(&task->blocked_ts, ATOMIC_RELAXED) != ts) {
        turnstile_unlock(chain);
        return;
    }

    requeue(ts, task);
    if (!update_top(ts)) {
        turnstile_unlock(chain);
        return;
    }
    propagate(chain, ts);
}