// SPDX-License-Identifier: BSD-2-Clause
/* Copyright (C) 2014 - 2020 Intel Corporation. */

#pragma once

#include <memkind.h>

#include <malloc.h>
#include <cassert>
#include "jemalloc/jemalloc.h"

// Malloc, jemalloc, memkind jemalloc and memkind memory operations definitions
namespace performance_tests
{
    using std::vector;
    using std::string;
    using std::thread;

#ifdef __DEBUG
#include <mutex>
    // Write entire text at once, avoiding switching to another thread
    extern int g_msgLevel;
    extern std::mutex g_coutMutex;
#define EMIT(LEVEL, TEXT) \
    if (g_msgLevel >= LEVEL) \
    { \
        g_coutMutex.lock(); std::cout << TEXT << std::endl; g_coutMutex.unlock(); \
    }
#else
#define EMIT(LEVEL, TEXT)
#endif

    // Use jemalloc, compiled with unique prefix (--with-jemalloc-prefix= configure option)
#ifdef SYSTEM_JEMALLOC_PREFIX
#define TOKENPASTE(x, y) x ## y
#define JE(x, y) TOKENPASTE(x, y)
#define jexx_malloc JE(SYSTEM_JEMALLOC_PREFIX, malloc)
#define jexx_calloc JE(SYSTEM_JEMALLOC_PREFIX, calloc)
#define jexx_memalign JE(SYSTEM_JEMALLOC_PREFIX, memalign)
#define jexx_realloc JE(SYSTEM_JEMALLOC_PREFIX, realloc)
#define jexx_free JE(SYSTEM_JEMALLOC_PREFIX, free)
    extern "C" {
        // jemalloc function prototypes
        // full header cannot be include due to conflict with memkind jemalloc
        extern void *jexx_malloc(size_t size);
        extern void *jexx_calloc(size_t num, size_t size);
        extern void *jexx_memalign(size_t alignment, size_t size);
        extern void *jexx_realloc(void *ptr, size_t size);
        extern void  jexx_free(void *ptr);
    }
#endif

    // Available memory operations
    enum OperationName {
        Malloc,
        Calloc,
        Realloc,
        Align,
        Free,
        Invalid
    };

    // Reprents a memory operation
    class Operation
    {
    public:
        // Each operation is assigned a bucket size from range (0, MaxBucketSize)
        static const unsigned MaxBucketSize = 100;
        // For memalign operation, alignment parameter will be a random value
        // from range (sizeof(void*), sizeof(void*) * MemalignMaxMultiplier)
        static const unsigned MemalignMaxMultiplier = 4;

    protected:
        OperationName m_name;
        // If random number from range (0, MaxBucketSize) is lower than m_bucketSize, an operation will be performed
        unsigned m_bucketSize;

    public:
        Operation()
        {
        }

        // Creates an operation with given name and bucket size
        // If no bucket size is given, operation will be always performed
        Operation(
            OperationName name,
            unsigned bucketSize = MaxBucketSize)
            : m_name(name)
            , m_bucketSize(bucketSize)
        {
            assert(bucketSize <= MaxBucketSize);
        }

        virtual ~Operation() {}
        ;

        // Check if operation should be performed (currently drawn random number lower than bucket size)
        bool checkCondition(unsigned ballSize) const
        {
            return (ballSize < m_bucketSize);
        }

        OperationName getName() const
        {
            return m_name;
        }

        string getNameStr() const
        {
            switch (m_name) {
                case OperationName::Malloc:
                    return "malloc";

                case OperationName::Calloc:
                    return "calloc";

                case OperationName::Realloc:
                    return "realloc";

                case OperationName::Align:
                    return "align";

                case OperationName::Free:
                    return "free";

                default:
                    return "<unknown>";
            }
        }

        // Get operation bucket size
        unsigned getBucketSize() const
        {
            return m_bucketSize;
        }

        // perform memory operation
        virtual void perform(const memkind_t &kind,
                             void *&mem,
                             size_t size = 0,
                             size_t offset=0,
                             size_t alignment=0)
        const = 0;

    };

    // Malloc memory operations
    class MallocOperation : public Operation
    {
    public:
        MallocOperation(
            OperationName name)
            : Operation(name)
        {
        }

        MallocOperation(
            OperationName name,
            unsigned bucketSize)
            : Operation(name, bucketSize)
        {
        }

        virtual void perform(const memkind_t &kind,
                             void *&mem,
                             size_t size,
                             size_t offset,
                             size_t alignment) const override
        {
            EMIT(2, "Entering Operation::" << getNameStr()
                 <<  ", size=" << size
                 << ", offset=" << offset
                 << ", alignment=" << alignment
                 << ", mem=" << mem)
            switch(m_name) {
                case Malloc: {
                    if (mem != nullptr) {
                        free(mem);
                    }
                    mem = malloc(size);
                    break;
                }
                case Calloc: {
                    if (mem != nullptr) {
                        free(mem);
                    }
                    // split allocation size randomly
                    // between number of elements and element size
                    mem = calloc((1 << offset), (size >> offset));
                    break;
                }
                case Realloc: {
                    mem = realloc(mem, size);
                    break;
                }
                case Align: {
                    if (mem != nullptr) {
                        free(mem);
                    }
                    // randomly choose alignment from (8, 8 * MemalignMaxMultiplie)
                    mem = memalign(alignment, size);
                    break;
                }
                case Free: {
                    free(mem);
                    mem = nullptr;
                    break;
                }

                default:
                    throw "Not implemented";
                    break;
            }
            EMIT(2, "Exiting Operation::" << getNameStr()
                 <<  ", size=" << size
                 << ", offset=" << offset
                 << ", alignment=" << alignment
                 << ", mem=" << mem)
        }
    };

#ifdef SYSTEM_JEMALLOC_PREFIX
    // Jemalloc memory operations
    class JemallocOperation : public Operation
    {
    public:
        JemallocOperation(
            OperationName name)
            : Operation(name)
        {
        }

