malloc

main.c

#define NDEBUG

#include <stdio.h>
#include <assert.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>
#include "solve.h"
#include "utils.h"

const int TOT_SIZE = 100 * 1024 * 1024;

const int SPAN_FACTOR = 3;
const int MAX_SPAN_EXP = 25; // 3 * 32M = 96M
const int MIN_SPAN_EXP = 3; // 3 * 8 = 24 (TODO: Small block cache)

typedef struct SpanHeader {
    uint8_t span_exp, is_free;
    struct SpanHeader *prev, *next;
} SpanHeader;

static int top_exp;
static uint8_t *base, *top, *span_st, *append_st, *append_top;
static SpanHeader *link_st;

static inline int expidx(int span_exp) {
    return span_exp - MIN_SPAN_EXP;
}

static inline int incr_top(int incr) {
    if (top + incr >= base + TOT_SIZE)
        return -1;
    if ((uint8_t*)mem_sbrk(incr) + 1 == NULL)
        return -1;
    top += incr;
    return 0;
}

static inline void insert_span_after(SpanHeader *x, SpanHeader *l) {
    SpanHeader *r = l->next;
    l->next = x, x->next = r;
    r->prev = x, x->prev = l;
}

static inline void insert_span(SpanHeader *x, int span_exp) {
    assert(((uint8_t*)x - span_st) % (SPAN_FACTOR << span_exp) == 0);
    x->span_exp = span_exp;
    x->is_free = true;
    insert_span_after(x, link_st + expidx(span_exp));
}

static inline void remove_span(SpanHeader *x) {
    SpanHeader *l = x->prev, *r = x->next;
    l->next = r, r->prev = l;
}

static inline void *append_only_malloc(size_t size) {
    size += (ALIGNMENT - (size % ALIGNMENT)) % ALIGNMENT;
    void *ret = append_top;
    append_top += size;
    if (append_top > top)
        if (incr_top(append_top - top) == -1)
            return NULL;
    return ret;
}

int my_init() {
    append_st = NULL;

    // Reset base
    if ((top = (uint8_t*)mem_sbrk(0)) + 1 == NULL)
        return -1;
    base = top;
    if (incr_top((ALIGNMENT - (size_t)top % ALIGNMENT) % ALIGNMENT) == -1)
        return -1;

    // Set the link list starting point
    link_st = (SpanHeader*)top;
    if (incr_top(sizeof(SpanHeader) * (MAX_SPAN_EXP - MIN_SPAN_EXP + 1)) == -1)
        return -1;
    for (int i = MIN_SPAN_EXP; i <= MAX_SPAN_EXP; i++)
        link_st[expidx(i)].prev = link_st[expidx(i)].next = link_st + expidx(i);

    // Allocate a minimal span
    span_st = top;
    if (incr_top(SPAN_FACTOR << MIN_SPAN_EXP) == -1)
        return -1;
    insert_span((SpanHeader*)(top - (SPAN_FACTOR << MIN_SPAN_EXP)), MIN_SPAN_EXP);
    top_exp = MIN_SPAN_EXP;

#ifndef NDEBUG
    puts("!!! DEBUG MODE !!!");
    printf("span_st = 0x%lx\n", (size_t)span_st);
#endif
    return 0;
}

