#define _GNU_SOURCE
#include <fcntl.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <string.h>
#include <unistd.h>
#include <stddef.h>
#include <stdbool.h>
#include <signal.h>
#include <stdint.h>

// Limits on page size and sizes that can be allocated by this library
#define PAGESIZE 4096
#define MIN_ALLOCATION_SIZE 2
#define MAX_ALLOCATION_SIZE 1024 // 1024 is maximum power of 2 (2^10) that can 
                                 // be fit multiple times in a page size of 
                                 // 4096 (2^12), accounting for page header
#define NUM_PAGE_SIZES 10
#define MAX_ALLOCATION_EXPONENT 11

// Struct that holds details about an individual freed list
typedef struct Free {

    struct Free *next;
    void *allocation;

} free_t;

// Struct that holds details about an individual page
typedef struct Page {

    int allocationSize;
    // numAllocations declared as unsigned int for space efficiency
    unsigned int numAllocations : MAX_ALLOCATION_EXPONENT; // Bit field since 
                                                           // there is a 
                                                           // maximum number of 
                                                           // allocations that 
                                                           // can fit in a page

} page_t;

// Array of pointers to newest page of each allocation size
page_t *lastPages[NUM_PAGE_SIZES];

// Array of pointers to the last-generated allocation of each allocation size
void *lastAllocations[NUM_PAGE_SIZES];

// Array of pointers to free lists of each allocation size
free_t *freeLists[NUM_PAGE_SIZES];

// Function to get the next-highest multiple of two of a number
// Inspired by: 
// https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
int getAllocationSize(int size) {

    size--;
    size |= size >> 1;
    size |= size >> 2;
    size |= size >> 4;
    size |= size >> 8;
    size |= size >> 16;
    size++;

    return size;
    
}

// Function to get the index of lastPages that should be used for a given 
// alloction size
int getPageExponent(int size) {

    // Return the exponent calculated as the number of leading zeros in a 
    // value subtracted from the number of bits total in ints, minus 
    // 1 since arrays start at 0
    return ((((sizeof(int) * 
               (sizeof(int) + sizeof(int))) - 1) - 
               __builtin_clz(size)) - 1);

}

// malloc re-implementation
void *malloc(size_t size) {
    
    // If the size passed to malloc is equal to or less than 0, return NULL
    // Technically, compiler warnings will fly if a negative size is passed 
    // to malloc, but it can be done nevertheless so I'm catching it here
    if (size <= 0) {

        return NULL;

    }

    // Create an allocation pointer that will be updated conditionally and used 
    // to return from malloc
    void *newAllocation;

    // If the size passed to malloc is equal to or less than the maximum 
    // (small) allocation size, add to the existing page structure
    if (size <= MAX_ALLOCATION_SIZE) {

        // Call helper function getAllocationSize to find the next-lowest power 
        // of 2 that can fit the requested size
        int allocationSize = getAllocationSize(size);

        // Call helper gunction getPageExponent to find the array index 
        // matching the pointer for the most recent page of this allocation size
        int pageExponent = getPageExponent(allocationSize);

        // If the free list for this allocation size has any items, use 
        // the first freed block and return early
        if (freeLists[pageExponent] != NULL) {
            
            newAllocation = freeLists[pageExponent]->allocation;

            // Advance the free list
            freeLists[pageExponent] = freeLists[pageExponent]->next;

            // Get the parent page of the freed allocation being re-used 
            // using the same bitwise mask as in free() and realloc()
            page_t *ptemp = (page_t *)
            ((uintptr_t)newAllocation & ~((uintptr_t)0xFFF));

            // Increment the number of allocations in the page
            ptemp->numAllocations++;

            return newAllocation;

        }

        // Create a temporary page pointer
        page_t *ptemp = NULL;

        // If a page of this allocation size exists and is not full, set 
        // the temporary page pointer to that existing page
        if (lastPages[pageExponent] != NULL && 
            (PAGESIZE - sizeof(page_t)) / allocationSize >= 
            (lastPages[pageExponent]->numAllocations + 1)) {

            ptemp = lastPages[pageExponent];

        }

        // If no page of this allocation size that can be re-used has been found
        if (ptemp == NULL) {

            // Get a new page
            ptemp = mmap(NULL, PAGESIZE, PROT_READ | PROT_WRITE, 
                         MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);

            // Update page metadata
            ptemp->allocationSize = allocationSize;
            ptemp->numAllocations = 1;

            // Update the pointer to the newest page of this allocation size
            lastPages[pageExponent] = ptemp;

            // Set the offset for the first allocation
            newAllocation = (void *)((char *)ptemp + sizeof(page_t));

            // Update the pointer to the last-generated allocation of this 
            // allocation size
            lastAllocations[pageExponent] = (void *)
            ((char *)newAllocation + allocationSize);

            // If a page has been found with an allocation size matching 
            // the size requested
        } else {
            
            // Get the next available allocation in the current page
            newAllocation = lastAllocations[pageExponent];

            // Set the offset for the next allocation space
            lastAllocations[pageExponent] = (void *)
            ((char *)newAllocation + allocationSize);

            // Increment the number of allocations in the page
            ptemp->numAllocations++;

        }

        // Otherwise, allocate a separate chunk of memory
    } else {
        
        // Get a new page
        page_t *newPage = mmap(NULL, (sizeof(page_t) + size), 
                               PROT_READ | PROT_WRITE, 
                               MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);

        // Update page metadata
        newPage->allocationSize = size;

        // Update the pointer to the large allocation
        newAllocation = (void *)((char *)newPage + sizeof(page_t));

    }

    return newAllocation;

}

