From 1ee23d33ac70fd93276a0c6fd5adda8975fa382a Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Sat, 23 Dec 2023 11:19:28 +0200
Subject: [PATCH] [NTOS:KE] Fix some KAFFINITYs

---
 ntoskrnl/include/internal/ke.h |  4 ++--
 ntoskrnl/ke/thrdobj.c          | 18 ++++++++++--------
 ntoskrnl/ke/thrdschd.c         |  4 ++--
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/ntoskrnl/include/internal/ke.h b/ntoskrnl/include/internal/ke.h
index 6a477545c44..ba3a5079a95 100644
--- a/ntoskrnl/include/internal/ke.h
+++ b/ntoskrnl/include/internal/ke.h
@@ -137,7 +137,7 @@ extern LIST_ENTRY KiProcessInSwapListHead, KiProcessOutSwapListHead;
 extern LIST_ENTRY KiStackInSwapListHead;
 extern KEVENT KiSwapEvent;
 extern PKPRCB KiProcessorBlock[];
-extern ULONG_PTR KiIdleSummary;
+extern KAFFINITY KiIdleSummary;
 extern PVOID KeUserApcDispatcher;
 extern PVOID KeUserCallbackDispatcher;
 extern PVOID KeUserExceptionDispatcher;
@@ -379,7 +379,7 @@ UCHAR
 NTAPI
 KeFindNextRightSetAffinity(
     IN UCHAR Number,
-    IN ULONG Set
+    IN KAFFINITY Set
 );
 
 VOID
diff --git a/ntoskrnl/ke/thrdobj.c b/ntoskrnl/ke/thrdobj.c
index 1fc58b99c9e..6fca3821411 100644
--- a/ntoskrnl/ke/thrdobj.c
+++ b/ntoskrnl/ke/thrdobj.c
@@ -20,9 +20,10 @@ extern LIST_ENTRY PspReaperListHead;
 UCHAR
 NTAPI
 KeFindNextRightSetAffinity(IN UCHAR Number,
-                           IN ULONG Set)
+                           IN KAFFINITY Set)
 {
-    ULONG Bit, Result;
+    KAFFINITY Bit;
+    ULONG Result;
     ASSERT(Set != 0);
 
     /* Calculate the mask */
@@ -32,7 +33,7 @@ KeFindNextRightSetAffinity(IN UCHAR Number,
     if (!Bit) Bit = Set;
 
     /* Now find the right set and return it */
-    BitScanReverse(&Result, Bit);
+    BitScanReverseAffinity(&Result, Bit);
     return (UCHAR)Result;
 }
 
@@ -491,7 +492,7 @@ KeStartThread(IN OUT PKTHREAD Thread)
 #ifdef CONFIG_SMP
     PKNODE Node;
     PKPRCB NodePrcb;
-    ULONG Set, Mask;
+    KAFFINITY Set, Mask;
 #endif
     UCHAR IdealProcessor = 0;
     PKPROCESS Process = Thread->ApcState.Process;
@@ -526,7 +527,7 @@ KeStartThread(IN OUT PKTHREAD Thread)
 #else
     Set = ~NodePrcb->MultiThreadProcessorSet;
 #endif
-    Mask = (ULONG)(Node->ProcessorMask & Process->Affinity);
+    Mask = Node->ProcessorMask & Process->Affinity;
     Set &= Mask;
     if (Set) Mask = Set;
 
@@ -1122,7 +1123,8 @@ KeSetSystemAffinityThread(IN KAFFINITY Affinity)
 #ifdef CONFIG_SMP
     if (!(Affinity & AFFINITY_MASK(CurrentThread->IdealProcessor)))
     {
-        ULONG AffinitySet, NodeMask;
+        KAFFINITY AffinitySet, NodeMask;
+        ULONG IdealProcessor;
 
         /* It's not! Get the PRCB */
         Prcb = KiProcessorBlock[CurrentThread->IdealProcessor];
@@ -1137,8 +1139,8 @@ KeSetSystemAffinityThread(IN KAFFINITY Affinity)
         }
 
         /* Calculate the ideal CPU from the affinity set */
-        BitScanReverse(&NodeMask, AffinitySet);
-        CurrentThread->IdealProcessor = (UCHAR)NodeMask;
+        BitScanReverseAffinity(&IdealProcessor, AffinitySet);
+        CurrentThread->IdealProcessor = (UCHAR)IdealProcessor;
     }
 #endif
 
diff --git a/ntoskrnl/ke/thrdschd.c b/ntoskrnl/ke/thrdschd.c
index a3235d221a4..40801ff5bee 100644
--- a/ntoskrnl/ke/thrdschd.c
+++ b/ntoskrnl/ke/thrdschd.c
@@ -22,8 +22,8 @@
 
 /* GLOBALS *******************************************************************/
 
-ULONG_PTR KiIdleSummary;
-ULONG_PTR KiIdleSMTSummary;
+KAFFINITY KiIdleSummary;
+KAFFINITY KiIdleSMTSummary;
 
 /* FUNCTIONS *****************************************************************/