/*
 * Copyright (C) 2020-2021 Aloysius Indrayanto
 *
 * This file is part of the eCxx library, see LICENSE file for the license details.
 */

/*
 * ----------------------------------------------------------------------------------
 * Some of the code in this file is copied (with a few modifications) from:
 *     heap.cpp - overrides of SDK heap handling functions
 * ----------------------------------------------------------------------------------
 * Copyright (C) 2016 Ivan Grokhotkov. All rights reserved.
 *
 * This file is distributed under MIT license.
 * ----------------------------------------------------------------------------------
 */

/*
 * ----------------------------------------------------------------------------------
 * Some of the code in this file is copied (with a few modifications) from:
 *     Esp.cpp - esp8266-specific APIs
 * ----------------------------------------------------------------------------------
 * Copyright (C) 2015 Ivan Grokhotkov. All rights reserved.
 *
 * This file is part of the esp8266 core for Arduino environment.
 *
 * This library is free software; you can redistribute it and/or modify it under the
 * terms of the GNU Lesser General Public License as published by the Free Software
 * Foundation; either version 2.1 of the License, or (at your option) any later
 * version.
 *
 * This library is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
 * PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License along
 * with this library; if not, write to the Free Software Foundation, Inc., 51
 * Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
 * ----------------------------------------------------------------------------------
 */


#include <umm_malloc.h>
#include <user_interface.h>

#include "CoreSys.h"


/*
 * About the ordering of integrity check, full poison check, and the specific *alloc functions
 *
 * 1. Integrity check:
 *        a. verifies the heap management information is not corrupt;
 *        b. this allows any other testing, that walks the heap, to run safely.
 *
 * 2. Place full poison check before or after a specific *alloc function:
 *        a. after  => when the *alloc function operates on an existing allocation;
 *        b. before => when the *alloc function creates a new, not modified, allocation.
 *
 *    In a free() or realloc() call, the focus is on their allocation. It is checked first and
 *    reported on first if an error exists. Full posion check is done after.
 *
 *    For malloc(), calloc(), and zalloc(), full posion check is done first because these
 *    functions do not modify an existing allocation.
 */


#if defined(UMM_POISON_CHECK) || defined(UMM_POISON_CHECK_LITE)

    #define UMM_MALLOC(S)              umm_poison_malloc(S)
    #define UMM_CALLOC(N, S)           umm_poison_calloc(N, S)
    #define UMM_REALLOC_FL(P, S, F, L) umm_poison_realloc_fl(P, S, F, L)
    #define UMM_FREE_FL(P, F, L)       umm_poison_free_fl(P, F, L)

    #undef realloc
    #undef free

#elif defined(DEBUG_ESP_OOM) || defined(UMM_INTEGRITY_CHECK)

    #define UMM_MALLOC(S)              umm_malloc(S)
    #define UMM_CALLOC(N, S)           umm_calloc(N, S)
    #define UMM_REALLOC_FL(P, S, F, L) umm_realloc(P, S)
    #define UMM_FREE_FL(P, F, L)       umm_free(P)

    #undef realloc
    #undef free

#else

    #define UMM_MALLOC(S)              malloc(S)
    #define UMM_CALLOC(N, S)           calloc(N, S)
    #define UMM_REALLOC_FL(P, S, F, L) realloc(P, S)
    #define UMM_FREE_FL(P, F, L)       free(P)

#endif


#if defined(UMM_POISON_CHECK)

    #define POISON_CHECK___ABORT()          \
        do {                                \
            if( ! POISON_CHECK() ) abort(); \
        } while(0)

    #define POISON_CHECK___PANIC_FL(FILE, LINE)                  \
        do {                                                     \
            if( ! POISON_CHECK() ) __panic_func(FILE, LINE, ""); \
        } while(0)

#else

    #define POISON_CHECK___ABORT() do {} while(0)
    #define POISON_CHECK___PANIC_FL(FILE, LINE) do { eCxx::ignoreUnused(FILE); eCxx::ignoreUnused(LINE); } while(0)