// calloc re-implementation
void *calloc(size_t numItems, size_t size) {
    
    // Create a new pointer by calling malloc
    void *temp = malloc(numItems * size);

    // Set all the bytes to 0
    memset(temp, 0, (numItems * size));

    return temp;

}

// free re-implementation
void free(void *ptr) {

    // If pointer being freed is NULL, return early
    if (ptr == NULL) {

        return;

    }
    
    // Get the parent page of the pointer passed to function using a bitwise 
    // mask that returns the memory address of the pointer given with the last 
    // 3 hexadecimals as 0
    // Syntax error about integrals was thrown, research rabbit hole lead me 
    // to a resource that explained uintptr_t might be useful for this 
    // operation: https://stackoverflow.com/a/1846648
    page_t *ptemp = (page_t *)((uintptr_t)ptr & ~((uintptr_t)0xFFF));

    // If the page's allocation size is greater than the maximum small 
    // allocation size, indicating the page is a large chunk, unmap 
    // the entire page and return
    if (ptemp->allocationSize > MAX_ALLOCATION_SIZE) {
        
        munmap(ptemp, (sizeof(page_t) + ptemp->allocationSize));

        return;

    }

    // Call helper gunction getPageExponent to find the array index 
    // matching this allocation size
    int pageExponent = getPageExponent(ptemp->allocationSize);

    // Create a new free list node
    free_t *ftemp = malloc(sizeof(free_t));

    // Set the new free list node's allocation to the allocation being freed
    ftemp->allocation = ptr;
    // Set the new free list node's next pointer to the current free list 
    // for this allocation size
    ftemp->next = freeLists[pageExponent];

    // Set the head of the free list for this allocation size to the new 
    // free list node
    freeLists[pageExponent] = ftemp;

    // Decrement the number of allocations in the page
    ptemp->numAllocations--;

}

// realloc re-implementation
void *realloc(void *ptr, size_t size) {

    // If size is less than or equal to 0, free the pointer given and return 
    // NULL per man page for realloc
    if (size <= 0) {

        free(ptr);

        return NULL;

        // If the pointer given is null, return a new pointer
    } else if (ptr == NULL) {

        return malloc(size);

    }

    // Get the parent page of the pointer passed to function using a bitwise 
    // mask that returns the memory address of the pointer given with the last 
    // 3 hexadecimals as 0
    // Same trick as in free()
    page_t *ptemp = (page_t *)((uintptr_t)ptr & ~((uintptr_t)0xFFF));

    // Allocate a new pointer of the specified size
    void *temp = malloc(size);

    // If the size is greater than or equal to the size of the original pointer 
    // allocation, copy memory up to the original allocation's size
    if (size >= ptemp->allocationSize) {

        memcpy(temp, ptr, ptemp->allocationSize);

        // If the size is less than the size of the original pointer 
        // allocation, copy memory up to the specified size
    } else {

        memcpy(temp, ptr, size);

    }

    // Free the old pointer
    free(ptr);

    return temp;

}