From 119f2eff6c61eb4f0c315773944b914af47d48ef Mon Sep 17 00:00:00 2001
From: daanx <daan@effp.org>
Date: Wed, 5 Mar 2025 09:51:40 -0800
Subject: [PATCH] use int for numa node count

---
 include/mimalloc/internal.h |  2 +-
 src/arena.c                 | 10 +++++-----
 src/os.c                    | 20 +++++++++++---------
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/include/mimalloc/internal.h b/include/mimalloc/internal.h
index b4515831..e8b1c919 100644
--- a/include/mimalloc/internal.h
+++ b/include/mimalloc/internal.h
@@ -160,7 +160,7 @@ bool          _mi_os_secure_guard_page_reset_at(void* addr);
 bool          _mi_os_secure_guard_page_reset_before(void* addr);
 
 int           _mi_os_numa_node(void);
-size_t        _mi_os_numa_node_count(void);
+int           _mi_os_numa_node_count(void);
 
 void*         _mi_os_alloc_aligned(size_t size, size_t alignment, bool commit, bool allow_large, mi_memid_t* memid);
 void*         _mi_os_alloc_aligned_at_offset(size_t size, size_t alignment, size_t align_offset, bool commit, bool allow_large, mi_memid_t* memid);
diff --git a/src/arena.c b/src/arena.c
index 70e1802b..2aa1f8fe 100644
--- a/src/arena.c
+++ b/src/arena.c
@@ -1524,17 +1524,17 @@ int mi_reserve_huge_os_pages_interleave(size_t pages, size_t numa_nodes, size_t
   if (pages == 0) return 0;
 
   // pages per numa node
-  size_t numa_count = (numa_nodes > 0 ? numa_nodes : _mi_os_numa_node_count());
-  if (numa_count <= 0) numa_count = 1;
+  int numa_count = (numa_nodes > 0 && numa_nodes <= INT_MAX ? (int)numa_nodes : _mi_os_numa_node_count());
+  if (numa_count <= 0) { numa_count = 1; }
   const size_t pages_per = pages / numa_count;
   const size_t pages_mod = pages % numa_count;
   const size_t timeout_per = (timeout_msecs==0 ? 0 : (timeout_msecs / numa_count) + 50);
 
   // reserve evenly among numa nodes
-  for (size_t numa_node = 0; numa_node < numa_count && pages > 0; numa_node++) {
+  for (int numa_node = 0; numa_node < numa_count && pages > 0; numa_node++) {
     size_t node_pages = pages_per;  // can be 0
-    if (numa_node < pages_mod) node_pages++;
-    int err = mi_reserve_huge_os_pages_at(node_pages, (int)numa_node, timeout_per);
+    if ((size_t)numa_node < pages_mod) { node_pages++; }
+    int err = mi_reserve_huge_os_pages_at(node_pages, numa_node, timeout_per);
     if (err) return err;
     if (pages < node_pages) {
       pages = 0;
diff --git a/src/os.c b/src/os.c
index 69decb71..2ee4c897 100644
--- a/src/os.c
+++ b/src/os.c
@@ -694,18 +694,19 @@ static void mi_os_free_huge_os_pages(void* p, size_t size) {
 Support NUMA aware allocation
 -----------------------------------------------------------------------------*/
 
-static _Atomic(size_t)  _mi_numa_node_count; // = 0   // cache the node count
+static _Atomic(int)  _mi_numa_node_count; // = 0   // cache the node count
 
-size_t _mi_os_numa_node_count(void) {
+int _mi_os_numa_node_count(void) {
   size_t count = mi_atomic_load_acquire(&_mi_numa_node_count);
   if mi_unlikely(count <= 0) {
     long ncount = mi_option_get(mi_option_use_numa_nodes); // given explicitly?
-    if (ncount > 0) {
-      count = (size_t)ncount;
+    if (ncount > 0 && ncount < INT_MAX) {
+      count = (int)ncount;
     }
     else {
-      count = _mi_prim_numa_node_count(); // or detect dynamically
-      if (count == 0) { count = 1; }
+      const size_t n = _mi_prim_numa_node_count(); // or detect dynamically
+      if (n == 0 || n > INT_MAX) { count = 1; }
+                            else { count = (int)n; }
     }
     mi_atomic_store_release(&_mi_numa_node_count, count); // save it
     _mi_verbose_message("using %zd numa regions\n", count);
@@ -715,12 +716,13 @@ size_t _mi_os_numa_node_count(void) {
 
 
 static int mi_os_numa_node_get(void) {
-  size_t numa_count = _mi_os_numa_node_count();
+  int numa_count = _mi_os_numa_node_count();
   if (numa_count<=1) return 0; // optimize on single numa node systems: always node 0
   // never more than the node count and >= 0
-  size_t numa_node = _mi_prim_numa_node();
+  const size_t n = _mi_prim_numa_node();
+  int numa_node = (n < INT_MAX ? (int)n : 0);
   if (numa_node >= numa_count) { numa_node = numa_node % numa_count; }
-  return (int)numa_node;
+  return numa_node;
 }
 
 int _mi_os_numa_node(void) {