#endif


#if defined(DEBUG_ESP_OOM)

    #define PTR_CHECK___LOG_LAST_FAIL_FL(P, S, F, L)              \
        if( (S) && !(P) ) {                                       \
          umm_last_fail_alloc_addr = __builtin_return_address(0); \
          umm_last_fail_alloc_size = S;                           \
          umm_last_fail_alloc_file = F;                           \
          umm_last_fail_alloc_line = L;                           \
        }

    #define PTR_CHECK___LOG_LAST_FAIL(P, S)                       \
        if( (S) && !(P) ) {                                       \
          umm_last_fail_alloc_addr = __builtin_return_address(0); \
          umm_last_fail_alloc_size = S;                           \
          umm_last_fail_alloc_file = nullptr;                     \
          umm_last_fail_alloc_line = 0;                           \
        }

#else

    #define PTR_CHECK___LOG_LAST_FAIL_FL(P, S, F, L)              \
        if( (S) && !(P) ) {                                       \
          umm_last_fail_alloc_addr = __builtin_return_address(0); \
          umm_last_fail_alloc_size = S;                           \
          eCxx::ignoreUnused(F);                                  \
          eCxx::ignoreUnused(L);                                  \
        }

    #define PTR_CHECK___LOG_LAST_FAIL(P, S)                       \
        if( (S) && !(P) ) {                                       \
          umm_last_fail_alloc_addr = __builtin_return_address(0); \
          umm_last_fail_alloc_size = S;                           \
        }

#endif


#if defined(UMM_INTEGRITY_CHECK)

    #define INTEGRITY_CHECK___ABORT()          \
        do {                                   \
            if( ! INTEGRITY_CHECK() ) abort(); \
        } while(0)

    #define INTEGRITY_CHECK___PANIC_FL(FILE, LINE)                  \
        do {                                                        \
            if( ! INTEGRITY_CHECK() ) __panic_func(FILE, LINE, ""); \
        } while(0)

#else

    #define INTEGRITY_CHECK___ABORT() do {} while(0)
    #define INTEGRITY_CHECK___PANIC_FL(FILE, LINE) do {} while(0)

#endif


