// Public nt-lite DDK synchronization declarations.

#pragma once

#include <stdint.h>
#include <intrin.h>
#include <ntlite/ddk/irql.h>

// Legacy spinlock storage used by current driver-visible objects.
typedef struct SPINLOCK {
    volatile long Locked;
} SPINLOCK;

#define SPINLOCK_INIT { 0 }

// Requests a synchronization wait without a timeout.
#define KE_WAIT_INFINITE UINT64_MAX

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

// Packs the interrupt flags and prior IRQL into a spinlock cookie.
static __forceinline uintptr_t KiPackSpinlockCookie(uintptr_t flags, KIRQL oldIrql) {
    return (flags & KI_SPINLOCK_FLAGS_MASK) |
           (((uintptr_t)oldIrql) << KI_SPINLOCK_IRQL_SHIFT);
}

// Extracts the saved interrupt flags from a spinlock cookie.
static __forceinline uintptr_t KiSpinlockCookieFlags(uintptr_t cookie) {
    return cookie & KI_SPINLOCK_FLAGS_MASK;
}

// Extracts the saved IRQL from a spinlock cookie.
static __forceinline KIRQL KiSpinlockCookieIrql(uintptr_t cookie) {
    return (KIRQL)(cookie >> KI_SPINLOCK_IRQL_SHIFT);
}

// Acquires a spinlock and raises to at least DISPATCH_LEVEL.
static __forceinline uintptr_t KeAcquireSpinlock(SPINLOCK* sl) {
    uintptr_t flags = KeDisableInterrupts();

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

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

    return KiPackSpinlockCookie(flags, oldIrql);
}

// Releases a spinlock and restores the caller's saved IRQL and interrupt 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);
    KeRestoreInterrupts(flags);
}

// Attempts to acquire a spinlock without blocking.
static __forceinline int KeTryAcquireSpinlock(SPINLOCK* sl, uintptr_t* outCookie) {
    uintptr_t flags = KeDisableInterrupts();

    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);
    KeRestoreInterrupts(flags);
    return 0;
}

// Acquires a spinlock for DPC-level code.
static __forceinline uintptr_t KeAcquireSpinlockRaiseToDpc(SPINLOCK* sl) {
    return KeAcquireSpinlock(sl);
}

// Releases a spinlock acquired for DPC-level code.
static __forceinline void KeReleaseSpinlockFromDpc(SPINLOCK* sl, uintptr_t cookie) {
    KeReleaseSpinlock(sl, cookie);
}