[NDK] Add Affinity helper inline functions

tkreuzer · Dec 3, 2023 · 9615891 · 9615891
1 parent 9ea2222
commit 9615891
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 1 deletion.
diff --git a/ntoskrnl/include/internal/ke.h b/ntoskrnl/include/internal/ke.h
@@ -155,7 +155,6 @@ extern VOID __cdecl KiInterruptTemplate(VOID);
 
 /* MACROS *************************************************************************/
 
-#define AFFINITY_MASK(ProcessorIndex) ((KAFFINITY)1 << (ProcessorIndex))
 #define PRIORITY_MASK(Priority) (1UL << (Priority))
 
 /* Tells us if the Timer or Event is a Syncronization or Notification Object */

diff --git a/sdk/include/ndk/kefuncs.h b/sdk/include/ndk/kefuncs.h
@@ -33,6 +33,54 @@ extern "C"
 
 #ifndef NTOS_MODE_USER
 
+//
+// Affinity helper
+//
+FORCEINLINE KAFFINITY AFFINITY_MASK(ULONG Index)
+{
+    ASSERT(Index < sizeof(KAFFINITY) * 8);
+    return (KAFFINITY)1 << Index;
+}
+#ifdef _WIN64
+FORCEINLINE BOOLEAN BitScanForwardAffinity(PULONG Index, KAFFINITY Mask)
+{
+    return BitScanForward64(Index, Mask);
+}
+FORCEINLINE BOOLEAN BitScanReverseAffinity(PULONG Index, KAFFINITY Mask)
+{
+    return BitScanReverse64(Index, Mask);
+}
+FORCEINLINE BOOLEAN InterlockedBitTestAndSetAffinity(volatile KAFFINITY *Affinity, ULONG Index)
+{
+    ASSERT(Index < sizeof(KAFFINITY) * 8);
+    return InterlockedBitTestAndSet64((PLONG64)Affinity, Index);
+}
+FORCEINLINE BOOLEAN InterlockedBitTestAndResetAffinity(volatile KAFFINITY *Affinity, ULONG Index)
+{
+    ASSERT(Index < sizeof(KAFFINITY) * 8);
+    return InterlockedBitTestAndReset64((PLONG64)Affinity, Index);
+}
+#else
+FORCEINLINE BOOLEAN BitScanForwardAffinity(PULONG Index, KAFFINITY Mask)
+{
+    return BitScanForward(Index, Mask);
+}
+FORCEINLINE BOOLEAN BitScanReverseAffinity(PULONG Index, KAFFINITY Mask)
+{
+    return BitScanReverse(Index, Mask);
+}
+FORCEINLINE BOOLEAN InterlockedBitTestAndSetAffinity(volatile KAFFINITY *Affinity, ULONG Index)
+{
+    ASSERT(Index < sizeof(KAFFINITY) * 8);
+    return InterlockedBitTestAndSet((PLONG)Affinity, Index);
+}
+FORCEINLINE BOOLEAN InterlockedBitTestAndResetAffinity(volatile KAFFINITY *Affinity, ULONG Index)
+{
+    ASSERT(Index < sizeof(KAFFINITY) * 8);
+    return InterlockedBitTestAndReset((PLONG)Affinity, Index);
+}
+#endif // _WIN64
+
 //
 // APC Functions
 //