From 6f1e6fdd06799c0dbf8b06f044960f31bd160ea2 Mon Sep 17 00:00:00 2001
From: SirOlaf <34164198+SirOlaf@users.noreply.github.com>
Date: Sun, 23 Aug 2026 07:36:10 +0200
Subject: [PATCH] Specialize `rawAlloc` for alignment (#26115)

Specialize `rawAlloc` for alignment (cherry-picked from the other PR).
This cuts the frame of the normal unaligned path down enough to regain
the performance lost from loading the cold page in #26110

Also cleans up `MemRegion` a bit, the regressions are either gone or
were measurement errors.
---
 lib/system/alloc.nim | 40 ++++++++++++++++++++++++++++------------
 1 file changed, 28 insertions(+), 12 deletions(-)

diff --git a/lib/system/alloc.nim b/lib/system/alloc.nim
index 2de1d9f7e5..161c3adae4 100644
--- a/lib/system/alloc.nim
+++ b/lib/system/alloc.nim
@@ -155,15 +155,17 @@ type
 
   MemRegion = object
     when usesRegionHandles:
+      # Keeping the handle here does change the layout, but until proven otherwise
+      # this layout is more readable and shouldn't regress performance.
       regionHandle: ptr RegionHandle
     when not defined(gcDestructors):
       minLargeObj, maxLargeObj: int
     freeSmallChunks: array[0..max(1, SmallChunkSize div MemAlign-1), PSmallChunk]
       # List of available chunks per size class. Only one is expected to be active per class.
-    when defined(gcDestructors):
+    when defined(gcDestructors) and not usesRegionHandles:
       sharedFreeLists: SharedFreeLists
-        # Used directly without threads. Threaded builds use RegionHandle but
-        # retain this 2 KiB spacer: removing it regresses 2-4 KiB allocations.
+        # Remote-free buckets live on the MemRegion when there is no
+        # RegionHandle. Threaded memory managers with handles keep them on the handle instead.
     flBitmap: uint32
     slBitmap: array[RealFli, uint32]
     matrix: array[RealFli, array[MaxSli, PBigChunk]]
@@ -963,13 +965,19 @@ proc bigChunkAlignOffset(alignment: int): int {.inline.} =
   else:
     result = align(sizeof(BigChunk) + sizeof(FreeCell), alignment) - sizeof(BigChunk) - sizeof(FreeCell)
 
-proc rawAlloc(a: var MemRegion, requestedSize: int, alignment: int = 0): pointer =
+template rawAllocAux(aligned: static bool) {.dirty.} =
   when defined(nimTypeNames):
     inc(a.allocCounter)
   sysAssert(allocInv(a), "rawAlloc: begin")
   sysAssert(roundup(65, 8) == 72, "rawAlloc: roundup broken")
-  var size = roundup(requestedSize, max(MemAlign, alignment))
-  let alignOff = smallChunkAlignOffset(alignment)
+  when aligned:
+    var size = roundup(requestedSize, max(MemAlign, alignment))
+    let alignOff = smallChunkAlignOffset(alignment)
+  else:
+    # Common `alloc` path: no custom alignment. Keep this a separate
+    # instantiation so clang does not emit `smallChunkAlignOffset(0)`.
+    var size = (requestedSize + (MemAlign - 1)) and not (MemAlign - 1)
+    const alignOff = 0
   sysAssert(size >= sizeof(FreeCell), "rawAlloc: requested size too small")
   sysAssert(size >= requestedSize, "insufficient allocated size!")
   #c_fprintf(stdout, "alloc; size: %ld; %ld\n", requestedSize, size)
@@ -986,11 +994,13 @@ proc rawAlloc(a: var MemRegion, requestedSize: int, alignment: int = 0): pointer
             if atomicLoadN(sharedHead, ATOMIC_RELAXED) != nil:
               tc.freeList = atomicExchangeN(sharedHead, nil, ATOMIC_ACQUIRE)
           else:
-            tc.freeList = a.sharedFreeLists[s]
-            a.sharedFreeLists[s] = nil
-          # If `tc.freeList` isn't nil, `tc` gains capacity. Calculate how
-          # much it gained and how many foreign cells are included.
-          compensateCounters(a, tc, size)
+            let sharedHead = addr a.sharedFreeLists[s]
+            tc.freeList = sharedHead[]
+            sharedHead[] = nil
+          # Empty peeks are the common local case; skip the walk and the
+          # `free += 0` / `occ -= 0` stores clang would otherwise keep.
+          if tc.freeList != nil:
+            compensateCounters(a, tc, size)
 
     # allocate a small block: for small chunks, we use only its next pointer
     let s = size div MemAlign
@@ -1071,7 +1081,7 @@ proc rawAlloc(a: var MemRegion, requestedSize: int, alignment: int = 0): pointer
     # For big chunks with custom alignment, allocate extra space.
     # Since chunks are page-aligned, the needed padding is a compile-time
     # deterministic value rather than a worst-case estimate.
-    let alignPad = bigChunkAlignOffset(alignment)
+    let alignPad = when aligned: bigChunkAlignOffset(alignment) else: 0
     size = requestedSize + bigChunkOverhead() + alignPad
     # allocate a large block
     var c = if size >= HugeChunkSize: getHugeChunk(a, size)
@@ -1096,6 +1106,12 @@ proc rawAlloc(a: var MemRegion, requestedSize: int, alignment: int = 0): pointer
   when defined(heaptrack):
     heaptrack_malloc(result, requestedSize)
 
+proc rawAlloc(a: var MemRegion, requestedSize: int): pointer =
+  rawAllocAux(false)
+
+proc rawAlloc(a: var MemRegion, requestedSize: int, alignment: int): pointer =
+  rawAllocAux(true)
+
 proc rawAlloc0(a: var MemRegion, requestedSize: int): pointer =
   result = rawAlloc(a, requestedSize)
   zeroMem(result, requestedSize)