#include "IroMalloc.h"
#include "compiler/CompilerTools.h"

#define FLAGMASK 0xF

typedef struct Block {
    struct Block *prev;
    struct Block *next;
    void *x8;
    void *xC;
    size_t remain;
    size_t x14;
} Block;

typedef struct SubBlockBase {
    size_t x0;
    Block *block;
} SubBlockBase;

typedef struct SubBlock {
    size_t x0;
    Block *block;
    struct SubBlock *x8;
    struct SubBlock *xC;
} SubBlock;

typedef struct SubBlockTail {
    size_t x0copy;
} SubBlockTail;

typedef struct BlockTail {
    size_t x14copy;
    SubBlock *unk;
} BlockTail;

typedef struct FixStart {
    struct FixBlock *a;
    struct FixSubBlock *b;
    long count;
} FixStart;

typedef struct FixBlock {
    struct FixBlock *prev;
    struct FixBlock *next;
    size_t entrysize;
} FixBlock;

typedef struct FixSubBlockBase {
    FixBlock *fixblock;
} FixSubBlockBase;

typedef struct FixSubBlock {
    FixBlock *fixblock;
    struct FixSubBlock *next;
} FixSubBlock;

static const size_t fix_pool_sizes[] = {0xC, 0x1C, 0x2C, 0x4C};
static Block *start_;
static FixStart fix_start[4];
static int initialized;

// forward decls
static void Block_link(Block *block, SubBlock *subblock);
static void Block_unlink(Block *block, SubBlock *subblock);
static void SubBlock_construct(SubBlock *subblock, size_t size, Block *block, int flag1, int flag2);
static SubBlock *SubBlock_split(SubBlock *subblock, size_t pos);
static SubBlock *SubBlock_merge_prev(SubBlock *subblock, SubBlock **sbptr);
static void SubBlock_merge_next(SubBlock *subblock, SubBlock **sbptr);

#define BLOCK_TAIL(block) ((BlockTail *) ((long) (block) + ((block)->x14 & ~FLAGMASK) - sizeof(BlockTail)))

#define BUF_LOCATOR(buf) (*((size_t *) ((long) (buf) - sizeof(size_t))))
#define SBB_BLOCK(sbb) ((Block *) ((long) ((sbb)->block) & ~1))

#define BUF_IS_VAR(buf) (BUF_LOCATOR(buf) & 1)
#define VARBUF_SBB(buf) ((SubBlockBase *) ((long) (buf) - sizeof(SubBlockBase)))
#define VARBUF_SB(buf) ((SubBlock *) ((long) (buf) - sizeof(SubBlockBase)))
#define VARBUF_BLOCKSIZE(buf) (VARBUF_SBB(buf)->x0 & ~FLAGMASK)
#define VARBUF_SIZE(buf) (((VARBUF_SBB(buf)->x0 & ~FLAGMASK) - sizeof(SubBlockBase)))
#define VARBUF_BLOCK(buf) SBB_BLOCK(VARBUF_SBB(buf))
#define FIXBUF_SIZE(buf) (((FixBlock *) BUF_LOCATOR(buf))->entrysize)

#define BUF_SIZE(buf) BUF_IS_VAR(buf) ? VARBUF_SIZE(buf) : FIXBUF_SIZE(buf)

static void Block_construct(Block *block, size_t size) {
    SubBlock *subblock;

    block->x14 = size | 3;
    ((BlockTail *) ((long) block + size - sizeof(BlockTail)))->x14copy = block->x14;

    subblock = (SubBlock *) (block + 1);
    SubBlock_construct(subblock, size - sizeof(Block) - sizeof(BlockTail), block, 0, 0);

    block->remain = size - sizeof(Block) - sizeof(BlockTail);
    BLOCK_TAIL(block)->unk = NULL;

    Block_link(block, subblock);
}

static SubBlock *Block_subBlock(Block *block, size_t size) {
    SubBlock *subblock;
    size_t check;
    size_t best;

    if (!BLOCK_TAIL(block)->unk)
        return NULL;

    subblock = BLOCK_TAIL(block)->unk;
    best = subblock->x0 & ~FLAGMASK;
    check = subblock->x0 & ~FLAGMASK;
    while (check < size) {
        subblock = subblock->xC;
        check = subblock->x0 & ~FLAGMASK;
        if (best < check)
            best = check;
        if (subblock == BLOCK_TAIL(block)->unk) {
            block->remain = best;
            return NULL;
        }
    }

    if ((check - size) >= 0x60)
        SubBlock_split(subblock, size);

    BLOCK_TAIL(block)->unk = subblock->xC;
    Block_unlink(block, subblock);
    return subblock;
}

