/*
 * Copyright 2024, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 */

#include <OS.h>
#include <malloc.h>
#include <pthread.h>
#include <sys/param.h>

#include <errno_private.h>
#include <libroot_private.h>
#include <shared/locks.h>
#include <system/tls.h>

#include "PagesAllocator.h"


/* generic stuff */
#if B_PAGE_SIZE == 4096
#define _MAX_PAGE_SHIFT 12
#endif

extern char* __progname;

#define MAP_CONCEAL                     (0)
#define __MAP_NOREPLACE         (0)

#define DEF_STRONG(X)
#define DEF_WEAK(x)


/* entropy routines, wrapped for malloc */

static uint32_t
malloc_arc4random()
{
        // TODO: Improve this?
        const uintptr_t address = (uintptr_t)&address;
        uint32_t random = (uint32_t)address * 1103515245;
        random += (uint32_t)system_time();
        return random;
}
#define arc4random malloc_arc4random


static void
malloc_arc4random_buf(void *_buf, size_t nbytes)
{
        uint8* buf = (uint8*)_buf;
        while (nbytes > 0) {
                uint32_t value = malloc_arc4random();
                const size_t copy = (nbytes > 4) ? 4 : nbytes;
                memcpy(buf, &value, copy);
                nbytes -= copy;
                buf += copy;
        }
}
#define arc4random_buf malloc_arc4random_buf


static uint32_t
malloc_arc4random_uniform(uint32_t upper_bound)
{
        return (malloc_arc4random() % upper_bound);
}
#define arc4random_uniform malloc_arc4random_uniform


/* memory routines, wrapped for malloc */

static inline void
malloc_explicit_bzero(void* buf, size_t len)
{
        memset(buf, 0, len);
}
#define explicit_bzero malloc_explicit_bzero


static int
malloc_mprotect(void* address, size_t length, int protection)
{
        /* do nothing */
        return 0;
}
#define mprotect malloc_mprotect


static int
mimmutable(void* address, size_t length)
{
        /* do nothing */
        return 0;
}


/* memory mapping */


static void*
malloc_mmap(void* address, size_t length, int protection, int flags, int fd, off_t offset)
{
        status_t status;
        uint8 cleared;

        length = (length + (B_PAGE_SIZE - 1)) & ~(B_PAGE_SIZE - 1);;

        status = __allocate_pages(&address, length, flags, &cleared);
        if (status != B_OK) {
                __set_errno(status);
                return MAP_FAILED;
        }

        if (!cleared)
                memset(address, 0, length);
        return address;
}
#define mmap malloc_mmap


static int
malloc_munmap(void* address, size_t length)
{
        status_t status = __free_pages(address, length);
        if (status < 0) {
                errno = status;
                return -1;
        }
        return 0;
}
#define munmap malloc_munmap


/* public methods */

void*
memalign(size_t align, size_t len)
{
        void* result = NULL;
        int status;

        if (align < sizeof(void*))
                align = sizeof(void*);

        status = posix_memalign(&result, align, len);
        if (status != 0) {
                errno = status;
                return NULL;
        }
        return result;
}


void*
valloc(size_t size)
{
        return memalign(B_PAGE_SIZE, size);
}


/* malloc implementation */

#define _MALLOC_MUTEXES 32
static mutex sMallocMutexes[_MALLOC_MUTEXES];
static pthread_once_t sThreadedMallocInitOnce = PTHREAD_ONCE_INIT;

static int32 sNextMallocThreadID = 1;

static u_int mopts_nmutexes();
static void _malloc_init(int from_rthreads);


static inline void
_MALLOC_LOCK(int32 index)
{
        mutex_lock(&sMallocMutexes[index]);
}


static inline void
_MALLOC_UNLOCK(int32 index)
{
        mutex_unlock(&sMallocMutexes[index]);
}


static void
init_threaded_malloc()
{
        u_int i;
        for (i = 2; i < _MALLOC_MUTEXES; i++)
                mutex_init(&sMallocMutexes[i], "heap mutex");

        _MALLOC_LOCK(0);
        _malloc_init(1);
        _MALLOC_UNLOCK(0);
}


status_t
__init_heap()
{
        tls_set(TLS_MALLOC_SLOT, (void*)0);
        __init_pages_allocator();
        mutex_init(&sMallocMutexes[0], "heap mutex");
        mutex_init(&sMallocMutexes[1], "heap mutex");
        _malloc_init(0);
        return B_OK;
}


static int32
get_thread_malloc_id()
{
        int32 result = (int32)(intptr_t)tls_get(TLS_MALLOC_SLOT);
        if (result == -1) {
                // thread has never called malloc() before; assign it an ID.
                result = atomic_add(&sNextMallocThreadID, 1);
                tls_set(TLS_MALLOC_SLOT, (void*)(intptr_t)result);
        }
        return result;
}


void
__heap_thread_init()
{
        pthread_once(&sThreadedMallocInitOnce, &init_threaded_malloc);
        tls_set(TLS_MALLOC_SLOT, (void*)(intptr_t)-1);
}


void
__heap_thread_exit()
{
        const int32 id = (int32)(intptr_t)tls_get(TLS_MALLOC_SLOT);
        if (id != -1 && id == (sNextMallocThreadID - 1)) {
                // Try to "de-allocate" this thread's ID.
                atomic_test_and_set(&sNextMallocThreadID, id, id + 1);
        }
}


void
__heap_before_fork()
{
        u_int i;
        u_int nmutexes = mopts_nmutexes();
        for (i = 0; i < nmutexes; i++)
                _MALLOC_LOCK(i);

        __pages_allocator_before_fork();
}


void
__heap_after_fork_child()
{
        u_int i;
        u_int nmutexes = mopts_nmutexes();
        for (i = 0; i < nmutexes; i++)
                mutex_init(&sMallocMutexes[i], "heap mutex");

        __pages_allocator_after_fork(0);
}


void
__heap_after_fork_parent()
{
        u_int i;
        u_int nmutexes = mopts_nmutexes();
        for (i = 0; i < nmutexes; i++)
                _MALLOC_UNLOCK(i);

        __pages_allocator_after_fork(1);
}


void
__heap_terminate_after()
{
}