From 4be79a0cd088e86998e0c17edc3a08163af9ec1c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Matos?= <joao@tritao.eu>
Date: Wed, 11 Sep 2024 22:04:31 +0100
Subject: [PATCH] Add `os.getnumcpus` function to get number of logical CPU
 cores. (#2241)

---
 src/host/os_getnumcpus.c      | 91 +++++++++++++++++++++++++++++++++++
 src/host/premake.c            |  1 +
 src/host/premake.h            |  1 +
 tests/base/test_os.lua        | 10 ++++
 website/docs/os.getnumcpus.md | 17 +++++++
 website/sidebars.js           |  1 +
 6 files changed, 121 insertions(+)
 create mode 100644 src/host/os_getnumcpus.c
 create mode 100644 website/docs/os.getnumcpus.md

diff --git a/src/host/os_getnumcpus.c b/src/host/os_getnumcpus.c
new file mode 100644
index 0000000000..4fc3f42f88
--- /dev/null
+++ b/src/host/os_getnumcpus.c
@@ -0,0 +1,91 @@
+/**
+ * \file   os_getnumcpus.c
+ * \brief  Retrieve the logical number of CPUs of the host system.
+ * \author Copyright (c) 2002-2024 Jason Perkins and the Premake project
+ */
+
+#if __linux__
+#define _GNU_SOURCE
+#endif
+
+#include "premake.h"
+
+#if PLATFORM_LINUX
+#include <sched.h>
+#elif PLATFORM_SOLARIS | PLATFORM_AIX | PLATFORM_MACOSX | PLATFORM_BSD
+#include <sys/sysctl.h>
+#elif PLATFORM_WINDOWS
+#define WIN32_LEAN_AND_MEAN
+#include <windows.h>
+#endif
+
+int do_getnumcpus()
+{
+#if PLATFORM_WINDOWS
+	SYSTEM_INFO sysinfo;
+	GetSystemInfo(&sysinfo);
+	return sysinfo.dwNumberOfProcessors;
+#elif PLATFORM_LINUX
+	cpu_set_t set;
+	int count, i;
+
+	if (sched_getaffinity(0, sizeof(cpu_set_t), &set) != -1)
+	{
+		count = 0;
+		for (i = 0; i < CPU_SETSIZE; i++)
+		{
+			if (CPU_ISSET(i, &set))
+			{
+				count++;
+			}
+		}
+
+		return count;
+	}
+	else
+	{
+		return 0;
+	}
+#elif PLATFORM_SOLARIS | PLATFORM_AIX | PLATFORM_MACOSX
+	return sysconf(_SC_NPROCESSORS_ONLN);
+#elif PLATFORM_BSD
+	int mib[4];
+	int numCPU;
+	size_t len = sizeof(numCPU);
+
+	/* set the mib for hw.ncpu */
+	mib[0] = CTL_HW;
+	mib[1] = HW_AVAILCPU;  // alternatively, try HW_NCPU;
+
+	/* get the number of CPUs from the system */
+	sysctl(mib, 2, &numCPU, &len, NULL, 0);
+
+	if (numCPU < 1)
+	{
+		mib[1] = HW_NCPU;
+		sysctl(mib, 2, &numCPU, &len, NULL, 0);
+		if (numCPU < 1)
+			return 0;
+	}
+
+	return numCPU;
+#else
+	#warning do_getnumcpus is not implemented for your platform yet
+	return 0;
+#endif
+}
+
+int os_getnumcpus(lua_State* L)
+{
+	int result = do_getnumcpus();
+	if (result > 0)
+	{
+		lua_pushinteger(L, result);
+		return 1;
+	}
+	else
+	{
+		return 0;
+	}
+}
+
diff --git a/src/host/premake.c b/src/host/premake.c
index 49d192f283..e304846223 100644
--- a/src/host/premake.c
+++ b/src/host/premake.c
@@ -66,6 +66,7 @@ static const luaL_Reg os_functions[] = {
 	{ "_is64bit",               os_is64bit              },
 	{ "isdir",                  os_isdir                },
 	{ "getcwd",                 os_getcwd               },
+	{ "getnumcpus",             os_getnumcpus           },
 	{ "getpass",                os_getpass              },
 	{ "getWindowsRegistry",     os_getWindowsRegistry   },
 	{ "listWindowsRegistry",    os_listWindowsRegistry  },
diff --git a/src/host/premake.h b/src/host/premake.h
index 00c29f6b42..8f49827379 100644
--- a/src/host/premake.h
+++ b/src/host/premake.h
@@ -121,6 +121,7 @@ int os_chmod(lua_State* L);
 int os_comparefiles(lua_State* L);
 int os_copyfile(lua_State* L);
 int os_getcwd(lua_State* L);
+int os_getnumcpus(lua_State* L);
 int os_getpass(lua_State* L);
 int os_getWindowsRegistry(lua_State* L);
 int os_listWindowsRegistry(lua_State* L);
diff --git a/tests/base/test_os.lua b/tests/base/test_os.lua
index ae2ea75c3f..dcb027a780 100644
--- a/tests/base/test_os.lua
+++ b/tests/base/test_os.lua
@@ -488,3 +488,13 @@
 		test.isequal(true, ok)
 		test.isnil(err)
 	end
+
+
+--
+-- os.getnumcpus() tests.
+--
+
+	function suite.numcpus()
+		local numcpus = os.getnumcpus()
+		test.istrue(numcpus > 0)
+	end
diff --git a/website/docs/os.getnumcpus.md b/website/docs/os.getnumcpus.md
new file mode 100644
index 0000000000..91978ecd86
--- /dev/null
+++ b/website/docs/os.getnumcpus.md
@@ -0,0 +1,17 @@
+Gets the number of logical CPU cores.
+
+```lua
+os.getnumcpus()
+```
+
+### Parameters ###
+
+None.
+
+### Return Value ###
+
+The number of logical CPU cores of the running system.
+
+### Availability ###
+
+Premake 5.0.0 beta 3 or later.
diff --git a/website/sidebars.js b/website/sidebars.js
index 04856f0059..7f0e169c65 100644
--- a/website/sidebars.js
+++ b/website/sidebars.js
@@ -375,6 +375,7 @@ module.exports = {
 						'os.findlib',
 						'os.get',
 						'os.getcwd',
+						'os.getnumcpus',
 						'os.getpass',
 						'os.getversion',
 						'os.host',