extern "C" {


// Debugging helper variables - last allocation which returned null
void*         umm_last_fail_alloc_addr = nullptr;
eCxx::int32_t umm_last_fail_alloc_size = 0;

#if defined(DEBUG_ESP_OOM)
const char*   umm_last_fail_alloc_file = nullptr;
eCxx::int32_t umm_last_fail_alloc_line = 0;
#endif


extern void* _malloc_r(_reent*, size_t size)
{
    void* ret = malloc(size);

    PTR_CHECK___LOG_LAST_FAIL(ret, size);

    return ret;
}


extern void _free_r(_reent*, void* ptr)
{ free(ptr); }


extern void* _realloc_r(_reent*, void* ptr, size_t size)
{
    void* ret = realloc(ptr, size);

    PTR_CHECK___LOG_LAST_FAIL(ret, size);

    return ret;
}


extern void* _calloc_r(_reent*, size_t count, size_t size)
{
    void* ret = calloc(count, size);

    PTR_CHECK___LOG_LAST_FAIL(ret, count * size);

    return ret;
}


#ifdef DEBUG_ESP_OOM

    #undef malloc
    #undef calloc
    #undef realloc

    #define DEBUG_HEAP_PRINTF ets_uart_printf

    extern void ICACHE_RAM_ATTR print_oom_size(eCxx::int32_t size)
    { if( system_get_os_print() ) DEBUG_HEAP_PRINTF(":oom(%d)@?\n", size ); }

    extern void ICACHE_RAM_ATTR print_loc(eCxx::int32_t size, const char* file, eCxx::int32_t line)
    {
        if( !system_get_os_print() ) return;

        DEBUG_HEAP_PRINTF(":oom(%d)@", size);

        const bool inISR = ETS_INTR_WITHINISR();

        if( inISR && reinterpret_cast<size_t>(file) >= 0x40200000UL ) {
            DEBUG_HEAP_PRINTF("File: %p", file);
        }
        else if( !inISR && reinterpret_cast<size_t>(file) >= 0x40200000UL ) {
            char buf[ ets_strlen(file) + 1] __ECXX_ATTR_ALIGNED_4;
            ets_strcpy(buf, file);
            DEBUG_HEAP_PRINTF(buf);
        }
        else {
            DEBUG_HEAP_PRINTF(file);
        }

        DEBUG_HEAP_PRINTF(":%d\n", line);
    }

    #define OOM_CHECK___PRINT_OOM(P, S) if(!P) print_oom_size(S)
    #define OOM_CHECK___PRINT_LOC(P, S, F, L) if(!P) print_loc(S, F, L)

#else

    // It cost 64 more bytes of IRAM to turn on (it was not previously enabled)
    #if 1
        #undef  PTR_CHECK___LOG_LAST_FAIL_FL
        #define PTR_CHECK___LOG_LAST_FAIL_FL(P, S, F, L)
        #undef  PTR_CHECK___LOG_LAST_FAIL
        #define PTR_CHECK___LOG_LAST_FAIL(P, S)
    #endif

    #define OOM_CHECK___PRINT_OOM(P, S)
    #define OOM_CHECK___PRINT_LOC(P, S, F, L)

#endif


#if defined(DEBUG_ESP_OOM) || defined(UMM_INTEGRITY_CHECK) || defined(UMM_POISON_CHECK) || defined(UMM_POISON_CHECK_LITE)

    extern void* ICACHE_RAM_ATTR malloc(size_t size)
    {
        INTEGRITY_CHECK___ABORT();
        POISON_CHECK___ABORT();

        void* ret = UMM_MALLOC(size);

        PTR_CHECK___LOG_LAST_FAIL(ret, size);
        OOM_CHECK___PRINT_OOM(ret, size);

        return ret;
    }

    extern void* ICACHE_RAM_ATTR calloc(size_t count, size_t size)
    {
        INTEGRITY_CHECK___ABORT();
        POISON_CHECK___ABORT();

        void* ret = UMM_CALLOC(count, size);

        PTR_CHECK___LOG_LAST_FAIL(ret, count * size);
        OOM_CHECK___PRINT_OOM(ret, size);

        return ret;
    }

    extern void* ICACHE_RAM_ATTR realloc(void* ptr, size_t size)
    {
        INTEGRITY_CHECK___ABORT();

        void* ret = UMM_REALLOC_FL(ptr, size, nullptr, 0);

        POISON_CHECK___ABORT();
        PTR_CHECK___LOG_LAST_FAIL(ret, size);
        OOM_CHECK___PRINT_OOM(ret, size);

        return ret;
    }

    extern void ICACHE_RAM_ATTR free(void* p)
    {
        INTEGRITY_CHECK___ABORT();

        UMM_FREE_FL(p, nullptr, 0);

        POISON_CHECK___ABORT();
    }

#endif


extern void* ICACHE_RAM_ATTR pvPortMalloc(size_t size, const char* file, int line)
{
    INTEGRITY_CHECK___PANIC_FL(file, line);
    POISON_CHECK___PANIC_FL(file, line);

    void* ret = UMM_MALLOC(size);

    PTR_CHECK___LOG_LAST_FAIL_FL(ret, size, file, line);
    OOM_CHECK___PRINT_LOC(ret, size, file, line);

    return ret;
}


extern void* ICACHE_RAM_ATTR pvPortCalloc(size_t count, size_t size, const char* file, int line)
{
    INTEGRITY_CHECK___PANIC_FL(file, line);
    POISON_CHECK___PANIC_FL(file, line);

    void* ret = UMM_CALLOC(count, size);

    PTR_CHECK___LOG_LAST_FAIL_FL(ret, count * size, file, line);
    OOM_CHECK___PRINT_LOC(ret, size, file, line);

    return ret;
}


extern void* ICACHE_RAM_ATTR pvPortRealloc(void *ptr, size_t size, const char* file, int line)
{
    INTEGRITY_CHECK___PANIC_FL(file, line);

    void* ret = UMM_REALLOC_FL(ptr, size, file, line);

    POISON_CHECK___PANIC_FL(file, line);
    PTR_CHECK___LOG_LAST_FAIL_FL(ret, size, file, line);
    OOM_CHECK___PRINT_LOC(ret, size, file, line);

    return ret;
}


extern void* ICACHE_RAM_ATTR pvPortZalloc(size_t size, const char* file, int line)
{
    INTEGRITY_CHECK___PANIC_FL(file, line);
    POISON_CHECK___PANIC_FL(file, line);

    void* ret = UMM_CALLOC(1, size);

    PTR_CHECK___LOG_LAST_FAIL_FL(ret, size, file, line);
    OOM_CHECK___PRINT_LOC(ret, size, file, line);

    return ret;
}


extern void ICACHE_RAM_ATTR vPortFree(void *ptr, const char* file, int line)
{
    INTEGRITY_CHECK___PANIC_FL(file, line);

    UMM_FREE_FL(ptr, file, line);

    POISON_CHECK___PANIC_FL(file, line);
}


extern size_t ICACHE_RAM_ATTR xPortWantedSizeAlign(size_t size)
{ return (size + 3) & ~((size_t) 3); }


extern void system_show_malloc(void)
{ umm_info(nullptr, true); }


////////////////////////////////////////////////////////////////////////////////////////////////////


#ifdef UMM_HEAP_EXTERNAL

static bool vmEnabled = false;

#endif


extern void enableVM()
{
#ifdef UMM_HEAP_EXTERNAL
    if(!vmEnabled) {
        install_vm_exception_handler();
        vmEnabled = true;
    }
#endif
}


extern void setExternalHeap()
{
#ifdef UMM_HEAP_EXTERNAL
    if(vmEnabled) {
        if( !umm_push_heap(UMM_HEAP_EXTERNAL) ) panic();
    }
#endif
}


extern void setIRAMHeap()
{
#ifdef UMM_HEAP_IRAM
    if( !umm_push_heap(UMM_HEAP_IRAM) ) panic();
#endif
}


extern void setDRAMHeap()
{
#if defined(UMM_HEAP_EXTERNAL) && !defined(UMM_HEAP_IRAM)
    if(vmEnabled) {
        if( !umm_push_heap(UMM_HEAP_DRAM) ) panic();
    }
#elif defined(UMM_HEAP_IRAM)
    if( !umm_push_heap(UMM_HEAP_DRAM) ) panic();
#endif
}


extern void resetHeap()
{
#if defined(UMM_HEAP_EXTERNAL) && !defined(UMM_HEAP_IRAM)
    if(vmEnabled) {
        if( !umm_pop_heap() ) panic();
    }
#elif defined(UMM_HEAP_IRAM)
    if( !umm_pop_heap() ) panic();
#endif
}


////////////////////////////////////////////////////////////////////////////////////////////////////


extern void* ICACHE_RAM_ATTR aligned_malloc(size_t alignment, size_t size)
{
    // Check the alignment
    if( alignment < sizeof(void*) ) alignment = sizeof(void*);

    // Calculate the offset
    const eCxx::int32_t offset = alignment - 1 + sizeof(void*);

    // Allocate memory block
    void* p1 = malloc(size + offset);
    if(!p1) return nullptr;

    // Calculate the aligned memory block address
    void** p2 = reinterpret_cast<void**>( ( reinterpret_cast<size_t>(p1) + offset ) & ~(alignment - 1) );
    p2[-1] = p1;

    // Return the aligned memory block address
    return reinterpret_cast<void*>(p2);
}


extern void ICACHE_RAM_ATTR aligned_free(void* ptr)
{
    // Free the original memory block
    free( reinterpret_cast<void**>(ptr)[-1] );
}


} // extern "C"