void *my_malloc(size_t size) {
    // Add span_exp overhead
    size += ALIGNMENT;

    // Find the span just >= size
    int span_exp = MIN_SPAN_EXP;
    while (span_exp <= MAX_SPAN_EXP && size > (SPAN_FACTOR << span_exp))
        span_exp++;
    if (span_exp > MAX_SPAN_EXP)
        return NULL;

    // Increase top
    append_top = top; // Prepare for append mode
    if ((uint8_t*)link_st[expidx(top_exp)].prev + (SPAN_FACTOR << top_exp) == top)
        append_top -= SPAN_FACTOR << top_exp;
    int split_exp = span_exp;
    while (split_exp < top_exp && link_st[expidx(split_exp)].prev == link_st + expidx(split_exp))
        split_exp++;
    for (; top_exp <= split_exp; top_exp++) {
        if (top_exp == split_exp && link_st[expidx(top_exp)].prev != link_st + expidx(top_exp))
            break;
        if (append_st)
            return append_only_malloc(size - ALIGNMENT);
        if (incr_top(SPAN_FACTOR << top_exp) == -1) {
            // Fall back to append only mode
            append_st = append_top;
            for (int i = MIN_SPAN_EXP; i <= top_exp; i++) // Undo
                if ((uint8_t*)link_st[expidx(i)].next >= append_st)
                    remove_span(link_st[expidx(i)].next);
            return append_only_malloc(size - ALIGNMENT);
        }
        if (link_st[expidx(top_exp)].prev == link_st + expidx(top_exp))
            insert_span((SpanHeader*)(top - (SPAN_FACTOR << top_exp)), top_exp);
        else {
            assert((uint8_t*)link_st[expidx(top_exp)].prev == top - (SPAN_FACTOR << (top_exp + 1)));
            remove_span(link_st[expidx(top_exp)].prev);
            insert_span((SpanHeader*)(top - (SPAN_FACTOR << (top_exp + 1))), top_exp + 1);
        }
    }

    // Split a larger span down
    SpanHeader *target_span = link_st[expidx(split_exp)].next;
    remove_span(target_span);
    for (split_exp--; split_exp >= span_exp; split_exp--) {
        uint8_t *buddy = (uint8_t*)target_span + (SPAN_FACTOR << split_exp);
        insert_span((SpanHeader*)buddy, split_exp);
    }

    // Setup the header (necessary because not stored into a list)
    target_span->span_exp = span_exp;
    target_span->is_free = false;
    return (uint8_t*)target_span + ALIGNMENT;
}

void my_free(void *_ptr) {
    if (_ptr == NULL)
        return;
    if (append_st && (uint8_t*)_ptr >= append_st)
        return;
    uint8_t *ptr = (uint8_t*)_ptr - ALIGNMENT;
    int span_exp = ((SpanHeader*)ptr)->span_exp;
    assert(!((SpanHeader*)ptr)->is_free);
    for (; span_exp < top_exp; span_exp++) {
        size_t ptr_no = (ptr - span_st) / SPAN_FACTOR;
        size_t buddy_no = ptr_no ^ (1 << span_exp);
        SpanHeader *buddy = (SpanHeader*)(buddy_no * SPAN_FACTOR + span_st);
        assert(buddy->span_exp <= span_exp);
        if (!buddy->is_free || buddy->span_exp != span_exp)
            break;
        remove_span(buddy);
        ptr = (ptr_no & buddy_no) * SPAN_FACTOR + span_st;
    }
    insert_span((SpanHeader*)ptr, span_exp);
}

void *my_realloc(void *_ptr, size_t size) {
    if (_ptr == NULL)
        return _ptr = my_malloc(size);

    if (size == 0) {
        my_free(_ptr);
        return _ptr = NULL;
    }

    if (!append_st || (uint8_t*)_ptr < append_st) {
        uint8_t *ptr = (uint8_t*)_ptr - ALIGNMENT;
        int span_exp = ((SpanHeader*)ptr)->span_exp;
        assert(!((SpanHeader*)ptr)->is_free);
        for (; span_exp < top_exp && (SPAN_FACTOR << span_exp) < size + ALIGNMENT; span_exp++) {
            size_t ptr_no = (ptr - span_st) / SPAN_FACTOR;
            size_t buddy_no = ptr_no ^ (1 << span_exp);
            SpanHeader *buddy = (SpanHeader*)(buddy_no * SPAN_FACTOR + span_st);
            assert(buddy->span_exp <= span_exp);
            if ((uint8_t*)buddy < ptr || !buddy->is_free || buddy->span_exp != span_exp)
                break;
            remove_span(buddy);
        }
        ((SpanHeader*)ptr)->span_exp = span_exp;
        if ((SPAN_FACTOR << span_exp) >= size + ALIGNMENT)
            return _ptr;
    }

    void *new_ptr = my_malloc(size);
    memcpy(new_ptr, _ptr, size); // Use memmove if overlapping
    my_free(_ptr);
    _ptr = new_ptr;
    return _ptr;
}