static void Block_link(Block *block, SubBlock *subblock) {
    size_t size;
    SubBlock **sbptr;

    size = subblock->x0 & ~FLAGMASK;
    subblock->x0 &= ~2;
    ((SubBlock *) ((long) subblock + size))->x0 &= ~4;
    ((SubBlockTail *) ((long) subblock + size - sizeof(SubBlockTail)))->x0copy = size;

    sbptr = &BLOCK_TAIL(block)->unk;
    if (*sbptr) {
        subblock->x8 = (*sbptr)->x8;
        subblock->x8->xC = subblock;
        subblock->xC = *sbptr;
        (*sbptr)->x8 = subblock;
        *sbptr = subblock;
        *sbptr = SubBlock_merge_prev(*sbptr, sbptr);
        SubBlock_merge_next(*sbptr, sbptr);
    } else {
        *sbptr = subblock;
        subblock->x8 = subblock;
        subblock->xC = subblock;
    }

    if (block->remain < ((*sbptr)->x0 & ~FLAGMASK))
        block->remain = (*sbptr)->x0 & ~FLAGMASK;
}

static void Block_unlink(Block *block, SubBlock *subblock) {
    size_t size;
    SubBlock **sbptr;

    size = subblock->x0 & ~FLAGMASK;
    subblock->x0 |= 2;
    ((SubBlock *) ((long) subblock + size))->x0 |= 4;

    sbptr = &BLOCK_TAIL(block)->unk;
    if (*sbptr == subblock)
        *sbptr = subblock->xC;

    if (*sbptr == subblock) {
        *sbptr = NULL;
        block->remain = 0;
    } else {
        subblock->xC->x8 = subblock->x8;
        subblock->x8->xC = subblock->xC;
    }
}

static void SubBlock_construct(SubBlock *subblock, size_t size, Block *block, int flag1, int flag2) {
    subblock->block = (Block *) ((long) block | 1);
    subblock->x0 = size;
    if (flag1)
        subblock->x0 |= 4;
    if (flag2) {
        subblock->x0 |= 2;
        ((SubBlock *) (((long) subblock) + size))->x0 |= 4;
    } else {
        ((SubBlockTail *) (((long) subblock) + size - sizeof(SubBlockTail)))->x0copy = size;
    }
}

static SubBlock *SubBlock_split(SubBlock *subblock, size_t pos) {
    size_t oldsize;
    int flag;
    SubBlock *splitright;
    Block *block;

    oldsize = subblock->x0 & ~FLAGMASK;
    flag = !(subblock->x0 & 2);
    splitright = (SubBlock *) ((long) subblock + pos);
    block = (Block *) ((long) subblock->block & ~1);
    SubBlock_construct(subblock, pos, block, subblock->x0 & 4, !flag);
    SubBlock_construct(splitright, oldsize - pos, block, !flag, !flag);
    if (flag) {
        splitright->xC = subblock->xC;
        splitright->xC->x8 = splitright;
        splitright->x8 = subblock;
        subblock->xC = splitright;
    }

    return splitright;
}

static SubBlock *SubBlock_merge_prev(SubBlock *subblock, SubBlock **sbptr) {
    size_t prevsize;
    SubBlock *prevblock;

    if (!(subblock->x0 & 4)) {
        prevsize = ((SubBlockTail *) ((long) subblock - sizeof(SubBlockTail)))->x0copy;
        if (prevsize & 2)
            return subblock;

        prevblock = (SubBlock *) ((long) subblock - prevsize);
        prevblock->x0 = prevblock->x0 & FLAGMASK;
        prevblock->x0 |= (prevsize + (subblock->x0 & ~FLAGMASK)) & ~FLAGMASK;
        if (!(prevblock->x0 & 2))
            ((SubBlockTail *) ((long) prevblock + prevsize + (subblock->x0 & ~FLAGMASK) - sizeof(SubBlockTail)))->x0copy = prevsize + (subblock->x0 & ~FLAGMASK);

        if (*sbptr == subblock)
            *sbptr = (*sbptr)->xC;

        subblock->xC->x8 = subblock->x8;
        subblock->xC->x8->xC = subblock->xC;
        return prevblock;
    }

    return subblock;
}

