handmade.network » Wiki » Tutorial/Making a Heap Allocator

The Plan

Let's make a heap allocator.

We'll make two functions, one to allocate memory, and another to free memory.

We'll divide allocation requests into two categories.

  • Large allocations (>32KiB)
  • Normal allocations

The large allocations will go straight through to the operating system. We'll deal with the normal allocations ourselves. This is because the OS will only deal with allocations of a fairly large granularity.

Normal allocations will be placed into 64KiB regions allocated from the OS, that'll get split up into smaller regions as necessary. When there are no more allocations within a given OS allocation, the memory will be returned to the OS.

We'll keep several linked lists of free regions for different allocation sizes.

To allocate memory, we'll take a free region, split it into two regions, store a reference to the unused region in the correct linked list, and then return the base address of the other region.

To free memory, we'll merge the freed region with its neighbours (if possible - they may be still in use), and then store the merged region in the correct linked list.

Here's a UML diagram:

HeapRegion

First, let's make a structure to store the metadata of an individual region.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
struct HeapRegion {
    union {
        // The number of bytes to the next region in this allocation block (and, therefore, the size of the region).
        /* 0x00 */ uint16_t next;
        /* 0x00 */ uint16_t size;
    };

    // The number of bytes to go backwards to get to the previous region in this allocation block.
    /* 0x02 */ uint16_t previous;

    // The offset of this region within the 64KiB allocation block from the OS.
    /* 0x04 */ uint16_t offset;

    // A logical value to indicate whether this region is in use or not.
    /* 0x06 */ uint16_t used;

    // Valid if the region is not in use.
    /* 0x08 */ HeapRegion *regionListNext;       // The next region in this linked list of unused region.
    /* 0x10 */ HeapRegion **regionListReference; // The previous region in this linked list's pointer to this region.
};

Only the first 8 bytes are needed for a unused region. We'll actually round the metadata block to 16-bytes so that SIMD types will be aligned on the heap.

We'll also define some macros to make traversal of the regions simpler.

1
2
3
4
#define HEAP_REGION_HEADER(region) ((HeapRegion *) ((uint8_t *) region - 0x10))
#define HEAP_REGION_DATA(region) ((uint8_t *) region + 0x10)
#define HEAP_REGION_NEXT(region) ((HeapRegion *) ((uint8_t *) region + region->next))
#define HEAP_REGION_PREVIOUS(region) (region->previous ? ((OSHeapRegion *) ((uint8_t *) region - region->previous)) : nullptr)

Messy pointer arithmetic. Yuck.

Globals

Our heap doesn't need many globals.

1
2
3
4
5
// The 12 linked lists of unused regions.
static HeapRegion *heapRegions[12];

// A mutex to synchronize heap operations.
static OSHandle heapMutex; // Replace this with whatever OS-specific code you need.

Unused Regions Lists

We need to be able to add and remove regions to and from the unused region linked lists.

First, let's also make a function to convert a region size into a index into the unused region array.

1
2
3
4
5
static uintptr_t HeapCalculateIndex(uintptr_t size) {
    int x = __builtin_clz(size);
    uintptr_t msb = sizeof(unsigned int) * 8 - x - 1;
    return msb - 4; // The allocation granularity is 32-bytes.
}

Now we can make the other functions.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
static void HeapAddFreeRegion(HeapRegion *region) {
    // Check the metadata is valid.
    if (region->used || region->size < 32) {
        HeapPanic();
    }

    // Get the index of the list we'll use.
    int index = HeapCalculateIndex(region->size);

    // Store it in the linked list.
    region->regionListNext = heapRegions[index];
    if (region->regionListNext) region->regionListNext->regionListReference = &region->regionListNext;
    heapRegions[index] = region;
    region->regionListReference = heapRegions + index;
}

static void HeapRemoveFreeRegion(HeapRegion *region) {
    // Check the metadata is valid.
    if (!region->regionListReference || region->used) {
        HeapPanic();
    }

    // Add remove the region from the linked list...

    *region->regionListReference = region->regionListNext;

    if (region->regionListNext) {
        region->regionListNext->regionListReference = region->regionListReference;
    }

    region->regionListReference = nullptr;
}

Allocating

Finally it's time to allocate memory.

First, if the caller passes in a size of 0, immediately return nullptr.

