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

#include "allocator_perf_tool/HugePageOrganizer.hpp"
#include <memkind/internal/memkind_hbw.h>

#include "check.h"
#include "trial_generator.h"
#include <numa.h>
#include <numaif.h>
#include <unistd.h>
#include <vector>

trial_t TrialGenerator ::create_trial_tuple(alloc_api_t api, size_t size,
                                            size_t alignment, int page_size,
                                            memkind_t memkind, int free_index)
{
    trial_t ltrial;
    ltrial.api = api;
    ltrial.size = size;
    ltrial.alignment = alignment;
    ltrial.page_size = page_size;
    ltrial.memkind = memkind;
    ltrial.free_index = free_index;
    return ltrial;
}

void TrialGenerator ::generate_gb(alloc_api_t api, int number_of_gb_pages,
                                  memkind_t memkind, alloc_api_t api_free,
                                  bool psize_strict, size_t align)
{
    std::vector<size_t> sizes_to_alloc;
    // When API = HBW_MEMALIGN_PSIZE: psize is set to HBW_PAGESIZE_1GB_STRICT
    // when allocation is a multiple of 1GB. Otherwise it is set to
    // HBW_PAGESIZE_1GB.
    for (int i = 1; i <= number_of_gb_pages; i++) {
        if (psize_strict || api != HBW_MEMALIGN_PSIZE)
            sizes_to_alloc.push_back(i * GB);
        else
            sizes_to_alloc.push_back(i * GB + 1);
    }
    int k = 0;
    trial_vec.clear();

    for (int i = 0; i < (int)sizes_to_alloc.size(); i++) {
        trial_vec.push_back(create_trial_tuple(api, sizes_to_alloc[i], align,
                                               2 * MB, memkind, -1));
        if (i > 0)
            k++;
        trial_vec.push_back(
            create_trial_tuple(api_free, 0, 0, 0, memkind, k++));
    }
}

int n_random(int i)
{
    return random() % i;
}

void TrialGenerator ::generate_size_2bytes_2KB_2MB(alloc_api_t api)
{
    size_t size[] = {2, 2 * KB, 2 * MB};

    int k = 0;
    trial_vec.clear();
    for (unsigned int i = 0; i < (int)(sizeof(size) / sizeof(size[0])); i++) {
        trial_vec.push_back(create_trial_tuple(
            api, size[i], 32, sysconf(_SC_PAGESIZE), MEMKIND_HBW, -1));

        if (i > 0)
            k++;
        trial_vec.push_back(
            create_trial_tuple(HBW_FREE, 0, 0, 0, MEMKIND_HBW, k));
        k++;
    }
}

void TrialGenerator ::print()
{

    std::vector<trial_t>::iterator it;

    std::cout << "*********** Size: " << trial_vec.size() << "********\n";
    std::cout << "SIZE PSIZE ALIGN FREE KIND" << std::endl;

    for (it = trial_vec.begin(); it != trial_vec.end(); it++) {
        std::cout << it->size << " " << it->page_size << " " << it->alignment
                  << " " << it->free_index << " " << it->memkind << " "
                  << std::endl;
    }
}

