monolithium/crt/src/malloc.c

/*
 * heap.c
 *
 * Copyright (C) 2017 Aleksandar Andrejevic <theflash@sdf.lonestar.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdlib.h>
#include <string.h>
#include <monolithium.h>

#define ALLOCATED (1 << 0)

typedef struct
{
    uint32_t magic;
    uint32_t flags;
    size_t size;
} heap_header_t;

__crt_heap_t *__crt_default_heap = NULL;

static void __crt_heap_coalesce(__crt_heap_t *heap)
{
    heap_header_t *ptr = (heap_header_t*)heap->base;
    heap_header_t *previous = NULL;

    while ((uintptr_t)ptr >= (uintptr_t)heap->base && (uintptr_t)ptr < (uintptr_t)heap->base + heap->size)
    {
        heap_header_t *next = (heap_header_t*)((uintptr_t)ptr + sizeof(heap_header_t) + ptr->size);

        if (previous && !(previous->flags & ALLOCATED) && !(ptr->flags & ALLOCATED))
        {
            if (((uintptr_t)ptr - (uintptr_t)heap->base) == heap->next_offset)
            {
                heap->next_offset = (uintptr_t)previous - (uintptr_t)heap->base;
            }

            previous->size += ptr->size + sizeof(heap_header_t);
        }
        else
        {
            previous = ptr;
        }

        ptr = next;
    }
}

void *__crt_heap_realloc(__crt_heap_t *heap, void *ptr, size_t alignment, size_t size)
{
    if (!alignment || (alignment & (alignment - 1))) return NULL;

    heap->lock_mutex_proc(heap->mutex);

    if (!(heap->flags & __CRT_HEAP_FLAG_READY))
    {
        heap_header_t *header = (heap_header_t*)heap->base;
        header->magic = heap->magic;
        header->flags = 0;
        header->size = heap->size - sizeof(heap_header_t);

        heap->flags |= __CRT_HEAP_FLAG_READY;
    }

    if (!size)
    {
        if (ptr)
        {
            heap_header_t *header = (heap_header_t*)((uintptr_t)ptr - sizeof(heap_header_t));

            if (header->flags & ALLOCATED)
            {
                header->flags &= ~ALLOCATED;
                __crt_heap_coalesce(heap);
            }
            else
            {
                heap->problem(__CRT_HEAP_DOUBLE_FREE);
            }
        }

        heap->unlock_mutex_proc(heap->mutex);
        return NULL;
    }

    heap_header_t *source_block = NULL;

    if (ptr)
    {
        heap_header_t *header = (heap_header_t*)((uintptr_t)ptr - sizeof(heap_header_t));

        if (!(header->flags & ALLOCATED) || header->magic != heap->magic)
        {
            heap->problem(__CRT_HEAP_BAD_POINTER);
            heap->unlock_mutex_proc(heap->mutex);
            return NULL;
        }

        if (size > header->size)
        {
             heap_header_t *next = (heap_header_t*)((uintptr_t)ptr + header->size);

             if (!(next->flags & ALLOCATED) && (header->size + next->size + sizeof(heap_header_t)) >= size)
             {
                 if (((uintptr_t)next - (uintptr_t)heap->base) == heap->next_offset)
                 {
                     heap->next_offset = (uintptr_t)header - (uintptr_t)heap->base;
                 }

                 header->size += next->size + sizeof(heap_header_t);
                 if (heap->flags & __CRT_HEAP_FLAG_ZEROFILL) memset(next, 0, size - header->size - size);
             }
             else
             {
                 source_block = header;
             }
        }

        if (size < (header->size - sizeof(heap_header_t)))
        {
            heap_header_t *new_block = (heap_header_t*)((uintptr_t)ptr + size);
            new_block->magic = heap->magic;
            new_block->flags = 0;
            new_block->size = header->size - size - sizeof(heap_header_t);
            header->size = size;
            __crt_heap_coalesce(heap);
        }

        if (!source_block)
        {
            heap->unlock_mutex_proc(heap->mutex);
            return ptr;
        }
    }

    heap_header_t *hole = NULL;

    if (heap->flags & __CRT_HEAP_FLAG_BESTFIT)
    {
        heap_header_t *current = (heap_header_t*)heap->base;
        size_t minimum_size;

        while ((uintptr_t)current < (uintptr_t)heap->base + heap->size)
        {
            if (current->magic != heap->magic)
            {
                heap->problem(__CRT_HEAP_CORRUPTED);
                heap->unlock_mutex_proc(heap->mutex);
                return NULL;
            }

            if (!(current->flags & ALLOCATED))
            {
                uintptr_t current_start = (uintptr_t)current + sizeof(heap_header_t);
                uintptr_t aligned_start = (current_start + alignment - 1) & ~(alignment - 1);
                size_t padding = aligned_start - current_start;

                if (current->size > padding)
                {
                    size_t adjusted_size = current->size - padding;

                    if (adjusted_size >= size && (!hole || adjusted_size < minimum_size))
                    {
                        hole = current;
                        minimum_size = adjusted_size;
                    }
                }
            }

            heap_header_t *next = (heap_header_t*)((uintptr_t)current + sizeof(heap_header_t) + current->size);

            if ((uintptr_t)next <= (uintptr_t)current)
            {
                heap->problem(__CRT_HEAP_CORRUPTED);
                heap->unlock_mutex_proc(heap->mutex);
                return NULL;
            }

            current = next;
        }
    }
    else
    {
        uintptr_t offset = heap->next_offset;
        int cycles = 0;

        do
        {
            heap_header_t *current = (heap_header_t*)(heap->base + offset);

            if (current->magic != heap->magic)
            {
                heap->problem(__CRT_HEAP_CORRUPTED);
                heap->unlock_mutex_proc(heap->mutex);
                return NULL;
            }

            if (!(current->flags & ALLOCATED))
            {
                uintptr_t current_start = (uintptr_t)current + sizeof(heap_header_t);
                uintptr_t aligned_start = (current_start + alignment - 1) & ~(alignment - 1);
                size_t padding = aligned_start - current_start;

                if (current->size >= padding && (current->size - padding) >= size)
                {
                    hole = current;
                    break;
                }
            }

            offset += sizeof(heap_header_t) + current->size;
            if (offset > heap->size)
            {
                offset = 0;
                if (++cycles > 1)
                {
                    heap->problem(__CRT_HEAP_CORRUPTED);
                    heap->unlock_mutex_proc(heap->mutex);
                    return NULL;
                }
            }
        }
        while (offset != heap->next_offset);

        heap->next_offset = offset;
    }

    if (!hole)
    {
        heap->problem(__CRT_HEAP_OUT_OF_MEMORY);
        heap->unlock_mutex_proc(heap->mutex);
        return NULL;
    }

    int coalesce = 0;
    uintptr_t start_address = (uintptr_t)hole + sizeof(heap_header_t);
    uintptr_t aligned_start = (start_address + alignment - 1) & ~(alignment - 1);
    size_t padding = aligned_start - start_address;

    if (padding > sizeof(heap_header_t))
    {
        heap_header_t *new_block = (heap_header_t*)(aligned_start - sizeof(heap_header_t));
        new_block->magic = heap->magic;
        new_block->flags = 0;
        new_block->size = hole->size - padding;
        hole->size -= padding + sizeof(heap_header_t);

        hole = new_block;
        coalesce = 1;
    }
    else if (padding)
    {
        heap_header_t *previous = (heap_header_t*)heap->base;

        while ((uintptr_t)previous < (uintptr_t)heap->base + heap->size)
        {
            heap_header_t *next = (heap_header_t*)((uintptr_t)previous + sizeof(heap_header_t) + previous->size);
            if (next == hole) break;

            if ((uintptr_t)next <= (uintptr_t)previous)
            {
                heap->problem(__CRT_HEAP_CORRUPTED);
                heap->unlock_mutex_proc(heap->mutex);
                return NULL;
            }

            previous = next;
        }

        if ((uintptr_t)previous < (uintptr_t)heap->base
            || (uintptr_t)previous >= (uintptr_t)heap->base + heap->size
            || previous->magic != heap->magic)
        {
            heap->problem(__CRT_HEAP_CORRUPTED);
            heap->unlock_mutex_proc(heap->mutex);
            return NULL;
        }

        previous->size += padding;

        heap_header_t *new_block = (heap_header_t*)(aligned_start - sizeof(heap_header_t));
        memmove(new_block, hole, sizeof(heap_header_t));
        if (heap->next_offset == (uintptr_t)hole - (uintptr_t)heap->base) heap->next_offset += padding;

        hole = new_block;
        hole->size -= padding;
    }

    hole->flags |= ALLOCATED;

    if (hole->size > sizeof(heap_header_t) && size < (hole->size - sizeof(heap_header_t)))
    {
        heap_header_t *new_block = (heap_header_t*)((uintptr_t)hole + size + sizeof(heap_header_t));
        new_block->magic = heap->magic;
        new_block->flags = 0;
        new_block->size = hole->size - size - sizeof(heap_header_t);
        hole->size = size;
        coalesce = 1;
    }

    void *destination = (void*)((uintptr_t)hole + sizeof(heap_header_t));
    if (heap->flags & __CRT_HEAP_FLAG_ZEROFILL) memset(destination, 0, size);

    if (source_block)
    {
        void *source = (void*)((uintptr_t)source_block + sizeof(heap_header_t));
        memcpy(destination, source, source_block->size);
        source_block->flags &= ~ALLOCATED;
        coalesce = 1;
    }

    if (coalesce) __crt_heap_coalesce(heap);
    heap->unlock_mutex_proc(heap->mutex);
    return destination;
}

void *aligned_alloc(size_t alignment, size_t size)
{
    return __crt_heap_realloc(__crt_default_heap, NULL, alignment, size);
}

void *realloc(void *ptr, size_t size)
{
    return __crt_heap_realloc(__crt_default_heap, ptr, sizeof(long), size);
}

void *malloc(size_t size)
{
    return __crt_heap_realloc(__crt_default_heap, NULL, sizeof(long), size);
}

void free(void *ptr)
{
    __crt_heap_realloc(__crt_default_heap, ptr, sizeof(long), 0);
}

void *calloc(size_t nmemb, size_t size)
{
    void *ptr = malloc(nmemb * size);
    if (!ptr) return NULL;

    memset(ptr, 0, nmemb * size);
    return ptr;
}

static void __crt_heap_lock(void *mutex)
{
    syscall_wait_mutex((handle_t)mutex, NO_TIMEOUT);
}

static void __crt_heap_unlock(void *mutex)
{
    syscall_release_mutex((handle_t)mutex);
}

static void __crt_heap_problem(int problem)
{
    // TODO
}

int __crt_initialize_heap(void)
{
    static __crt_heap_t heap;

    heap.magic             = 0x70616548;
    heap.base              = NULL;
    heap.size              = 0x10000000;
    heap.flags             = 0;
    heap.next_offset       = 0;
    heap.lock_mutex_proc   = __crt_heap_lock;
    heap.unlock_mutex_proc = __crt_heap_unlock;
    heap.problem           = __crt_heap_problem;

    sysret_t status = syscall_create_mutex(NULL, TRUE, (handle_t*)&heap.mutex);
    if (status != ERR_SUCCESS) return -1;

    status = syscall_alloc_memory(INVALID_HANDLE, &heap.base, heap.size, MEMORY_BLOCK_ACCESSIBLE | MEMORY_BLOCK_WRITABLE);
    if (status != ERR_SUCCESS) return -1;

    __crt_default_heap = &heap;
    return 0;
}