1
2
void *HeapAllocate(size_t size, bool zeroMemory) {
    if (!size) return nullptr;

Next, round up to the allocation granularity of 32-bytes, having added the size of the metadata block.

1
2
3
4
    size_t originalSize = size;

    size += 0x10; // Region metadata.
    size = (size + 0x1F) & ~0x1F; // Allocation granularity: 32 bytes.

Now we can tell if this is a large or normal allocation.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
    if (size >= 32768) {
        // This is a very large allocation, so allocate it by itself.
        // The OS will zero this memory for us - how nice!
        HeapRegion *region = (HeapRegion *) OSAllocate(size); // Replace this with OS specific code.
        region->used = 0xABCD; // Our magic signature for allocated regions.
        if (!region) return nullptr; // If the OS allocation failed, so will we.

        // Return the base address.     
        void *address = HEAP_REGION_DATA(region);
        return address;
    }

Best acquire the mutex now.

1
    OSAcquireMutex(heapMutex);

Now let's search for the smallest available heap region to service our allocation.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
    HeapRegion *region = nullptr;

    for (int i = HeapCalculateIndex(size) /*Start at the minimum size for the allocation*/; i < 12; i++) {
        // Are there any regions in this list?
        if (heapRegions[i] == nullptr || heapRegions[i]->size < size) {
            continue;
        }

        // Found a region! Remove it from the list.
        region = heapRegions[i];
        HeapRemoveFreeRegion(region);
        goto foundRegion;
    }

If this fails, we need to allocation another block of memory from the OS.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
    region = (HeapRegion *) OSAllocate(65536);

    // If the allocation failed, so will we.
    if (!region) {
        OSMutexRelease(heapMutex);
        return nullptr; 
    }

    // Set the size of the region.
    region->size = 65536 - 32;

    // Prevent HeapFree trying to merge off the end of the block by making a permanently used region.
    // The metadata will be in the last 32 byte of the allocated block.
    {
        HeapRegion *endRegion = HEAP_REGION_NEXT(region);
        endRegion->used = 0xABCD;
    }

Great! We've got a region.

1
2
3
4
5
    foundRegion:

    if (region->used || region->size < size) {
        HeapPanic();
    }

If the size of the region matches the allocation perfectly, return now.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
    if (region->size == size) {
        // If the size of this region is equal to the size of the region we're trying to allocate,
        // return this region immediately.
        region->used = 0xABCD;
        OSMutexRelease(heapMutex);

        // Zero the memory.
        if (zeroMemory) ZeroMemory(HEAP_REGION_DATA(region), originalSize);

        // Return the base address.
        void *address = HEAP_REGION_DATA(region);
        return address;
    }

Otherwise, we need to split up the region, and put the unused section back into the unused lists.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
    // Split the region into 2 parts.

    HeapRegion *allocatedRegion = region;
    size_t oldSize = allocatedRegion->size;
    allocatedRegion->size = size;
    allocatedRegion->used = 0xABCD;

    HeapRegion *freeRegion = HEAP_REGION_NEXT(allocatedRegion);
    freeRegion->size = oldSize - size;
    freeRegion->previous = size;
    freeRegion->offset = allocatedRegion->offset + size;
    freeRegion->used = false;
    HeapAddFreeRegion(freeRegion);

    HeapRegion *nextRegion = HEAP_REGION_NEXT(freeRegion);
    nextRegion->previous = freeRegion->size;

...and we can now return with the first half.

1
2
3
4
5
6
    OSMutexRelease(heapMutex);
    if (zeroMemory) ZeroMemory(HEAP_REGION_DATA(allocatedRegion), originalSize);

    void *address = HEAP_REGION_DATA(region);
    return address;
}

Freeing

The final step is to be able to free memory. Let's give it a shot.

1
2
3
4
5
6
7
void HeapFree(void *address) {
    if (!address) return;

    HeapRegion *region = HEAP_REGION_HEADER(address);

    // Check this is a valid heap allocation.
    if (region->used != 0xABCD) HeapPanic();

If this was a large allocation, we can free it easily.

1
2
3
4
5
    if (!region->size) {
        // The region was allocated by itself.
        OSFree(region);
        return;
    }

First, let's start by marking the region as free.

1
2
3
    HEAP_ACQUIRE_MUTEX();

    region->used = false;

Then, let's look at the previous and next regions to see if we can merge with them. If they're unused, we'll merge. We remove them from their unused list.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
    // Attempt to merge with the next region.

    HeapRegion *nextRegion = HEAP_REGION_NEXT(region);

    if (nextRegion && !nextRegion->used) {
        HeapRemoveFreeRegion(nextRegion);

        // Merge the regions.
        region->size += nextRegion->size;
        HEAP_REGION_NEXT(nextRegion)->previous = region->size;
    }

    // Attempt to merge with the previous region.

    HeapRegion *previousRegion = HEAP_REGION_PREVIOUS(region);

    if (previousRegion && !previousRegion->used) {
        HeapRemoveFreeRegion(previousRegion);

        // Merge the regions.
        previousRegion->size += region->size;
        HEAP_REGION_NEXT(region)->previous = previousRegion->size;
        region = previousRegion;
    }

Have we finished with this OS allocation block?

1
2
3
4
5
6
7
8
    if (region->size == 65536 - 32) {
        if (region->offset) HeapPanic();

        // The memory block is empty.
        OSFree(region);
        OSMutexRelease(heapMutex);
        return;
    }

Otherwise, let's put the region into a unused region list.

1
2
3
4
    // Put the free region in the region list.
    HeapAddFreeRegion(region);
    MutexRelease(heapMutex);
}

...and that's it.

Future Expansions

  • Resizing allocations.
  • Maintaining small lists of pointers to common allocation sizes for super-speedy allocations.
  • Keeping thread-local heap information to prevent always having to get the mutex.