// Public nt-lite DDK synchronization declarations.

#pragma once

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

// Legacy spinlock storage used by current driver-visible objects.
typedef struct SPINLOCK {
    volatile int32_t 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 (NtlAtomicCompareExchange32(&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);

    NtlAtomicExchange32(&sl->Locked, 0);

    // Restore the caller's hardware interrupt state before dropping to an
    // IRQL that can dispatch APCs. PASSIVE_LEVEL with IF still clear is not a
    // valid APC execution context.
    KeRestoreInterrupts(flags);
    KeLowerIrql(oldIrql);
}

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

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

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

    // Keep the failed-acquire unwind consistent with normal release: restore
    // IF before a PASSIVE transition can dispatch pending APCs.
    KeRestoreInterrupts(flags);
    KeLowerIrql(oldIrql);
    return false;
}

// 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);
}