// SPDX-License-Identifier: BSD-2-Clause
/* Copyright (C) 2016 - 2021 Intel Corporation. */
#include "allocator_perf_tool/HugePageOrganizer.hpp"
#include "common.h"
#include "hbwmalloc.h"
#include "memkind.h"
#include "memkind/internal/memkind_hbw.h"

#include <numa.h>
#include <numaif.h>
#include <stdint.h>
#include <sys/mman.h>

#define SHIFT_BYTES(ptr, bytes) ((char *)(ptr) + (bytes))

/*
 * Set of tests for hbw_verify_memory_region() function,
 * which intend to check if allocated memory fully fall into high bandwidth
 * memory. Note: in this tests we are using internal function
 * memkind_hbw_all_get_mbind_nodemask(). In future we intend to rewrite this
 * function when we will have replacement in new API.
 */
class HbwVerifyFunctionTest: public ::testing::Test
{
protected:
    const size_t BLOCK_SIZE = 64;
    const size_t page_size = sysconf(_SC_PAGESIZE);
    const int flags = MAP_ANONYMOUS | MAP_PRIVATE;
};

/*
 * Group of basic tests for different sizes of allocations
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_HBW_page_size_not_round)
{
    size_t size = page_size * 1024 + 5;
    char *ptr = (char *)hbw_malloc(size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), 0);
    hbw_free(ptr);
}

TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_HBW_page_size_round)
{
    size_t size = page_size * 1024;
    char *ptr = (char *)hbw_malloc(size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), 0);
    hbw_free(ptr);
}

TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_HBW_iterate_1_byte_to_8194_bytes)
{
    for (size_t size = 1; size <= (page_size * 2 + 2);
         size++) { // iterate through 2 pages and 2 bytes
        char *ptr = (char *)hbw_malloc(size);
        ASSERT_FALSE(ptr == NULL);
        EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), 0);
        hbw_free(ptr);
    }
}

TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_HBW_ext_5GB)
{
    size_t size = 5ull * (1 << 30); // 5GB - big allocation
    char *ptr = (char *)hbw_malloc(size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), 0);
    hbw_free(ptr);
}

/*
 * Test setting memory without HBW_TOUCH_PAGES flag
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_HBW_setting_memory_without_flag)
{
    size_t size = page_size;
    char *ptr = (char *)hbw_malloc(size);
    ASSERT_FALSE(ptr == NULL);
    memset(ptr, '.', size);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, 0), 0);
    for (size_t i = 0; i < size; i++) { // check that content is unchanged
        EXPECT_TRUE(ptr[i] == '.');
    }
    hbw_free(ptr);
}

/*
 * Test HBW_TOUCH_PAGES flag
 */
TEST_F(HbwVerifyFunctionTest,
       test_TC_MEMKIND_HBW_TOUCH_PAGES_check_overwritten_content)
{
    size_t size = 5 * page_size;
    char *ptr = (char *)hbw_malloc(size);
    ASSERT_FALSE(ptr == NULL);
    memset(ptr, '.', size);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), 0);
    for (size_t i = 0; i < size;
         i++) { // check that content of all pages doesn't change
        EXPECT_TRUE(ptr[i] == '.');
    }
    hbw_free(ptr);
}

/*
 * Tests check if number of 64-page blocks are working correctly
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_HBW_many_blocks_round)
{
    size_t size = 16 * (BLOCK_SIZE * page_size); // exactly 16 * 64-page blocks
    char *ptr = (char *)hbw_malloc(size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), 0);
    hbw_free(ptr);
}

TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_HBW_many_blocks_not_round)
{
    size_t size = (16 * BLOCK_SIZE * page_size) +
        (8 * page_size); // 16 * 64-page blocks and 1 * 8-paged block
    char *ptr = (char *)hbw_malloc(size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), 0);
    hbw_free(ptr);
}

/*
 * Check if 2 blocks and parts of the third block are working correctly.
 * With this test we can be sure that different sizes of allocation (not only
 * whole blocks) are ok.
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_HBW_iterate_not_round)
{
    for (size_t i = 1; i < BLOCK_SIZE; i++) {
        size_t size = (2 * BLOCK_SIZE * page_size) +
            (i * page_size); // 2 * 64-paged blocks and iterate through whole
                             // 3rd 64-paged block
        char *ptr = (char *)hbw_malloc(size);
        ASSERT_FALSE(ptr == NULL);
        EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), 0);
        hbw_free(ptr);
    }
}

/*
 * Tests for other kinds and malloc
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_2MBPages_HBW_HUGETLB)
{
    HugePageOrganizer huge_page_organizer(16);
    size_t size = 2 * 1024 * 1024 * 10; // 10 * 2MB pages
    char *ptr = (char *)memkind_malloc(MEMKIND_HBW_HUGETLB, size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), 0);
    memkind_free(MEMKIND_HBW_HUGETLB, ptr);
}

TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_DEFAULT)
{
    size_t size = page_size * 1024;
    char *ptr = (char *)memkind_malloc(MEMKIND_DEFAULT, size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), -1);
    memkind_free(MEMKIND_DEFAULT, ptr);
}

TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_malloc)
{
    size_t size = page_size * 1024;
    char *ptr = (char *)malloc(size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, HBW_TOUCH_PAGES), -1);
    free(ptr);
}

/*
 * Group of negative tests
 */