void TrialGenerator ::run(int num_bandwidth, std::vector<int> &bandwidth)
{

    int num_trial = trial_vec.size();
    int i, ret = 0;
    void **ptr_vec = NULL;

    ptr_vec = (void **)malloc(num_trial * sizeof(void *));
    if (NULL == ptr_vec) {
        fprintf(stderr, "Error in allocating ptr array\n");
        exit(-1);
    }

    for (i = 0; i < num_trial; ++i) {
        ptr_vec[i] = NULL;
    }
    for (i = 0; i < num_trial; ++i) {
        switch (trial_vec[i].api) {
            case HBW_FREE:
                if (i == num_trial - 1 || trial_vec[i + 1].api != HBW_REALLOC) {
                    hbw_free(ptr_vec[trial_vec[i].free_index]);
                    ptr_vec[trial_vec[i].free_index] = NULL;
                    ptr_vec[i] = NULL;
                } else {
                    ptr_vec[i + 1] =
                        hbw_realloc(ptr_vec[trial_vec[i].free_index],
                                    trial_vec[i + 1].size);
                    ptr_vec[trial_vec[i].free_index] = NULL;
                }
                break;
            case MEMKIND_FREE:
                if (i == num_trial - 1 ||
                    trial_vec[i + 1].api != MEMKIND_REALLOC) {
                    memkind_free(trial_vec[i].memkind,
                                 ptr_vec[trial_vec[i].free_index]);
                    ptr_vec[trial_vec[i].free_index] = NULL;
                    ptr_vec[i] = NULL;
                } else {
                    ptr_vec[i + 1] = memkind_realloc(
                        trial_vec[i].memkind, ptr_vec[trial_vec[i].free_index],
                        trial_vec[i + 1].size);
                    ptr_vec[trial_vec[i].free_index] = NULL;
                }
                break;
            case HBW_MALLOC:
                fprintf(stdout, "Allocating %zd bytes using hbw_malloc\n",
                        trial_vec[i].size);
                ptr_vec[i] = hbw_malloc(trial_vec[i].size);
                break;
            case HBW_CALLOC:
                fprintf(stdout, "Allocating %zd bytes using hbw_calloc\n",
                        trial_vec[i].size);
                ptr_vec[i] = hbw_calloc(trial_vec[i].size, 1);
                break;
            case HBW_REALLOC:
                fprintf(stdout, "Allocating %zd bytes using hbw_realloc\n",
                        trial_vec[i].size);
                fflush(stdout);
                if (NULL == ptr_vec[i]) {
                    ptr_vec[i] = hbw_realloc(NULL, trial_vec[i].size);
                }
                break;
            case HBW_MEMALIGN:
                fprintf(stdout, "Allocating %zd bytes using hbw_memalign\n",
                        trial_vec[i].size);
                ret = hbw_posix_memalign(&ptr_vec[i], trial_vec[i].alignment,
                                         trial_vec[i].size);
                break;
            case HBW_MEMALIGN_PSIZE:
                fprintf(stdout,
                        "Allocating %zd bytes using hbw_memalign_psize\n",
                        trial_vec[i].size);
                hbw_pagesize_t psize;
                if (trial_vec[i].page_size == (size_t)sysconf(_SC_PAGESIZE))
                    psize = HBW_PAGESIZE_4KB;
                else if (trial_vec[i].page_size == 2097152)
                    psize = HBW_PAGESIZE_2MB;
                else if (trial_vec[i].size % trial_vec[i].page_size > 0)
                    psize = HBW_PAGESIZE_1GB;
                else
                    psize = HBW_PAGESIZE_1GB_STRICT;

                ret = hbw_posix_memalign_psize(&ptr_vec[i],
                                               trial_vec[i].alignment,
                                               trial_vec[i].size, psize);

                break;
            case MEMKIND_MALLOC:
                fprintf(stdout, "Allocating %zd bytes using memkind_malloc\n",
                        trial_vec[i].size);
                ptr_vec[i] =
                    memkind_malloc(trial_vec[i].memkind, trial_vec[i].size);
                break;
            case MEMKIND_CALLOC:
                fprintf(stdout, "Allocating %zd bytes using memkind_calloc\n",
                        trial_vec[i].size);
                ptr_vec[i] =
                    memkind_calloc(trial_vec[i].memkind, trial_vec[i].size, 1);
                break;
            case MEMKIND_REALLOC:
                fprintf(stdout, "Allocating %zd bytes using memkind_realloc\n",
                        trial_vec[i].size);
                if (NULL == ptr_vec[i]) {
                    ptr_vec[i] = memkind_realloc(trial_vec[i].memkind,
                                                 ptr_vec[i], trial_vec[i].size);
                }
                break;
            case MEMKIND_POSIX_MEMALIGN:
                fprintf(stdout,
                        "Allocating %zd bytes using memkind_posix_memalign\n",
                        trial_vec[i].size);

                ret = memkind_posix_memalign(trial_vec[i].memkind, &ptr_vec[i],
                                             trial_vec[i].alignment,
                                             trial_vec[i].size);
                break;
        }
        if (trial_vec[i].api != HBW_FREE && trial_vec[i].api != MEMKIND_FREE &&
            trial_vec[i].memkind != MEMKIND_DEFAULT) {
            ASSERT_TRUE(ptr_vec[i] != NULL);
            memset(ptr_vec[i], 0, trial_vec[i].size);
            Check check(ptr_vec[i], trial_vec[i]);
            if (trial_vec[i].api == HBW_CALLOC) {
                EXPECT_EQ(0, check.check_zero());
            }
            if (trial_vec[i].api == HBW_MEMALIGN ||
                trial_vec[i].api == HBW_MEMALIGN_PSIZE ||
                trial_vec[i].api == MEMKIND_POSIX_MEMALIGN) {
                EXPECT_EQ(0, check.check_align(trial_vec[i].alignment));
                EXPECT_EQ(0, ret);
            }
            if (trial_vec[i].api == HBW_MEMALIGN_PSIZE ||
                (trial_vec[i].api == MEMKIND_MALLOC &&
                 (trial_vec[i].memkind == MEMKIND_HBW_HUGETLB ||
                  trial_vec[i].memkind == MEMKIND_HBW_PREFERRED_HUGETLB))) {
                EXPECT_EQ(0, check.check_page_size(trial_vec[i].page_size));
            }
        }
    }
    for (i = 0; i < num_trial; ++i) {
        if (ptr_vec[i]) {
            hbw_free(ptr_vec[i]);
        }
    }
}

void TGTest ::SetUp()
{
    size_t node;
    char *hbw_nodes_env, *endptr;
    tgen = std::unique_ptr<TrialGenerator>(new TrialGenerator());

    hbw_nodes_env = getenv("MEMKIND_HBW_NODES");
    if (hbw_nodes_env) {
        num_bandwidth = 128;
        for (node = 0; node < num_bandwidth; node++) {
            bandwidth.push_back(1);
        }
        node = strtol(hbw_nodes_env, &endptr, 10);
        bandwidth.push_back(2);
        while (*endptr == ':') {
            hbw_nodes_env = endptr + 1;
            node = strtol(hbw_nodes_env, &endptr, 10);
            if (endptr != hbw_nodes_env && node < num_bandwidth) {
                bandwidth.push_back(2);
            }
        }
    } else {
        num_bandwidth = NUMA_NUM_NODES;
        nodemask_t nodemask;
        struct bitmask nodemask_bm = {NUMA_NUM_NODES, nodemask.n};
        numa_bitmask_clearall(&nodemask_bm);

        memkind_hbw_all_get_mbind_nodemask(NULL, nodemask.n, NUMA_NUM_NODES);

        int i, max_node = numa_max_node();
        for (i = 0; i < NUMA_NUM_NODES; i++) {
            if (i > max_node) {
                bandwidth.push_back(0);
            } else if (numa_bitmask_isbitset(&nodemask_bm, i)) {
                bandwidth.push_back(2);
            } else {
                bandwidth.push_back(1);
            }
        }
    }
}

void TGTest ::TearDown()
{}
