// SPDX-License-Identifier: BSD-2-Clause

/* Copyright (C) 2017 - 2020 Intel Corporation. */

#pragma once

#include <stdio.h>

#include <assert.h>

#include <map>

#include <iterator>

class Allocator

{

public:

    virtual void *malloc(size_t size) = 0;

    virtual void *calloc(size_t num, size_t size) = 0;

    virtual void *realloc(void *ptr, size_t size) = 0;

    virtual int memalign(void **ptr, size_t alignment, size_t size) = 0;

    virtual void free(void *ptr) = 0;

    /* get_numa_policy() returns MPOL_INTERLEAVE, MPOL_BIND, MPOL_PREFERRED

     or -1 when the allocator is not providing NUMA policy */

    virtual int get_numa_policy() = 0;

    virtual bool is_high_bandwidth() = 0;

    virtual size_t get_page_size() = 0;

    virtual ~Allocator(void) {}

};

class MemkindAllocator : public Allocator

{

public:

    MemkindAllocator(memkind_memtype_t memtype, memkind_policy_t policy,

                     memkind_bits_t flags) :

        memtype(memtype),

        flags(flags)

    {

        int ret = memkind_create_kind(memtype, policy, flags, &kind);

        assert(ret == MEMKIND_SUCCESS);

        assert(kind != NULL);

    }

    MemkindAllocator(memkind_t kind) : kind(kind)

    {

        assert(kind != NULL);

    }

    virtual void *malloc(size_t size)

    {

        return memkind_malloc(kind, size);

    }

    virtual void *calloc(size_t num, size_t size)

    {

        return memkind_calloc(kind, num, size);

    }

    virtual void *realloc(void *ptr, size_t size)

    {

        return memkind_realloc(kind, ptr, size);

    }

    virtual int memalign(void **ptr, size_t alignment, size_t size)

    {

        return memkind_posix_memalign(kind, ptr, alignment, size);

    }

    virtual void free(void *ptr)

    {

        memkind_free(kind, ptr);

    }

    virtual int get_numa_policy()

    {

        static std::map<memkind_t, int> kind_policy = {

            std::make_pair(MEMKIND_HBW_INTERLEAVE, MPOL_INTERLEAVE),

            std::make_pair(MEMKIND_INTERLEAVE, MPOL_INTERLEAVE),

            std::make_pair(MEMKIND_HBW_PREFERRED, MPOL_PREFERRED),

            std::make_pair(MEMKIND_HBW_PREFERRED_HUGETLB, MPOL_PREFERRED),

            std::make_pair(MEMKIND_HBW, MPOL_BIND),

            std::make_pair(MEMKIND_HBW_HUGETLB, MPOL_BIND),

            std::make_pair(MEMKIND_REGULAR, MPOL_BIND),

            std::make_pair(MEMKIND_DEFAULT, MPOL_DEFAULT),

            std::make_pair(MEMKIND_HUGETLB, MPOL_DEFAULT),

            std::make_pair(MEMKIND_HBW_ALL_HUGETLB, MPOL_BIND),

            std::make_pair(MEMKIND_HBW_ALL, MPOL_BIND)

        };

        auto it = kind_policy.find(kind);

        if(it != std::end(kind_policy)) {

            return it->second;

        }

        return -1;

    }

    virtual bool is_high_bandwidth()

    {

        return  memtype == MEMKIND_MEMTYPE_HIGH_BANDWIDTH ||

                kind == MEMKIND_HBW ||

                kind == MEMKIND_HBW_HUGETLB ||

                kind == MEMKIND_HBW_PREFERRED ||

                kind == MEMKIND_HBW_INTERLEAVE;

    }

    virtual size_t get_page_size()

    {

        if (kind == MEMKIND_HUGETLB ||

            kind == MEMKIND_HBW_HUGETLB ||

            kind == MEMKIND_HBW_PREFERRED_HUGETLB ||

            (flags & MEMKIND_MASK_PAGE_SIZE_2MB)) {

            return 2*MB;

        }

        return 4*KB;

    }

    virtual ~MemkindAllocator()

    {

        int ret = memkind_destroy_kind(kind);

        assert(ret == MEMKIND_SUCCESS);

        kind = NULL;

    }

private:

    memkind_t kind = NULL;

    memkind_memtype_t memtype = memkind_memtype_t();

    memkind_bits_t flags = memkind_bits_t();

};

class HbwmallocAllocator : public Allocator

{

public:

    HbwmallocAllocator(hbw_policy_t hbw_policy)

    {

        hbw_set_policy(hbw_policy);

    }

    virtual void *malloc(size_t size)

    {

        return hbw_malloc(size);

    }

    virtual void *calloc(size_t num, size_t size)

    {

        return hbw_calloc(num, size);

    }

    virtual void *realloc(void *ptr, size_t size)

    {

        return hbw_realloc(ptr, size);

    }

    virtual void set_memalign_page_size(hbw_pagesize_t psize)

    {

        page_size = psize;

    }

    virtual int memalign(void **ptr, size_t alignment, size_t size)

    {

        if (page_size == HBW_PAGESIZE_4KB) {

            return hbw_posix_memalign(ptr, alignment, size);

        } else {

            return hbw_posix_memalign_psize(ptr, alignment, size, page_size);

        }

    }

    virtual void free(void *ptr)

    {

        hbw_free(ptr);

    }

    virtual int get_numa_policy()

    {

        if (hbw_get_policy() == HBW_POLICY_INTERLEAVE) {

            return MPOL_INTERLEAVE;

        } else if (hbw_get_policy() == HBW_POLICY_PREFERRED) {

            return MPOL_PREFERRED;

        } else if (hbw_get_policy() == HBW_POLICY_BIND ||

                   hbw_get_policy() == HBW_POLICY_BIND_ALL) {

            return MPOL_BIND;

        }

        return -1;

    }

    virtual bool is_high_bandwidth()

    {

        return  hbw_check_available() == 0;

    }

    virtual size_t get_page_size()

    {

        if (page_size == HBW_PAGESIZE_2MB) {

            return 2*MB;

        }

        return 4*KB;

    }

private:

    hbw_pagesize_t page_size = HBW_PAGESIZE_4KB;

};