TEST_F(HbwVerifyFunctionTest,
       test_TC_MEMKIND_Negative_size_0_and_SET_MEMORY_flag)
{
    void *ptr = hbw_malloc(page_size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, 0, HBW_TOUCH_PAGES), EINVAL);
    hbw_free(ptr);
}

TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_Negative_size_0_without_flag)
{
    void *ptr = hbw_malloc(page_size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, 0, 0), EINVAL);
    hbw_free(ptr);
}

TEST_F(HbwVerifyFunctionTest,
       test_TC_MEMKIND_Negative_Uninitialized_Memory_without_flag)
{
    void *ptr = NULL;
    EXPECT_EQ(hbw_verify_memory_region(ptr, page_size * 1024, 0), EINVAL);
}

TEST_F(HbwVerifyFunctionTest,
       test_TC_MEMKIND_Negative_Uninitialized_Memory_and_SET_MEMORY_flag)
{
    void *ptr = NULL;
    EXPECT_EQ(hbw_verify_memory_region(ptr, page_size * 1024, HBW_TOUCH_PAGES),
              EINVAL);
}

TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_Negative_Without_Memset)
{
    size_t size = page_size * 1024;
    void *ptr = hbw_malloc(size);
    ASSERT_FALSE(ptr == NULL);
    EXPECT_EQ(hbw_verify_memory_region(ptr, size, 0), -1);
    hbw_free(ptr);
}

/*
 * Corner cases: tests for half of pages
 * + HBM memory
 * # HBM memory and verified
 * - not HBM memory, but allocated
 * = not HBM memory and verified
 */

