//
// Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include <backendsCommon/MemoryManager.hpp>
#include <armnn/utility/IgnoreUnused.hpp>

#include <doctest/doctest.h>
#include <numeric>

namespace armnn
{
/// @brief Class that implements a sample custom allocator.
class SampleCustomAllocator : public armnn::ICustomAllocator
{
public:
    SampleCustomAllocator() = default;

    void* allocate(size_t size, size_t alignment) override
    {
        IgnoreUnused(alignment);
        CHECK(size == m_Values.size());
        m_CounterAllocate+=1;
        return m_Values.data();
    }

    void free(void* ptr) override
    {
        CHECK(ptr == m_Values.data());
        m_CounterFree+=1;
    }

    armnn::MemorySource GetMemorySourceType() override
    {
        return armnn::MemorySource::Malloc;
    }

    virtual void* GetMemoryRegionAtOffset(void* buffer, size_t offset, size_t alignment = 0 ) override
    {
        IgnoreUnused(alignment);
        return (static_cast<char*>(buffer) + offset);
    }

    /// Holds the data in the tensors. Create for testing purposes.
    std::vector<uint8_t> m_Values;
    /// Counts the number of times the function allocate is called.
    unsigned long m_CounterAllocate= 0;
    /// Counts the number of times the function free is called.
    unsigned long m_CounterFree = 0;
};

TEST_SUITE("MemoryManagerTests")
{
/// Unit test Storing, Allocating and Deallocating with a custom allocator.
TEST_CASE("MemoryManagerTest")
{
    using namespace armnn;

    // Create mock up bufferStorageVector with 2 BufferStorage with the same TensorMemory
    size_t numTensors = 5;
    std::vector<TensorMemory*> tensorMemoryPointerVector(numTensors);
    std::vector<TensorMemory> tensorMemoryVector;
    tensorMemoryVector.reserve(numTensors);

    std::vector<size_t> offsets(numTensors);
    std::iota(std::begin(offsets), std::end(offsets), 0);

    for (uint32_t idx = 0; idx < tensorMemoryPointerVector.size(); ++idx)
    {
        tensorMemoryVector.emplace_back(TensorMemory{offsets[idx], nullptr, 0});
        tensorMemoryPointerVector[idx] = &tensorMemoryVector[idx];
    }

    std::vector<BufferStorage> bufferStorageVector;
    bufferStorageVector.emplace_back(BufferStorage{tensorMemoryPointerVector, numTensors});
    bufferStorageVector.emplace_back(BufferStorage{tensorMemoryPointerVector, numTensors});

    // Create an instance of the SampleCustomAllocator
    SampleCustomAllocator customAllocator = SampleCustomAllocator();
    customAllocator.m_Values = {10, 11, 12, 13, 14};
    // Check that the test was set up correctly
    CHECK(customAllocator.m_Values.size() == numTensors);

    // Utilise 3 functions in the MemoryManager. Check the counters and the pointer to the values are correct.
    MemoryManager memoryManager;
    memoryManager.StoreMemToAllocate(bufferStorageVector, &customAllocator);

    memoryManager.Allocate();
    CHECK(customAllocator.m_CounterAllocate == bufferStorageVector.size());
    for (const auto& bufferStorage : bufferStorageVector)
    {
        uint32_t idx = 0;
        for (auto tensorMemory : bufferStorage.m_TensorMemoryVector)
        {
            auto value = reinterpret_cast<uint8_t *>(tensorMemory->m_Data);
            CHECK(customAllocator.m_Values[idx] == *value);
            idx += 1;
        }
    }

    memoryManager.Deallocate();
    CHECK(customAllocator.m_CounterFree == bufferStorageVector.size());
}
}

} // namespace armnn