From e6bd41ba3ffc643c2921b094b017f382f4460220 Mon Sep 17 00:00:00 2001
From: daan <daan@effp.org>
Date: Tue, 1 Nov 2022 16:26:17 -0700
Subject: [PATCH] update arena_id api

---
 include/mimalloc.h |  9 +++++----
 src/arena.c        | 22 +++++++++++++++++-----
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/include/mimalloc.h b/include/mimalloc.h
index 67ae8374..e9213ce4 100644
--- a/include/mimalloc.h
+++ b/include/mimalloc.h
@@ -275,12 +275,13 @@ mi_decl_export bool mi_manage_os_memory(void* start, size_t size, bool is_commit
 
 // Experimental: heaps associated with specific memory arena's
 typedef int mi_arena_id_t;
-mi_decl_export int  mi_reserve_huge_os_pages_at_ex(size_t pages, int numa_node, size_t timeout_msecs, bool exclusive, mi_arena_id_t* arena_id) mi_attr_noexcept;
-mi_decl_export int  mi_reserve_os_memory_ex(size_t size, bool commit, bool allow_large, bool exclusive, mi_arena_id_t* arena_id) mi_attr_noexcept;
-mi_decl_export bool mi_manage_os_memory_ex(void* start, size_t size, bool is_committed, bool is_large, bool is_zero, int numa_node, bool exclusive, mi_arena_id_t* arena_id) mi_attr_noexcept;
+mi_decl_export void* mi_arena_area(mi_arena_id_t arena_id, size_t* size);
+mi_decl_export int   mi_reserve_huge_os_pages_at_ex(size_t pages, int numa_node, size_t timeout_msecs, bool exclusive, mi_arena_id_t* arena_id) mi_attr_noexcept;
+mi_decl_export int   mi_reserve_os_memory_ex(size_t size, bool commit, bool allow_large, bool exclusive, mi_arena_id_t* arena_id) mi_attr_noexcept;
+mi_decl_export bool  mi_manage_os_memory_ex(void* start, size_t size, bool is_committed, bool is_large, bool is_zero, int numa_node, bool exclusive, mi_arena_id_t* arena_id) mi_attr_noexcept;
 
 #if MI_MALLOC_VERSION >= 200
-mi_decl_nodiscard mi_decl_export mi_heap_t* mi_heap_new_in_arena(mi_arena_id_t arena_id, bool exclusive);
+mi_decl_nodiscard mi_decl_export mi_heap_t* mi_heap_new_in_arena(mi_arena_id_t arena_id);
 #endif
 
 // deprecated
diff --git a/src/arena.c b/src/arena.c
index b7aa47bc..9a309b01 100644
--- a/src/arena.c
+++ b/src/arena.c
@@ -95,8 +95,9 @@ mi_arena_id_t _mi_arena_id_none(void) {
   return 0;
 }
 
-static bool mi_arena_id_suitable(mi_arena_id_t arena_id, bool exclusive, mi_arena_id_t req_arena_id) {
-  return (!exclusive || arena_id == req_arena_id);
+static bool mi_arena_id_is_suitable(mi_arena_id_t arena_id, bool exclusive, mi_arena_id_t req_arena_id) {
+  return ((!exclusive && req_arena_id == _mi_arena_id_none()) || 
+          (arena_id == req_arena_id));
 }
 
 
@@ -122,11 +123,11 @@ static bool mi_arena_memid_indices(size_t arena_memid, size_t* arena_index, mi_b
   return ((arena_memid & 0x80) != 0);
 }
 
-bool _mi_arena_memid_suitable(size_t arena_memid, mi_arena_id_t request_arena_id) {
+bool _mi_arena_memid_is_suitable(size_t arena_memid, mi_arena_id_t request_arena_id) {
   mi_assert_internal(arena_memid != MI_MEMID_OS);
   mi_arena_id_t id = (int)(arena_memid & 0x7F);
   bool exclusive = ((arena_memid & 0x80) != 0);
-  return mi_arena_id_suitable(id, exclusive, request_arena_id);
+  return mi_arena_id_is_suitable(id, exclusive, request_arena_id);
 }
 
 static size_t mi_block_count_of_size(size_t size) {
@@ -157,7 +158,7 @@ static void* mi_arena_alloc_from(mi_arena_t* arena, size_t arena_index, size_t n
 {
   MI_UNUSED(arena_index);
   mi_assert_internal(mi_arena_id_index(arena->id) == arena_index);
-  if (!mi_arena_id_suitable(arena->id, arena->exclusive, req_arena_id)) return NULL;
+  if (!mi_arena_id_is_suitable(arena->id, arena->exclusive, req_arena_id)) return NULL;
 
   mi_bitmap_index_t bitmap_index;
   if (!mi_arena_alloc(arena, needed_bcount, &bitmap_index)) return NULL;
@@ -265,6 +266,17 @@ void* _mi_arena_alloc(size_t size, bool* commit, bool* large, bool* is_pinned, b
   return _mi_arena_alloc_aligned(size, MI_ARENA_BLOCK_SIZE, commit, large, is_pinned, is_zero, arena_id, memid, tld);
 }
 
+
+void* mi_arena_area(mi_arena_id_t arena_id, size_t* size) {
+  if (size != NULL) *size = 0;
+  size_t arena_index = mi_arena_id_index(arena_id);
+  if (arena_index >= MI_MAX_ARENAS) return NULL;
+  mi_arena_t* arena = mi_atomic_load_ptr_relaxed(mi_arena_t, &mi_arenas[arena_index]);
+  if (arena == NULL) return NULL;
+  if (size != NULL) *size = arena->block_count * MI_ARENA_BLOCK_SIZE;
+  return arena->start;
+}
+
 /* -----------------------------------------------------------
   Arena free
 ----------------------------------------------------------- */