static void SubBlock_merge_next(SubBlock *subblock, SubBlock **sbptr) {
    SubBlock *nextblock;
    size_t nextsize;

    nextblock = (SubBlock *) ((long) subblock + (subblock->x0 & ~FLAGMASK));
    if (!(nextblock->x0 & 2)) {
        nextsize = (subblock->x0 & ~FLAGMASK) + (nextblock->x0 & ~FLAGMASK);
        subblock->x0 = subblock->x0 & FLAGMASK;
        subblock->x0 |= nextsize & ~FLAGMASK;

        if (!(subblock->x0 & 2))
            ((SubBlockTail *) ((long) subblock + nextsize - sizeof(SubBlockTail)))->x0copy = nextsize;

        if (!(subblock->x0 & 2))
            ((SubBlock *) ((long) subblock + nextsize))->x0 &= ~4;
        else
            ((SubBlock *) ((long) subblock + nextsize))->x0 |= 4;

        if (*sbptr == nextblock)
            *sbptr = (*sbptr)->xC;
        if (*sbptr == nextblock)
            *sbptr = NULL;

        nextblock->xC->x8 = nextblock->x8;
        nextblock->x8->xC = nextblock->xC;
    }
}

static void link_block(Block *block) {
    if (start_) {
        block->prev = start_->prev;
        block->prev->next = block;
        block->next = start_;

        start_->prev = block;
        start_ = block;
    } else {
        start_ = block;
        block->prev = block;
        block->next = block;
    }
}

static Block *unlink_block(Block *block) {
    Block *newblock;

    newblock = block->next;
    if (newblock == block)
        newblock = NULL;

    if (start_ == block)
        start_ = newblock;

    if (newblock) {
        newblock->prev = block->prev;
        newblock->prev->next = newblock;
    }

    block->prev = block->next = NULL;
    return newblock;
}

static Block *link_new_block(size_t size) {
    Block *block;

    size = (size + 0x1000 + sizeof(Block) + sizeof(BlockTail) - 1) & ~0xFFF;
    if (size < 0x10000)
        size = 0x10000;

    if (!(block = galloc(size)))
        return NULL;

    Block_construct(block, size);
    link_block(block);
    return block;
}

static void *allocate_from_var_pools(size_t size) {
    Block *block;
    SubBlock *subblock;

    size = (size + sizeof(Block) - 1) & ~0xF;
    if (size < 0x60)
        size = 0x60;

    block = start_ ? start_ : link_new_block(size);
    if (!block)
        return NULL;

    do {
        if (size <= block->remain) {
            if ((subblock = Block_subBlock(block, size))) {
                start_ = block;
                goto allocated;
            }
        }
    } while ((block = block->next) != start_);

    block = link_new_block(size);
    if (!block)
        return NULL;
    subblock = Block_subBlock(block, size);
allocated:
    return (SubBlockBase *) subblock + 1;
}

static void deallocate_from_var_pools(void *buf) {
    Block *block;
    SubBlock *subblock;
    int flag;

    subblock = (SubBlock *) ((long) buf - sizeof(SubBlockBase));
    block = (Block *) ((long) subblock->block & ~1);
    Block_link(block, subblock);

    flag = 0;
    subblock = (SubBlock *) (block + 1);
    if (!(subblock->x0 & 2)) {
        if ((subblock->x0 & ~FLAGMASK) == ((block->x14 & ~FLAGMASK) - sizeof(Block) - sizeof(BlockTail)))
            flag = 1;
    }

    if (flag)
        unlink_block(block);
}

static void FixBlock_construct(FixBlock *fixblock, FixBlock *prev, FixBlock *next, int poolIndex, void *buffer, size_t buffersize) {
    size_t entrysize;
    size_t entrycount;
    size_t i;
    FixSubBlock *fsb;
    FixSubBlock *fsbnext;

    fixblock->prev = prev;
    fixblock->next = next;
    fixblock->entrysize = fix_pool_sizes[poolIndex];

    entrysize = fix_pool_sizes[poolIndex] + sizeof(void *);
    entrycount = (buffersize / entrysize);
    fsb = buffer;
    for (i = 0; i < (entrycount - 1); i++) {
        fsbnext = (FixSubBlock *) ((long) fsb + entrysize);
        fsb->fixblock = fixblock;
        fsb->next = fsbnext;
        fsb = fsbnext;
    }

    fsb->fixblock = fixblock;
    fsb->next = fix_start[poolIndex].b;
    fix_start[poolIndex].b = buffer;
}

static void *allocate_from_fixed_pools(size_t size) {
    int poolIndex;
    FixSubBlock *fsb;

    poolIndex = 0;
    while (size > fix_pool_sizes[poolIndex])
        poolIndex++;

    if (!fix_start[poolIndex].b) {
        void *buf;
        size_t bufsize;

        buf = allocate_from_var_pools(4000);
        if (!buf)
            return NULL;

        bufsize = !BUF_IS_VAR(buf) ? FIXBUF_SIZE(buf) : VARBUF_SIZE(buf);
        FixBlock_construct(
                buf, NULL, fix_start[poolIndex].a, poolIndex,
                ((FixBlock *) buf) + 1,
                bufsize - sizeof(FixBlock)
        );
        fix_start[poolIndex].a = buf;
    }

    fsb = fix_start[poolIndex].b;
    fix_start[poolIndex].b = fsb->next;
    fix_start[poolIndex].count++;
    return ((FixSubBlockBase *) fsb) + 1;
}

