1 files changed, 111 insertions, 17 deletions

diff --git a/CPP/Windows/System.cpp b/CPP/Windows/System.cpp
index 5fa87f3..4745785 100644
--- a/CPP/Windows/System.cpp
+++ b/CPP/Windows/System.cpp
@@ -25,6 +25,69 @@ namespace NSystem {
 
 #ifdef _WIN32
 
+/*
+note: returned value in 32-bit version can be limited by value 32.
+      while 64-bit version returns full value.
+GetMaximumProcessorCount(groupNumber) can return higher value than
+GetActiveProcessorCount(groupNumber) in some cases, because CPUs can be added.
+*/
+// typedef DWORD (WINAPI *Func_GetMaximumProcessorCount)(WORD GroupNumber);
+typedef DWORD (WINAPI *Func_GetActiveProcessorCount)(WORD GroupNumber);
+typedef WORD (WINAPI *Func_GetActiveProcessorGroupCount)(VOID);
+/*
+#if 0 && defined(ALL_PROCESSOR_GROUPS)
+#define MY_ALL_PROCESSOR_GROUPS   ALL_PROCESSOR_GROUPS
+#else
+#define MY_ALL_PROCESSOR_GROUPS   0xffff
+#endif
+*/
+
+Z7_DIAGNOSTIC_IGNORE_CAST_FUNCTION
+
+bool CCpuGroups::Load()
+{
+  NumThreadsTotal = 0;
+  GroupSizes.Clear();
+  const HMODULE hmodule = ::GetModuleHandleA("kernel32.dll");
+  // Is_Win11_Groups = GetProcAddress(hmodule, "SetThreadSelectedCpuSetMasks") != NULL;
+  const
+      Func_GetActiveProcessorGroupCount
+        fn_GetActiveProcessorGroupCount = Z7_GET_PROC_ADDRESS(
+      Func_GetActiveProcessorGroupCount, hmodule,
+          "GetActiveProcessorGroupCount");
+  const
+      Func_GetActiveProcessorCount
+        fn_GetActiveProcessorCount = Z7_GET_PROC_ADDRESS(
+      Func_GetActiveProcessorCount, hmodule,
+          "GetActiveProcessorCount");
+  if (!fn_GetActiveProcessorGroupCount ||
+      !fn_GetActiveProcessorCount)
+    return false;
+
+  const unsigned numGroups = fn_GetActiveProcessorGroupCount();
+  if (numGroups == 0)
+    return false;
+  UInt32 sum = 0;
+  for (unsigned i = 0; i < numGroups; i++)
+  {
+    const UInt32 num = fn_GetActiveProcessorCount((WORD)i);
+    /*
+    if (num == 0)
+    {
+      // it means error
+      // but is it possible that some group is empty by some reason?
+      // GroupSizes.Clear();
+      // return false;
+    }
+    */
+    sum += num;
+    GroupSizes.Add(num);
+  }
+  NumThreadsTotal = sum;
+  // NumThreadsTotal = fn_GetActiveProcessorCount(MY_ALL_PROCESSOR_GROUPS);
+  return true;
+}
+
 UInt32 CountAffinity(DWORD_PTR mask)
 {
   UInt32 num = 0;
@@ -38,31 +101,62 @@ UInt32 CountAffinity(DWORD_PTR mask)
 
 BOOL CProcessAffinity::Get()
 {
-  #ifndef UNDER_CE
-  return GetProcessAffinityMask(GetCurrentProcess(), &processAffinityMask, &systemAffinityMask);
-  #else
-  return FALSE;
-  #endif
+  IsGroupMode = false;
+  Groups.Load();
+  // SetThreadAffinityMask(GetCurrentThread(), 1);
+  // SetProcessAffinityMask(GetCurrentProcess(), 1);
+  BOOL res = GetProcessAffinityMask(GetCurrentProcess(),
+      &processAffinityMask, &systemAffinityMask);
+  /* DOCs: On a system with more than 64 processors, if the threads
+     of the calling process  are in a single processor group, the
+     function sets the variables pointed to by lpProcessAffinityMask
+     and lpSystemAffinityMask to the process affinity mask and the
+     processor mask of active logical processors for that group.
+     If the calling process contains threads in multiple groups,
+     the function returns zero for both affinity masks
+
+     note: tested in Win10: GetProcessAffinityMask() doesn't return 0
+           in (processAffinityMask) and (systemAffinityMask) masks.
+     We need to test it in Win11: how to get mask==0 from GetProcessAffinityMask()?
+  */
+  if (!res)
+  {
+    processAffinityMask = 0;
+    systemAffinityMask = 0;
+  }
+  if (Groups.GroupSizes.Size() > 1 && Groups.NumThreadsTotal)
+    if (// !res ||
+        processAffinityMask == 0 || // to support case described in DOCs and for (!res) case
+        processAffinityMask == systemAffinityMask) // for default nonchanged affinity
+    {
+      // we set IsGroupMode only if processAffinity is default (not changed).
+      res = TRUE;
+      IsGroupMode = true;
+    }
+  return res;
 }
 
 
+UInt32 CProcessAffinity::Load_and_GetNumberOfThreads()
+{
+  if (Get())
+  {
+    const UInt32 numProcessors = GetNumProcessThreads();
+    if (numProcessors)
+      return numProcessors;
+  }
+  SYSTEM_INFO systemInfo;
+  GetSystemInfo(&systemInfo);
+  // the number of logical processors in the current group
+  return systemInfo.dwNumberOfProcessors;
+}
+
 UInt32 GetNumberOfProcessors()
 {
   // We need to know how many threads we can use.
   // By default the process is assigned to one group.
-  // So we get the number of logical processors (threads)
-  // assigned to current process in the current group.
-  // Group size can be smaller than total number logical processors, for exammple, 2x36
-
   CProcessAffinity pa;
-
-  if (pa.Get() && pa.processAffinityMask != 0)
-    return pa.GetNumProcessThreads();
-
-  SYSTEM_INFO systemInfo;
-  GetSystemInfo(&systemInfo);
-  // the number of logical processors in the current group
-  return (UInt32)systemInfo.dwNumberOfProcessors;
+  return pa.Load_and_GetNumberOfThreads();
 }
 
 #else