From 41f984f6c85ee1281305a2f2fe0c705095073c9c Mon Sep 17 00:00:00 2001
From: Joao Matos <joao@tritao.eu>
Date: Sun, 1 Sep 2024 18:01:31 +0100
Subject: [PATCH] Add `os.getnumcpus` function to get number of logical CPU
 cores.

---
 Bootstrap.mak                 |  2 +-
 premake5.lua                  |  4 ++
 src/host/os_getnumcpus.c      | 84 +++++++++++++++++++++++++++++++++++
 src/host/premake.c            |  1 +
 src/host/premake.h            |  1 +
 tests/base/test_os.lua        | 10 +++++
 website/docs/os.getnumcpus.md | 17 +++++++
 7 files changed, 118 insertions(+), 1 deletion(-)
 create mode 100644 src/host/os_getnumcpus.c
 create mode 100644 website/docs/os.getnumcpus.md

diff --git a/Bootstrap.mak b/Bootstrap.mak
index 5934916ad5..7a4eb16e34 100644
--- a/Bootstrap.mak
+++ b/Bootstrap.mak
@@ -107,7 +107,7 @@ linux-clean: nix-clean
 
 linux: linux-clean
 	mkdir -p build/bootstrap
-	$(CC) -o build/bootstrap/premake_bootstrap -DPREMAKE_NO_BUILTIN_SCRIPTS -DLUA_USE_POSIX -DLUA_USE_DLOPEN -I"$(LUA_DIR)" -I"$(LUASHIM_DIR)" $(SRC) -lm -ldl -lrt -luuid
+	$(CC) -o build/bootstrap/premake_bootstrap -DPREMAKE_NO_BUILTIN_SCRIPTS -D_GNU_SOURCE -DLUA_USE_POSIX -DLUA_USE_DLOPEN -I"$(LUA_DIR)" -I"$(LUASHIM_DIR)" $(SRC) -lm -ldl -lrt -luuid
 	./build/bootstrap/premake_bootstrap embed
 	./build/bootstrap/premake_bootstrap --to=build/bootstrap gmake2
 	$(MAKE) -C build/bootstrap -j`getconf _NPROCESSORS_ONLN` config=$(CONFIG)
diff --git a/premake5.lua b/premake5.lua
index 7818c1d9ad..678cfffda9 100644
--- a/premake5.lua
+++ b/premake5.lua
@@ -248,6 +248,10 @@
 		filter "system:linux or hurd"
 			links       { "dl", "rt" }
 
+		filter "system:linux"
+			defines     { "_GNU_SOURCE" }
+
+
 		filter { "system:not windows", "system:not macosx" }
 			if not _OPTIONS["no-curl"] then
 				links   { "mbedtls-lib" }
diff --git a/src/host/os_getnumcpus.c b/src/host/os_getnumcpus.c
new file mode 100644
index 0000000000..ad55738a72
--- /dev/null
+++ b/src/host/os_getnumcpus.c
@@ -0,0 +1,84 @@
+/**
+ * \file   os_getcwd.c
+ * \brief  Retrieve the current working directory.
+ * \author Copyright (c) 2002-2008 Jason Perkins and the Premake project
+ */
+
+#include "premake.h"
+
+#if PLATFORM_LINUX
+#include <sched.h>
+#elif PLATFORM_SOLARIS | PLATFORM_AIX | PLATFORM_MACOSX
+#include <sys/sysctl.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)
+			numCPU = 1;
+	}
+
+	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..02301c0a16
--- /dev/null
+++ b/website/docs/os.getnumcpus.md
@@ -0,0 +1,17 @@
+Gets the number of logical CPU cores.
+
+```lua
+cwd = os.getnumcpus()
+```
+
+### Parameters ###
+
+None.
+
+### Return Value ###
+
+The number of logical CPU cores of the running system.
+
+### Availability ###
+
+Premake 5.0 or later.