/*
 * PROJECT: Alcyone System Kernel
 * LICENSE: BSD Clause 3
 * PURPOSE: Slim Futex 
 * NT KERNEL: 5.11.9360
 * COPYRIGHT:  2023-2029 Dibymartanda Samanta <>
 */ 

#include <ntoskrnl.h>
#define NTDEBUG

extern "C"

typedef struct _KE_SLIM_FUTEX {
    volatile LONG Value;
} KE_SLIM_FUTEX, *PKE_SLIM_FUTEX;

/* IS_HELD_BIT: Indicates if the futex is currently held.
   HAS_PARKED_BIT : Indicates if there are any parked (waiting) threads. */
constexpr ULONG IS_HELD_BIT = 1;
constexpr ULONG HAS_PARKED_BIT = 2; 

NTSTATUS
KeAcquireSlimFutexSlow(
    _Inout_ PKE_SLIM_FUTEX Futex
)
{
    NTSTATUS Status;
    LARGE_INTEGER Timeout;
    Timeout.QuadPart = -50000; // 5ms

    for (;;) {
        LONG CurrentValue = Futex->Value;
        if (!(CurrentValue & IS_HELD_BIT)) {
            if (InterlockedCompareExchange(&Futex->Value, CurrentValue | IS_HELD_BIT, CurrentValue) == CurrentValue) {
                return STATUS_SUCCESS;
            }
            continue;
        }

        if (!(CurrentValue & HAS_PARKED_BIT)) {
            if (InterlockedCompareExchange(&Futex->Value, CurrentValue | HAS_PARKED_BIT, CurrentValue) != CurrentValue) {
                continue;
            }
        }

        Status = KeWaitForSingleObject(Futex, Executive, KernelMode, FALSE, &Timeout);
        if (!NT_SUCCESS(Status) && Status != STATUS_TIMEOUT) {
            return Status;
        }
    }
}

NTSTATUS
KeReleaseSlimFutexSlow(
    _Inout_ PKE_SLIM_FUTEX Futex
)
{
    LONG OldValue;
    do {
        OldValue = Futex->Value;
        LONG NewValue = OldValue & ~IS_HELD_BIT;
        if (OldValue & HAS_PARKED_BIT)
            NewValue &= ~HAS_PARKED_BIT;
        if (InterlockedCompareExchange(&Futex->Value, NewValue, OldValue) == OldValue) {
            if (OldValue & HAS_PARKED_BIT) {
                KeSetEvent((PRKEVENT)Futex, IO_NO_INCREMENT, FALSE);
            }
            return STATUS_SUCCESS;
        }
    } while (TRUE);
}

/*Exported Function*/ 

/* Initializes PKE_SLIM_FUTEX  Pointer*/ 
NTSTATUS
KeInitializeSlimFutex(
    _Out_ PKE_SLIM_FUTEX Futex
)
{
    Futex->Value = 0;
    return STATUS_SUCCESS;
}

/* Attempts to acquire the slim futex, return STATUS_SUCCESS  if succeeds */  
NTSTATUS
KeAcquireSlimFutex(
    _Inout_ PKE_SLIM_FUTEX Futex
)
{
    if (InterlockedCompareExchange(&Futex->Value, IS_HELD_BIT, 0) == 0) {
        // Futex acquired!
        return STATUS_SUCCESS;
    }

    return KeAcquireSlimFutexSlow(Futex);
}

/*TRUE if the futex is successfully acquired, FALSE otherwise*/

BOOLEAN
KeTryAcquireSlimFutex(
    _Inout_ PKE_SLIM_FUTEX Futex
)
{
    for (;;) {
        LONG CurrentValue = Futex->Value;
        if (CurrentValue & IS_HELD_BIT)
            return FALSE;
        if (InterlockedCompareExchange(&Futex->Value, CurrentValue | IS_HELD_BIT, CurrentValue) == CurrentValue)
            return TRUE;
    }
}

/* Releases the slim futex*/
NTSTATUS
KeReleaseSlimFutex(
    _Inout_ PKE_SLIM_FUTEX Futex
)
{
    if (InterlockedCompareExchange(&Futex->Value, 0, IS_HELD_BIT) == IS_HELD_BIT) {
        // Futex released and nobody was waiting!
        return STATUS_SUCCESS;
    }

    return KeReleaseSlimFutexSlow(Futex);
}

/* Return TRUE if provided Slim Futex is held , else false */
BOOLEAN
KeIsSlimFutexHeld(
    _In_ PKE_SLIM_FUTEX Futex
)
{
    return (Futex->Value & IS_HELD_BIT) != 0;
}
/* Return true  if Futex Locked, else return False */
BOOLEAN
KeIsSlimFutexLocked(
    _In_ PKE_SLIM_FUTEX Futex
)
{
    return KeIsSlimFutexHeld(Futex);
}