enum {
    PAGE_SIZE = 4096,
    MINIMUM_RANGE_SIZE = PAGE_SIZE
};

void *AllocateRange(uint32_t *size) {
    *size = (*size + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (*size < MINIMUM_RANGE_SIZE) {
        *size = MINIMUM_RANGE_SIZE;
    }
    return VirtualAlloc(0, *size, MEM_COMMIT, PAGE_READWRITE);
}

void FreeRange(void *pointer, uint32_t size) {
    VirtualFree(pointer, size, MEM_RELEASE);
}

struct Block {
    uint32_t free : 1;
    uint32_t size : 31;
    uint32_t previous_size;
};

struct FreeBlock {
    uint32_t free : 1;
    uint32_t size : 31;
    uint32_t previous_size;
    FreeBlock *previous_free;
    FreeBlock *next_free;
};

enum {
    LOG2_MINIMUM_BLOCK_SIZE = 5,
    MINIMUM_BLOCK_SIZE = 1 << LOG2_MINIMUM_BLOCK_SIZE,
    LOG2_CLASSES_PER_EXPONENT = 3,
    CLASSES_PER_EXPONENT = 1 << LOG2_CLASSES_PER_EXPONENT,
    MAX_CLASSES = (32 - LOG2_MINIMUM_BLOCK_SIZE) * CLASSES_PER_EXPONENT
};

FreeBlock *free_blocks[MAX_CLASSES];

uint32_t GetFreeBlockIndex(uint32_t size) {
    Assert(size >= MINIMUM_BLOCK_SIZE);
    DWORD exponent;
    _BitScanReverse(&exponent, size);
    uint32_t remainder = size - (1 << exponent);
    int32_t shift = exponent - LOG2_CLASSES_PER_EXPONENT;
    if (shift < 0) {
        shift = 0;
    }
    uint32_t offset = remainder >> shift;
    Assert(offset < CLASSES_PER_EXPONENT);
    return (exponent - LOG2_MINIMUM_BLOCK_SIZE) * CLASSES_PER_EXPONENT + offset;
}

void InsertFreeBlock(FreeBlock *block) {
    uint32_t index = GetFreeBlockIndex(block->size);
    FreeBlock *first_free = free_blocks[index];
    if (first_free) {
        first_free->previous_free = block;
    }
    block->previous_free = 0;
    block->next_free = first_free;
    free_blocks[index] = block;
}

void RemoveFreeBlock(FreeBlock *block) {
    if (block->previous_free) {
        block->previous_free->next_free = block->next_free;
    } else {
        uint32_t index = GetFreeBlockIndex(block->size);
        Assert(free_blocks[index] == block);
        free_blocks[index] = block->next_free;
    }
    if (block->next_free) {
        block->next_free->previous_free = block->previous_free;
    }
}

void AllocateFreeBlock(FreeBlock *block, uint32_t size) {
    Assert(block->free);
    Assert(size >= MINIMUM_BLOCK_SIZE);
    Assert(size <= block->size);
    RemoveFreeBlock(block);
    if (block->size - size >= MINIMUM_BLOCK_SIZE) {
        uint32_t new_size = block->size - size;
        FreeBlock *new_block = (FreeBlock *)((char *)block + size);
        Block *next_block = (Block *)((char*)block + block->size);
        next_block->previous_size = new_size;
        block->size = size;
        new_block->free = 1;
        new_block->size = new_size;
        new_block->previous_size = size;
        new_block->next_free = block->next_free;
        InsertFreeBlock(new_block);
    }
    block->free = 0;
}

void *TryAllocate(uint32_t size) {
    size += sizeof(Block);
    if (size < MINIMUM_BLOCK_SIZE) {
        size = MINIMUM_BLOCK_SIZE;
    }
    bool retried = false;
    for (uint32_t index = GetFreeBlockIndex(size); index < MAX_CLASSES; index++) {
        FreeBlock *free_block = free_blocks[index];
        if (free_block) {
            if (free_block->size >= size) {
                AllocateFreeBlock(free_block, size);
                return (char *)free_block + sizeof(Block);
            } else {
                Assert(!retried);
                retried = true;
            }
        }
    }
    return 0;
}

void FreeAllocatedBlock(Block *block) {
    Block *previous_block = (Block *)((char *)block - block->previous_size);
    Block *next_block = (Block *)((char *)block + block->size);
    if (previous_block->free) {
        RemoveFreeBlock((FreeBlock *)previous_block);
        previous_block->size += block->size;
        next_block->previous_size = previous_block->size;
        Block *previous_previous_block = (Block *)((char *)previous_block - previous_block->previous_size);
        block = previous_block;
        previous_block = previous_previous_block;
    }
    if (next_block->free) {
        RemoveFreeBlock((FreeBlock *)next_block);
        block->size += next_block->size;
        Block *next_next_block = (Block *)((char *)next_block + next_block->size);
        next_next_block->previous_size = block->size;
    }
    block->free = 1;
    InsertFreeBlock((FreeBlock *)block);
}

void Free(void *pointer) {
    if (pointer) {
        Block *block = (Block *)((char *)pointer - sizeof(Block));
        Assert(block->free == 0);
        FreeAllocatedBlock(block);
    }
}

void InsertFreeRange(void *base, uint32_t size) {
    Assert(size % PAGE_SIZE == 0);
    Assert(size >= 3 * PAGE_SIZE);
    Block *head = (Block *)base;
    head->free = 0;
    head->size = PAGE_SIZE;
    head->previous_size = 0;
    FreeBlock *block = (FreeBlock *)((char *)base + PAGE_SIZE);
    block->free = 1;
    block->size = size - 2 * PAGE_SIZE;
    block->previous_size = head->size;
    Block *tail = (Block *)((char *)base + size - PAGE_SIZE);
    tail->free = 0;
    tail->size = PAGE_SIZE;
    tail->previous_size = block->size;
    uint32_t index = GetFreeBlockIndex(block->size);
    block->previous_free = 0;
    block->next_free = free_blocks[index];
    free_blocks[index] = block;
}

void *Allocate(uint32_t size) {
    void *pointer = TryAllocate(size);
    if (!pointer) {
        uint32_t range_size = (size + 4 * PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
        void *range_base = AllocateRange(&range_size);
        Printf("Allocating new range.\n");
        InsertFreeRange(range_base, range_size);
        pointer = TryAllocate(size);
        Assert(pointer);
    }
    return pointer;
}