#pragma once
// ke_spinlock.h - spinlock for short, non-sleeping critical sections
//
// Step 3 IRQL bring-up:
//   * Acquiring a spinlock disables maskable interrupts and raises the software
//     IRQL to at least DISPATCH_LEVEL.
//   * Releasing restores both the previous IRQL and the previous RFLAGS.IF
//     state through the returned cookie.
//
// The public type remains uintptr_t so existing callers do not need to change.
// Never hold a spinlock across a wait/blocking operation.
#include <stdint.h>
#include <intrin.h>
#include <ke_irql.h>

typedef struct SPINLOCK {
    volatile long Locked;
} SPINLOCK;

#define SPINLOCK_INIT  { 0 }

#define KI_SPINLOCK_IRQL_SHIFT  ((unsigned)(sizeof(uintptr_t) * 8u - 8u))
#define KI_SPINLOCK_FLAGS_MASK  (~(((uintptr_t)0xFFu) << KI_SPINLOCK_IRQL_SHIFT))

static __forceinline uintptr_t KiPackSpinlockCookie(uintptr_t flags, KIRQL oldIrql) {
    return (flags & KI_SPINLOCK_FLAGS_MASK) |
           (((uintptr_t)oldIrql) << KI_SPINLOCK_IRQL_SHIFT);
}

static __forceinline uintptr_t KiSpinlockCookieFlags(uintptr_t cookie) {
    return cookie & KI_SPINLOCK_FLAGS_MASK;
}

static __forceinline KIRQL KiSpinlockCookieIrql(uintptr_t cookie) {
    return (KIRQL)(cookie >> KI_SPINLOCK_IRQL_SHIFT);
}

// Save RFLAGS, mask interrupts, raise to DISPATCH_LEVEL if needed, acquire.
// If already running at DIRQL, the higher IRQL is preserved.
static __forceinline uintptr_t KeAcquireSpinlock(SPINLOCK* sl) {
    uintptr_t flags = __readeflags();
    _disable();

    KIRQL oldIrql = KeGetCurrentIrql();
    if (oldIrql < DISPATCH_LEVEL)
        KeRaiseIrql(DISPATCH_LEVEL);

    while (_InterlockedCompareExchange(&sl->Locked, 1, 0) != 0)
        _mm_pause();

    return KiPackSpinlockCookie(flags, oldIrql);
}

// Release and restore the caller's previous software IRQL/RFLAGS state.
static __forceinline void KeReleaseSpinlock(SPINLOCK* sl, uintptr_t cookie) {
    uintptr_t flags = KiSpinlockCookieFlags(cookie);
    KIRQL oldIrql   = KiSpinlockCookieIrql(cookie);

    _InterlockedExchange(&sl->Locked, 0);
    KeLowerIrql(oldIrql);
    __writeeflags(flags);
}

static __forceinline int KeTryAcquireSpinlock(SPINLOCK* sl, uintptr_t* outCookie) {
    uintptr_t flags = __readeflags();
    _disable();

    KIRQL oldIrql = KeGetCurrentIrql();
    if (oldIrql < DISPATCH_LEVEL)
        KeRaiseIrql(DISPATCH_LEVEL);

    if (_InterlockedCompareExchange(&sl->Locked, 1, 0) == 0) {
        *outCookie = KiPackSpinlockCookie(flags, oldIrql);
        return 1;
    }

    KeLowerIrql(oldIrql);
    __writeeflags(flags);
    return 0;
}

// Transitional names for the IRQL-aware driver model work that follows.
static __forceinline uintptr_t KeAcquireSpinlockRaiseToDpc(SPINLOCK* sl) {
    return KeAcquireSpinlock(sl);
}

static __forceinline void KeReleaseSpinlockFromDpc(SPINLOCK* sl, uintptr_t cookie) {
    KeReleaseSpinlock(sl, cookie);
}