From 3b057c08b4a50e6315418676cbef9ec299004874 Mon Sep 17 00:00:00 2001
From: Daan <daanl@outlook.com>
Date: Sat, 11 May 2024 10:07:15 -0700
Subject: [PATCH] refactor in-place expansion

---
 src/os.c                   | 92 +++++++++++++++++++++++---------------
 src/prim/emscripten/prim.c |  5 ++-
 src/prim/windows/prim.c    |  1 +
 3 files changed, 60 insertions(+), 38 deletions(-)

diff --git a/src/os.c b/src/os.c
index 4ff1a815..8acc3338 100644
--- a/src/os.c
+++ b/src/os.c
@@ -403,6 +403,11 @@ void* _mi_os_alloc_aligned_at_offset(size_t size, size_t alignment, size_t offse
 ----------------------------------------------------------- */
 
 void* _mi_os_alloc_expandable(size_t size, size_t alignment, size_t future_reserve, mi_memid_t* memid, mi_stats_t* stats) {
+  if (!mi_os_mem_config.has_virtual_reserve) {
+    // don't allocate expandable if the OS does not support virtual reservation
+    return _mi_os_alloc_aligned(size, alignment, false, false, memid, stats);
+  }
+
   size = mi_os_get_alloc_size(size);
   if (future_reserve < 2*size) { future_reserve = 2*size; }
   void* p = _mi_os_alloc_aligned(future_reserve, alignment, false, false, memid, stats);
@@ -415,9 +420,56 @@ void* _mi_os_alloc_expandable(size_t size, size_t alignment, size_t future_reser
   return p;
 }
 
+static bool mi_os_try_expand_inplace( void* p, size_t size, size_t newsize, mi_memid_t* memid, mi_stats_t* stats) {
+  // try to expand the existing virtual range "in-place"    
+  if (p != NULL && size > 0 && newsize > size && 
+      !mi_os_mem_config.must_free_whole && !memid->is_pinned && memid->mem.os.prim_info == NULL) 
+  {
+    void* expand = (uint8_t*)p + size;
+    size_t extra = newsize - size;
+    mi_assert_internal(extra > 0 && (extra %  _mi_os_page_size()) == 0);
+    bool os_is_large = false;
+    bool os_is_zero = false;
+    void* newp = mi_os_prim_alloc_at(expand, extra, 1, false /* commit? */, false /* allow large */, &os_is_large, &os_is_zero, stats);
+    if (newp == expand) {
+      // success! we expanded the virtual address space in-place
+      if (_mi_os_commit(newp, extra, &os_is_zero, stats)) {
+        _mi_verbose_message("expanded in place (address: %p, from %zu bytes to %zu bytes\n", p, size, newsize);
+        memid->is_pinned = os_is_large;
+        memid->mem.os.size += extra;
+        return true;
+      }
+    }
+    
+    // failed, free reserved space 
+    if (newp != NULL) {
+      mi_os_prim_free(newp, extra, false, stats);
+    }
+    return false;
+  }
+  else if (p != NULL && newsize > 0 && newsize < size && 
+           !mi_os_mem_config.must_free_whole && !memid->is_pinned && memid->mem.os.prim_info == NULL)
+  {
+     // we can shrink in-place by free-ing the upper part
+     void* shrink = (uint8_t*)p + newsize;
+     size_t extra = size - newsize;
+     mi_assert_internal(extra > 0 && (extra %  _mi_os_page_size()) == 0);
+     mi_os_prim_free(shrink, extra, true, stats);
+     _mi_verbose_message("shrunk OS memory in place (address: %p, from %zu bytes to %zu bytes\n", p, size, newsize);
+     memid->mem.os.size -= extra;     
+     return true;     
+  }
+  else {
+    return false;
+  }
+}
+
 bool  _mi_os_expand(void* p, size_t size, size_t newsize, mi_memid_t* memid, mi_stats_t* stats) {
   if (p == NULL) return false;
-  if (memid->memkind != MI_MEM_OS_EXPAND) return false;
+  if (memid->memkind != MI_MEM_OS_EXPAND) {
+    return mi_os_try_expand_inplace(p,size,newsize,memid,stats);
+  }
+  // expandable memory
   if (newsize > size) {
     if (memid->mem.os.size < newsize) {
       return false;
@@ -456,42 +508,10 @@ static void* mi_os_remap_copy(void* p, size_t size, size_t newsize, size_t align
   newsize = mi_os_get_alloc_size(newsize);
 
   // first try to expand the existing virtual range "in-place"    
-  if (p != NULL && size > 0 && newsize > size && 
-      !mi_os_mem_config.must_free_whole && !memid->is_pinned && memid->mem.os.prim_info == NULL) 
-  {
-    void* expand = (uint8_t*)p + size;
-    size_t extra = newsize - size;
-    mi_assert_internal(extra > 0 && (extra %  _mi_os_page_size()) == 0);
-    bool os_is_large = false;
-    bool os_is_zero = false;
-    void* newp = mi_os_prim_alloc_at(expand, extra, 1, false /* commit? */, false, &os_is_large, &os_is_zero, stats);
-    if (newp == expand) {
-      // success! we expanded the virtual address space in-place
-      if (_mi_os_commit(newp, extra, &os_is_zero, stats)) {
-        _mi_verbose_message("expanded in place (address: %p, from %zu bytes to %zu bytes\n", p, size, newsize);
-        memid->is_pinned = os_is_large;
-        memid->mem.os.size += newsize;
-        return p;
-      }
-    }
-    
-    // failed, free reserved space and fall back to a copy
-    if (newp != NULL) {
-      mi_os_prim_free(newp, extra, false, stats);
-    }
+  if (mi_os_try_expand_inplace(p,size,newsize,memid,stats)) {
+    return p;
   }
-  else if (p != NULL && newsize > 0 && newsize < size && 
-           !mi_os_mem_config.must_free_whole && !memid->is_pinned && memid->mem.os.prim_info == NULL)
-  {
-     // we can shrink in-place by free-ing the upper part
-     void* shrink = (uint8_t*)p + newsize;
-     size_t extra = size - newsize;
-     mi_assert_internal(extra > 0 && (extra %  _mi_os_page_size()) == 0);
-     mi_os_prim_free(shrink, extra, true, stats);
-     _mi_verbose_message("shrunk OS memory in place (address: %p, from %zu bytes to %zu bytes\n", p, size, newsize);
-     return p;     
-  }
-
+  
   // otherwise: copy into a fresh area
   void* newp = _mi_os_alloc_aligned(newsize, alignment, true /* commit */, false /* allow_large */, &newmemid, stats);
   if (newp == NULL) return NULL;
diff --git a/src/prim/emscripten/prim.c b/src/prim/emscripten/prim.c
index 1f60a1bb..fcf577d9 100644
--- a/src/prim/emscripten/prim.c
+++ b/src/prim/emscripten/prim.c
@@ -53,6 +53,7 @@ void _mi_prim_mem_init( mi_os_mem_config_t* config) {
   config->has_overcommit = false;
   config->must_free_whole = true;
   config->has_virtual_reserve = false;
+  config->has_remap = false;
 }
 
 extern void emmalloc_free(void*);
@@ -71,8 +72,8 @@ int _mi_prim_free(void* addr, size_t size) {
 extern void* emmalloc_memalign(size_t alignment, size_t size);
 
 // Note: the `try_alignment` is just a hint and the returned pointer is not guaranteed to be aligned.
-int _mi_prim_alloc(size_t size, size_t try_alignment, bool commit, bool allow_large, bool* is_large, bool* is_zero, void** addr) {
-  MI_UNUSED(try_alignment); MI_UNUSED(allow_large); MI_UNUSED(commit);
+int _mi_prim_alloc(void* hint, size_t size, size_t try_alignment, bool commit, bool allow_large, bool* is_large, bool* is_zero, void** addr) {
+  MI_UNUSED(hint); MI_UNUSED(allow_large); MI_UNUSED(commit);
   *is_large = false;
   // TODO: Track the highest address ever seen; first uses of it are zeroes.
   //       That assumes no one else uses sbrk but us (they could go up,
diff --git a/src/prim/windows/prim.c b/src/prim/windows/prim.c
index 49861318..695ad36c 100644
--- a/src/prim/windows/prim.c
+++ b/src/prim/windows/prim.c
@@ -113,6 +113,7 @@ void _mi_prim_mem_init( mi_os_mem_config_t* config )
   config->has_overcommit = false;
   config->must_free_whole = true;
   config->has_virtual_reserve = true;
+  config->has_remap = false;
   // get the page size
   SYSTEM_INFO si;
   GetSystemInfo(&si);