static void deallocate_from_fixed_pools(void *buf, size_t size) {
    int poolIndex;
    FixBlock *fixblock;
    FixBlock *nextfb;
    FixSubBlock *fsb;

    poolIndex = 0;
    while (size > fix_pool_sizes[poolIndex])
        poolIndex++;

    fsb = (FixSubBlock *) ((long) buf - sizeof(FixSubBlockBase));
    fsb->next = fix_start[poolIndex].b;
    fix_start[poolIndex].b = fsb;
    if (--fix_start[poolIndex].count == 0) {
        fixblock = fix_start[poolIndex].a;
        while (fixblock) {
            nextfb = fixblock->next;
            deallocate_from_var_pools(fixblock);
            fixblock = nextfb;
        }
        fix_start[poolIndex].a = NULL;
        fix_start[poolIndex].b = NULL;
    }
}

size_t IRO_msize(void *buf) {
    return BUF_SIZE(buf);
}

void *IRO_malloc(size_t size) {
    if (!size)
        return NULL;

    if (size <= 0x4C)
        return allocate_from_fixed_pools(size);
    else
        return allocate_from_var_pools(size);
}

void IRO_free(void *buf) {
    if (buf) {
        size_t size = !BUF_IS_VAR(buf) ? FIXBUF_SIZE(buf) : VARBUF_SIZE(buf);

        if (size <= 0x4C)
            deallocate_from_fixed_pools(buf, size);
        else
            deallocate_from_var_pools(buf);
    }
}

void *IRO_realloc(void *buf, size_t newsize) {
    size_t oldsize;
    size_t newblocksize;
    void *newbuf;

    if (!buf)
        return IRO_malloc(newsize);

    if (!newsize) {
        IRO_free(buf);
        return NULL;
    }

    oldsize = !BUF_IS_VAR(buf) ? FIXBUF_SIZE(buf) : VARBUF_SIZE(buf);
    if (newsize > oldsize) {
        if (BUF_IS_VAR(buf)) {
            newblocksize = (newsize + sizeof(Block) - 1) & ~FLAGMASK;
            if (newblocksize < 0x60)
                newblocksize = 0x60;
            SubBlock_merge_next(VARBUF_SB(buf), &BLOCK_TAIL(VARBUF_BLOCK(buf))->unk);
            if (VARBUF_BLOCKSIZE(buf) >= newblocksize) {
                if ((VARBUF_BLOCKSIZE(buf) - newblocksize) >= 0x60) {
                    Block_link(VARBUF_BLOCK(buf), SubBlock_split(VARBUF_SB(buf), newblocksize));
                }
                return buf;
            }
        }

        newbuf = IRO_malloc(newsize);
        if (!newbuf)
            return NULL;

        memcpy(newbuf, buf, oldsize);
        IRO_free(buf);
        return newbuf;
    }

    if (BUF_IS_VAR(buf)) {
        newsize = (newsize + sizeof(Block) - 1) & ~FLAGMASK;
        if (newsize < 0x60)
            newsize = 0x60;
        if ((VARBUF_BLOCKSIZE(buf) - newsize) >= 0x60) {
            Block_link(VARBUF_BLOCK(buf), SubBlock_split(VARBUF_SB(buf), newsize));
        }
    }

    return buf;
}

void *IRO_calloc(size_t a, size_t b) {
    void *buf;
    size_t len = b * a;
    buf = IRO_malloc(len);
    if (buf)
        memset(buf, 0, len);
    return buf;
}

void IRO_pool_free_all(void) {
    int i;
    Block *block;
    if ((block = start_)) {
        do {
            block = block->next;
        } while (block != start_);
        start_ = NULL;

        for (i = 0; i < 4; i++) {
            fix_start[i].a = NULL;
            fix_start[i].b = NULL;
            fix_start[i].count = 0;
        }
    }
}

void IRO_InitializeAllocator(void) {
    if (!initialized) {
        int i;
        for (i = 0; i < 4; i++) {
            fix_start[i].a = NULL;
            fix_start[i].b = NULL;
            fix_start[i].count = 0;
        }
        start_ = NULL;
        initialized = 1;
    }
}

void IRO_TerminateAllocator(void) {
    initialized = 0;
}