Avoid over-allocation for naturally aligned requests

src/alloc-aligned.c

@@ -22,7 +22,7 @@ static bool mi_malloc_is_naturally_aligned( size_t size, size_t alignment ) {
   if (alignment > size) return false;
   const size_t bsize = mi_good_size(size);
   const bool ok = (bsize <= MI_PAGE_MAX_START_BLOCK_ALIGN2 && _mi_is_power_of_two(bsize)) ||             // power-of-two under N
-                  (alignment==MI_PAGE_OSPAGE_BLOCK_ALIGN2 && (bsize % MI_PAGE_OSPAGE_BLOCK_ALIGN2)==0);  // or multiple of N
+                  (alignment<=MI_PAGE_OSPAGE_BLOCK_ALIGN2 && (bsize % MI_PAGE_OSPAGE_BLOCK_ALIGN2)==0);  // or multiple of N
   if (ok) { mi_assert_internal((bsize & (alignment-1)) == 0); } // since both power of 2 and alignment <= size
   return ok;
 }
@@ -182,13 +182,19 @@ static mi_decl_cold mi_decl_noinline void* mi_error_bad_alignment(size_t size, s
 }
 
 // Primitive aligned allocation
-static inline void* mi_theap_malloc_zero_aligned_at(mi_theap_t* const theap, const size_t size, const size_t alignment, const size_t offset, const bool zero, mi_page_t** ppage) mi_attr_noexcept
+static inline void* mi_theap_malloc_zero_aligned_at(mi_theap_t* const theap, size_t size, const size_t alignment, const size_t offset, const bool zero, mi_page_t** ppage) mi_attr_noexcept
 {
   // note: we don't require `size > offset`, we just guarantee that the address at offset is aligned regardless of the allocated size.
   if mi_unlikely(!mi_alignment_is_valid(alignment)) { // require power-of-two and multiple of void* (see <https://en.cppreference.com/w/c/memory/aligned_alloc#Notes>)
     return mi_error_bad_alignment(size, alignment, offset);
   }
 
+  // Use a naturally aligned size class instead of over-allocating and adjusting the pointer.
+  // Keep the requested size when padding is enabled for byte-precise overflow checks.
+  if (MI_PADDING_SIZE == 0 && offset == 0 && size < alignment && alignment <= MI_PAGE_MAX_START_BLOCK_ALIGN2) {
+    size = alignment;
+  }
+
   // try first if there happens to be a small block available with just the right alignment
   // since most small power-of-2 blocks (under MI_PAGE_MAX_BLOCK_START_ALIGN2) are already
   // naturally aligned this can be often the case.

test/test-api.c

@@ -70,6 +70,44 @@ static bool mem_is_zero(const void* p, size_t size) {
   return mem_has_vals((const uint8_t*)p,size,0);
 }
 
+// Keep several blocks live so alignment is checked beyond the first block on a page.
+static bool test_natural_alignment(size_t size, size_t alignment, size_t expected_usable_size) {
+  void* blocks[16];
+  const size_t block_count = sizeof(blocks)/sizeof(blocks[0]);
+  bool ok = true;
+
+  for (int round = 0; round < 2; round++) {
+    const bool zero = (round != 0);
+    for (size_t i = 0; i < block_count; i++) {
+      blocks[i] = (zero ? mi_zalloc_aligned(size, alignment) : mi_malloc_aligned(size, alignment));
+      if (blocks[i] == NULL) {
+        ok = false;
+        continue;
+      }
+
+      const size_t usable_size = mi_usable_size(blocks[i]);
+      if ((uintptr_t)blocks[i] % alignment != 0 || usable_size < size) { ok = false; }
+
+      // Padding and sampled guard pages can change the reported usable size.
+      #if !MI_PADDING && !MI_GUARDED
+      if (usable_size != expected_usable_size) { ok = false; }
+      #else
+      (void)expected_usable_size;
+      #endif
+      if (zero && !mem_is_zero(blocks[i], size)) { ok = false; }
+
+      // Dirty the blocks to exercise zeroing when memory is reused.
+      memset(blocks[i], 0xAB, size);
+    }
+
+    for (size_t i = 0; i < block_count; i++) {
+      mi_free(blocks[i]);
+    }
+  }
+
+  return ok;
+}
+
 // ---------------------------------------------------------------------------
 // Main testing
 // ---------------------------------------------------------------------------
@@ -190,6 +228,34 @@ int main(void) {
     }
     result = ok;
   };
+  CHECK_BODY("malloc-aligned-natural") {
+    const struct {
+      size_t size;
+      size_t alignment;
+      size_t expected_usable_size;
+    } cases[] = {
+      {    0,   16,    16},
+      {    8,   16,    16},
+      {    1,   32,    32},
+      {   64,  256,   256},
+      {    1, 4096,  4096},
+      {  300, 4096,  4096},
+      { 4000, 4096,  4096},
+      { 4096, 4096,  4096},
+      { 8192,   32,  8192},
+      {65536,  256, 65536},
+      {65536, 4096, 65536}
+    };
+    bool ok = true;
+
+    for (size_t i = 0; i < sizeof(cases)/sizeof(cases[0]); i++) {
+      if (!test_natural_alignment(cases[i].size, cases[i].alignment, cases[i].expected_usable_size)) {
+        ok = false;
+      }
+    }
+
+    result = ok;
+  };
   CHECK_BODY("malloc-aligned5") {
     void* p = mi_malloc_aligned(4097,4096);
     size_t usable = mi_usable_size(p);