From a200291ae576e1e766be8526c0d058bd9766d480 Mon Sep 17 00:00:00 2001
From: daan <daan@effp.org>
Date: Sun, 6 Nov 2022 08:26:17 -0800
Subject: [PATCH] further progress on removing aligned limit

---
 src/alloc-aligned.c | 14 +++++++-------
 src/page.c          |  4 ++--
 test/test-api.c     | 21 +++++++++++++++++++--
 3 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/src/alloc-aligned.c b/src/alloc-aligned.c
index f4ddbf99..63acd58c 100644
--- a/src/alloc-aligned.c
+++ b/src/alloc-aligned.c
@@ -29,22 +29,21 @@ static mi_decl_noinline void* mi_heap_malloc_zero_aligned_at_fallback(mi_heap_t*
     mi_assert_internal(p == NULL || ((uintptr_t)p % alignment) == 0);
     return p;
   }
-  
+
   void* p;
   size_t oversize;
   if mi_unlikely(alignment > MI_ALIGNMENT_MAX) {
     // use OS allocation for very large alignment and allocate inside a huge page (dedicated segment with 1 page)
     if mi_unlikely(offset != 0) {
       // todo: cannot support offset alignment for very large alignments yet
-      #if MI_DEBUG > 0
+#if MI_DEBUG > 0
       _mi_error_message(EOVERFLOW, "aligned allocation with a very large alignment cannot be used with an alignment offset (size %zu, alignment %zu, offset %zu)\n", size, alignment, offset);
-      #endif
+#endif
       return NULL;
     }
-    oversize = size + MI_SEGMENT_SIZE - 1;
-    p = _mi_heap_malloc_zero_ex(heap, oversize, zero, alignment);
+    oversize = (size <= MI_SMALL_SIZE_MAX ? MI_SMALL_SIZE_MAX + 1 /* ensure we use generic malloc path */ : size);
+    p = _mi_heap_malloc_zero_ex(heap, oversize, zero, alignment); // the page block size should be large enough to align in the single huge page block
     if (p == NULL) return NULL;
-    //mi_assert_internal(_mi_is_aligned(p, alignment));    
   }
   else {
     // otherwise over-allocate
@@ -57,7 +56,8 @@ static mi_decl_noinline void* mi_heap_malloc_zero_aligned_at_fallback(mi_heap_t*
   uintptr_t adjust = alignment - (((uintptr_t)p + offset) & align_mask);
   mi_assert_internal(adjust <= alignment);
   void* aligned_p = (adjust == alignment ? p : (void*)((uintptr_t)p + adjust));
-  if (aligned_p != p) mi_page_set_has_aligned(_mi_ptr_page(p), true);
+  if (aligned_p != p) { mi_page_set_has_aligned(_mi_ptr_page(p), true);  }
+  mi_assert_internal(mi_page_usable_block_size(_mi_ptr_page(p)) >= adjust + size);
   mi_assert_internal(p == _mi_page_ptr_unalign(_mi_ptr_segment(aligned_p), _mi_ptr_page(aligned_p), aligned_p));
   mi_assert_internal(((uintptr_t)aligned_p + offset) % alignment == 0);
   
diff --git a/src/page.c b/src/page.c
index 5e2ec826..6b54eb2c 100644
--- a/src/page.c
+++ b/src/page.c
@@ -262,7 +262,7 @@ static mi_page_t* mi_page_fresh_alloc(mi_heap_t* heap, mi_page_queue_t* pq, size
   }
   // a fresh page was found, initialize it
   mi_assert_internal(pq==NULL || _mi_page_segment(page)->page_kind != MI_PAGE_HUGE);
-  mi_page_init(heap, page, block_size, heap->tld);
+  mi_page_init(heap, page, (pq == NULL ? MI_HUGE_BLOCK_SIZE : block_size), heap->tld);
   mi_heap_stat_increase(heap, pages, 1);
   if (pq!=NULL) mi_page_queue_push(heap, pq, page); // huge pages use pq==NULL
   mi_assert_expensive(_mi_page_is_valid(page));
@@ -643,7 +643,7 @@ static void mi_page_init(mi_heap_t* heap, mi_page_t* page, size_t block_size, mi
   mi_track_mem_noaccess(page_start,page_size);
   page->xblock_size = (block_size < MI_HUGE_BLOCK_SIZE ? (uint32_t)block_size : MI_HUGE_BLOCK_SIZE);
   mi_assert_internal(page_size / block_size < (1L<<16));
-  page->reserved = (uint16_t)(page_size / block_size);
+  page->reserved = (block_size < MI_HUGE_BLOCK_SIZE ? (uint16_t)(page_size / block_size) : 1);
   #ifdef MI_ENCODE_FREELIST
   page->keys[0] = _mi_heap_random_next(heap);
   page->keys[1] = _mi_heap_random_next(heap);
diff --git a/test/test-api.c b/test/test-api.c
index 650056e2..312b3f1b 100644
--- a/test/test-api.c
+++ b/test/test-api.c
@@ -163,10 +163,27 @@ int main(void) {
   CHECK_BODY("malloc-aligned7") {
     void* p = mi_malloc_aligned(1024,MI_ALIGNMENT_MAX);
     mi_free(p);
+    result = ((uintptr_t)p % MI_ALIGNMENT_MAX) == 0;
   };
   CHECK_BODY("malloc-aligned8") {
-    void* p = mi_malloc_aligned(1024,2*MI_ALIGNMENT_MAX); 
-    mi_free(p);
+    bool ok = true;
+    for (int i = 0; i < 5 && ok; i++) {
+      int n = (1 << i);
+      void* p = mi_malloc_aligned(1024, n * MI_ALIGNMENT_MAX);
+      ok = ((uintptr_t)p % (n*MI_ALIGNMENT_MAX)) == 0;
+      mi_free(p);
+    }
+    result = ok;
+  };
+  CHECK_BODY("malloc-aligned9") {
+    bool ok = true;
+    for (int i = 0; i < 5 && ok; i++) {
+      int n = (1 << i);
+      void* p = mi_malloc_aligned( 2*n*MI_ALIGNMENT_MAX, n*MI_ALIGNMENT_MAX);
+      ok = ((uintptr_t)p % (n*MI_ALIGNMENT_MAX)) == 0;
+      mi_free(p);
+    }
+    result = ok;
   };
   CHECK_BODY("malloc-aligned-at1") {
     void* p = mi_malloc_aligned_at(48,32,0); result = (p != NULL && ((uintptr_t)(p) + 0) % 32 == 0); mi_free(p);