/* 3 pages:
 * |++##|####|##++|
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_Half_Pages)
{
    size_t size = 3 * page_size;
    nodemask_t nodemask;
    struct bitmask hbw_nodemask = {NUMA_NUM_NODES, nodemask.n};
    // function memkind_hbw_all_get_mbind_nodemask() has to be rewritten, when
    // we will have replacement in new API.
    memkind_hbw_all_get_mbind_nodemask(NULL, hbw_nodemask.maskp,
                                       hbw_nodemask.size);
    char *ptr = (char *)mmap(NULL, size, PROT_READ | PROT_WRITE, flags, -1, 0);
    ASSERT_FALSE(ptr == NULL);

    // all pages should fall on HBM
    mbind(ptr, size, MPOL_BIND, hbw_nodemask.maskp, NUMA_NUM_NODES, 0);

    // verified are: half of the first page, second page and half of the third
    // page
    EXPECT_EQ(hbw_verify_memory_region(SHIFT_BYTES(ptr, page_size / 2),
                                       size - page_size, HBW_TOUCH_PAGES),
              0);
    EXPECT_EQ(munmap(ptr, size), 0);
}

/* 3 pages:
 * |####|####|----|
 * |++##|####|==--|
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_Half_Pages_1_and_2_page)
{
    size_t size = 3 * page_size;
    nodemask_t nodemask;
    struct bitmask hbw_nodemask = {NUMA_NUM_NODES, nodemask.n};
    // function memkind_hbw_all_get_mbind_nodemask() has to be rewritten, when
    // we will have replacement in new API.
    memkind_hbw_all_get_mbind_nodemask(NULL, hbw_nodemask.maskp,
                                       hbw_nodemask.size);
    char *ptr = (char *)mmap(NULL, size, PROT_READ | PROT_WRITE, flags, -1, 0);
    ASSERT_FALSE(ptr == NULL);

    // first and second page should fall on HBM
    mbind(ptr, size - page_size, MPOL_BIND, hbw_nodemask.maskp, NUMA_NUM_NODES,
          0);

    // check if mbind is successful |####|####|----|
    EXPECT_EQ(hbw_verify_memory_region(ptr, size - page_size, HBW_TOUCH_PAGES),
              0);

    // verified are: half of the first page, second page and half of the third
    // page |++##|####|==--|
    EXPECT_EQ(hbw_verify_memory_region(SHIFT_BYTES(ptr, page_size / 2),
                                       size - page_size, HBW_TOUCH_PAGES),
              -1);
    EXPECT_EQ(munmap(ptr, size), 0);
}

/* 3 pages:
 * |----|####|####|
 * |--==|####|##++|
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_Half_Pages_2_and_3_page)
{
    size_t size = 3 * page_size;
    nodemask_t nodemask;
    struct bitmask hbw_nodemask = {NUMA_NUM_NODES, nodemask.n};
    // function memkind_hbw_all_get_mbind_nodemask() has to be rewritten, when
    // we will have replacement in new API.
    memkind_hbw_all_get_mbind_nodemask(NULL, hbw_nodemask.maskp,
                                       hbw_nodemask.size);
    char *ptr = (char *)mmap(NULL, size, PROT_READ | PROT_WRITE, flags, -1, 0);
    ASSERT_FALSE(ptr == NULL);

    // second and third page should fall on HBM
    mbind(SHIFT_BYTES(ptr, page_size), size - page_size, MPOL_BIND,
          hbw_nodemask.maskp, NUMA_NUM_NODES, 0);

    // check if mbind is successful |----|####|####|
    EXPECT_EQ(hbw_verify_memory_region(SHIFT_BYTES(ptr, page_size),
                                       size - page_size, HBW_TOUCH_PAGES),
              0);

    // verified are: half of the first page, second page and half of the third
    // page
    EXPECT_EQ(hbw_verify_memory_region(SHIFT_BYTES(ptr, page_size / 2),
                                       size - page_size, HBW_TOUCH_PAGES),
              -1);
    EXPECT_EQ(munmap(ptr, size), 0);
}

/* 3 pages:
 * |####|----|++++| and  |++++|----|####|
 * |++##|====|##++|
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_Half_Pages_1_and_3_page)
{
    size_t size = 3 * page_size;
    nodemask_t nodemask;
    struct bitmask hbw_nodemask = {NUMA_NUM_NODES, nodemask.n};
    // function memkind_hbw_all_get_mbind_nodemask() has to be rewritten, when
    // we will have replacement in new API.
    memkind_hbw_all_get_mbind_nodemask(NULL, hbw_nodemask.maskp,
                                       hbw_nodemask.size);
    char *ptr = (char *)mmap(NULL, size, PROT_READ | PROT_WRITE, flags, -1, 0);
    ASSERT_FALSE(ptr == NULL);

    // first and third page should fall on HBM
    mbind(ptr, page_size, MPOL_BIND, hbw_nodemask.maskp, NUMA_NUM_NODES, 0);
    mbind(SHIFT_BYTES(ptr, 2 * page_size), page_size, MPOL_BIND,
          hbw_nodemask.maskp, NUMA_NUM_NODES, 0);

    // check if mbind is successful |####|----|++++| and  |++++|----|####|
    EXPECT_EQ(
        hbw_verify_memory_region(ptr, size - 2 * page_size, HBW_TOUCH_PAGES),
        0);
    EXPECT_EQ(hbw_verify_memory_region(SHIFT_BYTES(ptr, 2 * page_size),
                                       size - 2 * page_size, HBW_TOUCH_PAGES),
              0);

    // verified are: half of the first page, second page and half of the third
    // page
    EXPECT_EQ(hbw_verify_memory_region(SHIFT_BYTES(ptr, page_size / 2),
                                       size - page_size, HBW_TOUCH_PAGES),
              -1);
    EXPECT_EQ(munmap(ptr, size), 0);
}

/* 5 pages:
 * |----|+###|####|###+|----|
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_Boundaries_CornerCase)
{
    size_t size = 5 * page_size;
    nodemask_t nodemask;
    struct bitmask hbw_nodemask = {NUMA_NUM_NODES, nodemask.n};
    // function memkind_hbw_all_get_mbind_nodemask() has to be rewritten, when
    // we will have replacement in new API.
    memkind_hbw_all_get_mbind_nodemask(NULL, hbw_nodemask.maskp,
                                       hbw_nodemask.size);
    char *ptr = (char *)mmap(NULL, size, PROT_READ | PROT_WRITE, flags, -1, 0);
    ASSERT_FALSE(ptr == NULL);

    size_t tested_pages_size = size - 2 * page_size;

    // second, third and fourth page should fall on HBM
    mbind(SHIFT_BYTES(ptr, page_size), tested_pages_size, MPOL_BIND,
          hbw_nodemask.maskp, NUMA_NUM_NODES, 0);

    // verified are: second byte of the second page to last byte -1 of fourth
    // page
    EXPECT_EQ(hbw_verify_memory_region(SHIFT_BYTES(ptr, page_size + 1),
                                       tested_pages_size - 2, HBW_TOUCH_PAGES),
              0);
    EXPECT_EQ(munmap(ptr, size), 0);
}

/* 5 pages:
 * |-###|####|####|####|###-|
 */
TEST_F(HbwVerifyFunctionTest, test_TC_MEMKIND_HBW_partial_verification)
{
    size_t size = 5 * page_size;

    // all pages should fall on HBM
    void *ptr = hbw_malloc(size);
    ASSERT_FALSE(ptr == NULL);

    // but only second byte of the first page to last byte -1 of the last page
    // are touched
    memset(SHIFT_BYTES(ptr, 1), 0, size - 2);

    // verified are: second byte of the first page to last byte -1 of last page
    EXPECT_EQ(hbw_verify_memory_region(SHIFT_BYTES(ptr, 1), size - 2, 0), 0);
    hbw_free(ptr);
}