/*
* Copyright (C) 2017 - 2019 Intel Corporation.
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
* 1. Redistributions of source code must retain the above copyright notice(s),
*    this list of conditions and the following disclaimer.
* 2. Redistributions in binary form must reproduce the above copyright notice(s),
*    this list of conditions and the following disclaimer in the documentation
*    and/or other materials provided with the distribution.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER(S) ``AS IS'' AND ANY EXPRESS
* OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
* MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO
* EVENT SHALL THE COPYRIGHT HOLDER(S) BE LIABLE FOR ANY DIRECT, INDIRECT,
* INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
* PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
* LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
* OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
* ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <stdio.h>
#include <numa.h>
#include <hbwmalloc.h>
#include <memkind.h>
#include <vector>
#include <memory>
#include <gtest/gtest.h>
#include "allocator_perf_tool/Allocation_info.hpp"
#include "allocator_perf_tool/GTestAdapter.hpp"

typedef std::unique_ptr<void, void(*)(void *)> hbw_mem_ptr;


class HBWPreferredLocalityTest: public ::testing::Test
{
private:
    int find_closest_node(int node, const std::vector<int> &nodes)
    {
        int min_distance = 0;
        int closest_node = -1;
        for (int i = 0; i < nodes.size(); i++) {
            int distance = numa_distance(node, nodes[i]);
            if (distance && (distance < min_distance || min_distance == 0)) {
                min_distance = distance;
                closest_node = nodes[i];
            }
        }
        return closest_node;
    }

    bool pin_to_cpu(int cpu_id)
    {
        cpu_set_t cpu_set;
        CPU_ZERO(&cpu_set);
        CPU_SET(cpu_id, &cpu_set);
        return sched_setaffinity(0, sizeof(cpu_set_t), &cpu_set) != -1;
    }

    void check_ptr_numa(void *ptr, int expected_numa_id, int cpu_id, size_t size)
    {
        memset(ptr, 1, size);

        int numa_id = get_numa_node_id(ptr);
        EXPECT_EQ(numa_id, expected_numa_id);

        char property_name[50];
        snprintf(property_name, 50, "actual_numa_for_cpu_%d_expected_numa_%d", cpu_id,
                 expected_numa_id);
        GTestAdapter::RecordProperty(property_name, numa_id);
    }

public:
    void pin_memory_in_requesting_mem_thread(size_t size,
                                             const std::vector<int> &cpu_ids,
                                             const std::vector<int> &mcdram_nodes)
    {
#ifdef _OPENMP
        int threads_num = cpu_ids.size();
        int ret = hbw_set_policy(HBW_POLICY_PREFERRED);
        ASSERT_EQ(ret, 0);

        #pragma omp parallel for num_threads(threads_num)
        for (int i = 0; i < threads_num; i++) {
            if (!pin_to_cpu(cpu_ids[i])) {
                ADD_FAILURE();
                continue;
            }

            void *internal_ptr = hbw_malloc(size);
            if (!internal_ptr) {
                ADD_FAILURE();
                continue;
            }
            hbw_mem_ptr ptr(internal_ptr, hbw_free);
            int expected_numa_id = find_closest_node(numa_node_of_cpu(cpu_ids[i]),
                                                     mcdram_nodes);
            check_ptr_numa(ptr.get(), expected_numa_id, cpu_ids[i], size);
        }
#else
        std::cout << "[ SKIPPED ] Feature OPENMP not supported" << std::endl;
#endif
    }

    void pin_memory_in_other_thread_than_requesting_mem(size_t size,
                                                        const std::vector<int> &cpu_ids,
                                                        const std::vector<int> &mcdram_nodes)
    {
#ifdef _OPENMP
        int threads_num = cpu_ids.size();
        int ret = hbw_set_policy(HBW_POLICY_PREFERRED);
        ASSERT_EQ(ret, 0);

        int main_thread_cpu_id = 0;
        int expected_numa_id = find_closest_node(main_thread_cpu_id, mcdram_nodes);
        ASSERT_TRUE(pin_to_cpu(main_thread_cpu_id));

        std::vector<hbw_mem_ptr> ptrs;
        for (int i = 0; i < threads_num; i++) {
            void *internal_ptr = hbw_malloc(size);
            ASSERT_TRUE(internal_ptr);
            ptrs.emplace_back(internal_ptr, hbw_free);
        }

        #pragma omp parallel for num_threads(threads_num)
        for (int i = 0; i < threads_num; i++) {
            if (!pin_to_cpu(cpu_ids[i])) {
                ADD_FAILURE();
                continue;
            }
            check_ptr_numa(ptrs[i].get(), expected_numa_id, cpu_ids[i], size);
        }
#else
        std::cout << "[ SKIPPED ] Feature OPENMP not supported" << std::endl;
#endif
    }
};

TEST_F(HBWPreferredLocalityTest,
       test_TC_MEMKIND_KNL_SNC4_pin_memory_in_requesting_mem_thread_4_threads_100_bytes)
{
    pin_memory_in_requesting_mem_thread(100u, std::vector<int> {0, 18, 36, 54},
                                        std::vector<int> {4, 5, 6, 7});
}

TEST_F(HBWPreferredLocalityTest,
       test_TC_MEMKIND_KNL_SNC4_pin_memory_in_other_thread_than_requesting_mem_4_threads_100_bytes)
{
    pin_memory_in_other_thread_than_requesting_mem(100u, std::vector<int> {0, 18, 36, 54},
                                                   std::vector<int> {4, 5, 6, 7});
}