        JemallocOperation(
            OperationName name,
            unsigned bucketSize)
            : Operation(name, bucketSize)
        {
        }

        virtual void perform(const memkind_t &kind,
                             void *&mem,
                             size_t size,
                             size_t offset,
                             size_t alignment) const override
        {
            EMIT(2, "Entering Operation::" << getNameStr()
                 <<  ", size=" << size
                 << ", offset=" << offset
                 << ", alignment=" << alignment
                 << ", mem=" << mem)
            switch(m_name) {
                case Malloc: {
                    if (mem != nullptr) {
                        jexx_free(mem);
                    }
                    mem = jexx_malloc(size);
                    break;
                }
                case Calloc: {
                    if (mem != nullptr) {
                        jexx_free(mem);
                    }
                    // split allocation size randomly
                    // between number of elements and element size
                    mem = jexx_calloc((1 <<, offset), (size >>, offset));
                    break;
                }
                case Realloc: {
                    mem = jexx_realloc(mem, size);
                    break;
                }
                case Align: {
                    if (mem != nullptr) {
                        jexx_free(mem);
                    }
                    // randomly choose alignment from (8, 8 * MemalignMaxMultiplie)
                    mem = jexx_memalign(alignment, size);
                    break;
                }
                case Free: {
                    jexx_free(mem);
                    mem = nullptr;
                    break;
                }

                default:
                    throw "Not implemented";
                    break;
            }
            EMIT(2, "Exiting Operation::" << getNameStr()
                 <<  ", size=" << size
                 << ", offset=" << offset
                 << ", alignment=" << alignment
                 << ", mem=" << mem)
        }
    };
#endif
    // Jemkmalloc memory operations
    class JemkmallocOperation : public Operation
    {
    public:
        JemkmallocOperation(
            OperationName name)
            : Operation(name)
        {
        }

        JemkmallocOperation(
            OperationName name,
            unsigned bucketSize)
            : Operation(name, bucketSize)
        {
        }

        virtual void perform(const memkind_t &kind,
                             void *&mem,
                             size_t size,
                             size_t offset,
                             size_t alignment) const override
        {
#ifdef JEMK
            EMIT(2, "Entering Operation::" << getNameStr()
                 <<  ", size=" << size
                 << ", offset=" << offset
                 << ", alignment=" << alignment
                 << ", mem=" << mem)
            switch(m_name) {
                case Malloc: {
                    if (mem != nullptr) {
                        jemk_free(mem);
                    }
                    mem = jemk_malloc(size);
                    break;
                }
                case Calloc: {
                    if (mem != nullptr) {
                        jemk_free(mem);
                    }
                    // split allocation size randomly
                    // between number of elements and element size
                    mem = jemk_calloc((1 << offset), (size >> offset));
                    break;
                }
                case Realloc: {
                    mem = jemk_realloc(mem, size);
                    break;
                }
                case Align: {
                    if (mem != nullptr) {
                        jemk_free(mem);
                    }
                    // randomly choose alignment from (8, 8 * MemalignMaxMultiplie)
                    mem = jemk_memalign(alignment, size);
                    break;
                }
                case Free: {
                    jemk_free(mem);
                    mem = nullptr;
                    break;
                }

                default:
                    throw "Not implemented";
                    break;
            }
            EMIT(2, "Exiting Operation::" << getNameStr()
                 <<  ", size=" << size
                 << ", offset=" << offset
                 << ", alignment=" << alignment
                 << ", mem=" << mem)
#endif // JE_MK
        }
    };

    // Memkind memory operations
    class MemkindOperation : public Operation
    {
    public:
        MemkindOperation()
        {}

        MemkindOperation(
            OperationName name)
            : Operation(name)
        {
        }

        MemkindOperation(
            OperationName name,
            size_t bucketSize)
            : Operation(name, bucketSize)
        {
        }

        virtual void perform(const memkind_t &kind,
                             void *&mem,
                             size_t size,
                             size_t offset,
                             size_t alignment) const override
        {
            EMIT(2, "Entering Operation::" << getNameStr()
                 <<  ", size=" << size
                 << ", offset=" << offset
                 << ", alignment=" << alignment
                 << ", mem=" << mem)
            switch (m_name) {
                case Malloc: {
                    if (mem != nullptr) {
                        memkind_free(kind, mem);
                    }
                    mem = memkind_malloc(kind, size);
                    break;
                }
                case Calloc: {
                    if (mem != nullptr) {
                        memkind_free(kind, mem);
                    }
                    // split allocation size randomly
                    // between number of elements and element size
                    mem = memkind_calloc(kind, 1 << offset, size >> offset);
                    break;
                }
                case Realloc: {
                    mem = memkind_realloc(kind, mem, size);
                    break;
                }

                case Align: {
                    if (mem != nullptr) {
                        memkind_free(kind, mem);
                        mem = nullptr;
                    }
                    // randomly choose alignment from (sizeof(void*), sizeof(void*) * MemalignMaxMultiplie)
                    if (memkind_posix_memalign(kind, &mem, alignment, size) != 0) {
                        // failure
                        mem = nullptr;
                    }
                    break;
                }
                case Free: {
                    memkind_free(kind, mem);
                    mem = nullptr;
                    break;
                }

                default:
                    break;
            }
            EMIT(2, "Exiting Operation::" << getNameStr()
                 <<  ", size=" << size
                 << ", offset=" << offset
                 << ", alignment=" << alignment
                 << ", mem=" << mem)
        }
    };
}
