/*
 * Copyright (C) 2018-present ScyllaDB
 */

/*
 * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
 */

#define BOOST_TEST_MODULE small_vector

#include <boost/test/unit_test.hpp>
#include <fmt/ranges.h>
#include <algorithm>
#include <functional>
#include <sstream>

#include "utils/small_vector.hh"

template<typename T, size_t N>
void check_equivalent(const utils::small_vector<T, N>& a, const std::vector<T>& b) {
    BOOST_REQUIRE_EQUAL(a.size(), b.size());
    BOOST_REQUIRE_GE(a.capacity(), a.size());
    for (auto i = 0u; i < b.size(); i++) {
        BOOST_REQUIRE(a[i] == b[i]);
    }
    auto idx = 0;
    for (auto&& v : a) {
        BOOST_REQUIRE(v == b[idx++]);
    }
}

template<typename Iterator>
class input_iterator_wrapper {
    Iterator _it;
public:
    using value_type = typename std::iterator_traits<Iterator>::value_type;
    using pointer = typename std::iterator_traits<Iterator>::pointer;
    using reference = typename std::iterator_traits<Iterator>::reference;
    using difference_type = typename std::iterator_traits<Iterator>::difference_type;
    using iterator_category = std::input_iterator_tag;

    input_iterator_wrapper(Iterator it) : _it(it) { }

    value_type operator*() const { return *_it; }

    input_iterator_wrapper& operator++() {
        ++_it;
        return *this;
    }
    input_iterator_wrapper operator++(int) {
        auto it = *this;
        operator++();
        return it;
    }

    bool operator==(const input_iterator_wrapper& other) const = default;
};

template<typename T>
void test_random_walk(std::function<T()> make_element) {
    utils::small_vector<T, 8> actual;
    std::vector<T> expected;

    auto emplace_back = [&] (T x) {
        actual.emplace_back(x);
        expected.emplace_back(x);
        check_equivalent(actual, expected);
    };

    auto test_range_ctor = [&] (utils::small_vector<T, 8>& actual) {
        auto a1 = utils::small_vector<T, 8>(actual.begin(), actual.end());
        check_equivalent(a1, expected);

        auto a2 = utils::small_vector<T, 8>(input_iterator_wrapper(actual.begin()), input_iterator_wrapper(actual.end()));
        check_equivalent(a2, expected);
    };

    auto move_ctor = [&] (utils::small_vector<T, 8>& actual) {
        auto a = std::move(actual);
        check_equivalent(a, expected);
        return a;
    };

    auto copy_ctor = [&] (utils::small_vector<T, 8>& actual) {
        auto a = actual;
        check_equivalent(a, expected);
        return a;
    };

    auto move_assign = [&] (utils::small_vector<T, 8>& src, utils::small_vector<T, 8>& dst) {
        dst = std::move(src);
        check_equivalent(dst, expected);
    };

    auto copy_assign = [&] (utils::small_vector<T, 8>& src, utils::small_vector<T, 8>& dst) {
        dst = src;
        check_equivalent(dst, expected);
    };


    for (auto i = 0; i < 64; i++) {
        emplace_back(make_element());

        test_range_ctor(actual);

        auto a1 = copy_ctor(actual);
        auto a2 = move_ctor(actual);

        // Assigning to moved-from containers
        move_assign(a1, actual);
        copy_assign(a2, a1);

        // Assigning to a container with other elements
        copy_assign(a2, actual);
        move_assign(a2, actual);

        // Assigning to a container with pre-existing elements, forcing slow
        // path
        auto a3 = utils::small_vector<T, 8>();
        for (auto j = 0; j <= i / 2; j++) {
            a3.emplace_back();
        }
        copy_assign(a1, a3);
    }

    auto another_actual = utils::small_vector<T, 8>(expected.begin(), expected.end());
    check_equivalent(another_actual, expected);

    for (auto i = 0u; i <= actual.size(); i++) {
        auto test_range_insertion = [&] (auto first, auto last) {
            auto a1 = actual;
            auto a2 = actual;
            check_equivalent(a1, expected);
            check_equivalent(a2, expected);

            auto e = expected;
            a1.insert(a1.begin() + i, first, last);
            a2.insert(a2.begin() + i, input_iterator_wrapper(first), input_iterator_wrapper(last));
            e.insert(e.begin() + i, first, last);
            check_equivalent(a1, e);
            check_equivalent(a2, e);
        };

        test_range_insertion(actual.begin(), actual.begin());
        test_range_insertion(actual.begin(), actual.end());
        test_range_insertion(actual.begin(), actual.begin() + 1);
        test_range_insertion(actual.begin(), actual.begin() + 4);

        {
            auto a = actual;
            auto e = expected;
            a.insert(a.begin() + i, actual[0]);
            e.insert(e.begin() + i, actual[0]);
            check_equivalent(a, e);
        }

        {
            auto a = actual;
            auto e = expected;
            a.erase(a.begin(), a.begin() + i);
            e.erase(e.begin(), e.begin() + i);
            check_equivalent(a, e);
        }

        if (i < actual.size()) {
            auto a = actual;
            auto e = expected;
            a.erase(a.begin() + i, a.end());
            e.erase(e.begin() + i, e.end());
            check_equivalent(a, e);
        }

        if (i < actual.size()) {
            auto a = actual;
            auto e = expected;
            a.erase(a.begin() + i);
            e.erase(e.begin() + i);
            check_equivalent(a, e);
        }
    }
}

BOOST_AUTO_TEST_CASE(random_walk_trivial) {
    test_random_walk<int>([x = 0] () mutable { return x++; });
}

BOOST_AUTO_TEST_CASE(random_walk_nontrivial) {
    test_random_walk<std::shared_ptr<int>>([x = 0] () mutable { return std::make_shared<int>(x++); });
}

template<typename T>
void test_insert(std::function<T()> make_element) {
    utils::small_vector<T, 8> actual;
    std::vector<T> expected;

    auto emplace_back = [&] {
        auto e = make_element();
        actual.emplace_back(e);
        expected.emplace_back(e);
    };

    auto insert_middle = [&] (size_t count) {
        auto a = actual;
        auto e = expected;
        check_equivalent(a, e);

        std::vector<T> vec;
        std::generate_n(std::back_inserter(vec), count, make_element);
        a.insert(a.begin() + 1, vec.begin(), vec.end());
        e.insert(e.begin() + 1, vec.begin(), vec.end());
        check_equivalent(a, e);
    };

    auto insert_end = [&] (size_t count) {
        auto a = actual;
        auto e = expected;
        check_equivalent(a, e);

        std::vector<T> vec;
        std::generate_n(std::back_inserter(vec), count, make_element);
        a.insert(a.end(), vec.begin(), vec.end());
        e.insert(e.end(), vec.begin(), vec.end());
        check_equivalent(a, e);
    };

    auto test_inserts = [&] {
        insert_middle(2);
        insert_middle(4);
        insert_middle(6);
        insert_middle(8);
        insert_middle(64);

        insert_end(2);
        insert_end(4);
        insert_end(6);
        insert_end(8);
        insert_end(64);
    };

    for (auto i = 0u; i < 2u; i++) {
        emplace_back();
    }

    test_inserts();

    for (auto i = 0u; i < 2u; i++) {
        emplace_back();
    }

    test_inserts();

    for (auto i = 0u; i < 4u; i++) {
        emplace_back();
    }

    test_inserts();
}

BOOST_AUTO_TEST_CASE(insert_trivial) {
    test_insert<int>([x = 0] () mutable { return x++; });
}

BOOST_AUTO_TEST_CASE(insert_nontrivial) {
    test_insert<std::shared_ptr<int>>([x = 0] () mutable { return std::make_shared<int>(x++); });
}

namespace  {

class fails_on_copy {
    size_t _counter;
public:
    static inline thread_local size_t live = 0;
public:
    explicit fails_on_copy(size_t counter) : _counter(counter) {
        live++;
    }
    ~fails_on_copy() {
        // Make sure the compiler doesn't do anything clever
        *reinterpret_cast<volatile size_t*>(&_counter) = -1;
        live--;
    }

    fails_on_copy(fails_on_copy&& other) noexcept : _counter(other._counter) {
        live++;
    }
    fails_on_copy(const fails_on_copy& other) : _counter(other._counter - 1) {
        if (!_counter) {
            throw 1;
        }
        live++;
    }
    fails_on_copy& operator=(fails_on_copy&&) = default;
    fails_on_copy& operator=(const fails_on_copy& other) {
        _counter = other._counter - 1;
        if (!_counter) {
            throw 1;
        }
        return *this;
    }

    size_t counter() const { return _counter; }
};

}

BOOST_AUTO_TEST_CASE(exception_safety) {
    std::vector<fails_on_copy> vec;
    vec.emplace_back(4);
    vec.emplace_back(1);

    BOOST_REQUIRE_THROW((utils::small_vector<fails_on_copy, 1>(vec.begin(), vec.end())), int);
    BOOST_REQUIRE_THROW((utils::small_vector<fails_on_copy, 4>(vec.begin(), vec.end())), int);

    utils::small_vector<fails_on_copy, 2> v;
    v.emplace_back(0);
    v.emplace_back(1);
    v.emplace_back(2);
    v.emplace_back(3);

    auto verify_unchanged = [&] {
        BOOST_REQUIRE_EQUAL(v.size(), 4);
        for (auto i = 0; i < 4; i++) {
            BOOST_REQUIRE_EQUAL(v[i].counter(), i);
        }
    };

    BOOST_REQUIRE_THROW(v.insert(v.begin(), vec.begin(), vec.end()), int);
    verify_unchanged();

    BOOST_REQUIRE_THROW(v.insert(v.end(), vec.begin(), vec.end()), int);
    verify_unchanged();

    vec.emplace_back(4);
    vec.emplace_back(4);
    vec.emplace_back(4);
    vec.emplace_back(4);
    BOOST_REQUIRE_THROW(v.insert(v.begin(), vec.begin(), vec.end()), int);
    verify_unchanged();

    vec.clear();
    vec.emplace_back(4);
    vec.emplace_back(4);
    vec.emplace_back(4);
    vec.emplace_back(4);
    vec.emplace_back(4);
    vec.emplace_back(1);
    BOOST_REQUIRE_THROW(v.insert(v.begin(), vec.begin(), vec.end()), int);
    verify_unchanged();

    {
        auto fc = fails_on_copy(1);
        BOOST_REQUIRE_THROW(v.insert(v.begin(), fc), int);
        verify_unchanged();

        fc = fails_on_copy(1);
        BOOST_REQUIRE_THROW(v.insert(v.end(), fc), int);
        verify_unchanged();

        fc = fails_on_copy(1);
        BOOST_REQUIRE_THROW(v.push_back(fc), int);
        verify_unchanged();
    }

    utils::small_vector<fails_on_copy, 2> v2;
    v2.emplace_back(4);
    v2.emplace_back(1);
    BOOST_REQUIRE_THROW(v = v2, int);
    verify_unchanged();

    v2.emplace_back(4);
    v2.emplace_back(4);
    BOOST_REQUIRE_THROW(v = v2, int);
    verify_unchanged();

    v2.clear();
    v.clear();
    vec.clear();

    BOOST_REQUIRE_EQUAL(fails_on_copy::live, 0);
}

BOOST_AUTO_TEST_CASE(compare) {
    {
        auto lhs = utils::small_vector<int, 4>({1, 2, 3, 4});
        auto rhs = utils::small_vector<int, 4>({1, 3, 3, 4});
        BOOST_CHECK_LT(lhs, rhs);
    }
    {
        auto lhs = utils::small_vector<int, 4>({1, 2, 3, 4});
        auto rhs = utils::small_vector<int, 4>({1, 3, 4});
        BOOST_CHECK_LT(lhs, rhs);
    }
    {
        auto lhs = utils::small_vector<int, 4>({1, 2, 3, 4});
        auto rhs = utils::small_vector<int, 4>({1, 3, 4});
        BOOST_CHECK_LT(lhs, rhs);
    }
    {
        auto lhs = utils::small_vector<int, 4>({4, 2, 1});
        auto rhs = utils::small_vector<int, 4>({1, 3, 3, 4});
        BOOST_CHECK_GT(lhs, rhs);
    }
    {
        auto lhs = utils::small_vector<int, 4>({4, 2, 1});
        auto rhs = utils::small_vector<int, 4>({1, 3, 3, 4});
        BOOST_CHECK_GE(lhs, rhs);
    }
}

BOOST_AUTO_TEST_CASE(resize) {
    auto vec = utils::small_vector<int, 4>();
    vec.emplace_back(1);
    vec.resize(1024);

    BOOST_REQUIRE_EQUAL(vec.size(), 1024);
    BOOST_REQUIRE_EQUAL(vec[0], 1);
    for (auto i = 0; i < 1023; i++) {
        BOOST_REQUIRE_EQUAL(vec[i + 1], 0);
    }

    vec.resize(1024);
    BOOST_REQUIRE_EQUAL(vec.size(), 1024);
    BOOST_REQUIRE_EQUAL(vec[0], 1);
    for (auto i = 0; i < 1023; i++) {
        BOOST_REQUIRE_EQUAL(vec[i + 1], 0);
    }

    vec.resize(512);
    BOOST_REQUIRE_EQUAL(vec.size(), 512);
    BOOST_REQUIRE_EQUAL(vec[0], 1);
    for (auto i = 0; i < 511; i++) {
        BOOST_REQUIRE_EQUAL(vec[i + 1], 0);
    }

    vec.resize(0);
    BOOST_REQUIRE_EQUAL(vec.size(), 0);
    for (auto&& v : vec) {
        (void)v;
        BOOST_FAIL("should not reach");
    }
}

class ctor_assn_measuring {
    static inline size_t _counter = 0;
    int _x;
public:
    ctor_assn_measuring(int x) : _x(x) {
        ++_counter;
    }
    ctor_assn_measuring(const ctor_assn_measuring& other) : _x(other._x) {
        _counter++;
    }
    ctor_assn_measuring(ctor_assn_measuring&& other) noexcept : _x(other._x) {
        _counter++;
    }
    ctor_assn_measuring& operator=(const ctor_assn_measuring& other) {
        _x = other._x;
        _counter++;
        return *this;
    }
    ctor_assn_measuring& operator=(ctor_assn_measuring&& other) noexcept {
        _x = other._x;
        _counter++;
        return *this;
    }
    static size_t counter() {
        return _counter;
    }
    bool operator==(const ctor_assn_measuring& other) const = default;
};

static
size_t
do_test_from_range(std::ranges::range auto src, std::ranges::range auto src_copy) {
    auto plus1 = std::views::transform(std::bind_front(std::plus(), 1));
    auto before_start = ctor_assn_measuring::counter();
    auto dst = src | plus1 | std::ranges::to<utils::small_vector<ctor_assn_measuring, 4>>();
    auto after_end = ctor_assn_measuring::counter();
    // We cannot compare against src if it's an input_range, since it can only be
    // iterated once. Compare against src_copy.
    BOOST_REQUIRE(std::ranges::equal(dst, src_copy | plus1));
    return after_end - before_start;
}

BOOST_AUTO_TEST_CASE(from_range) {
    auto iota_view = std::views::iota(0, 10);
    // Verify that no copies are needed if the range is sized.
    BOOST_REQUIRE_EQUAL(do_test_from_range(iota_view, iota_view), 10);
    auto iota_str = "0 1 2 3 4 5 6 7 8 9";
    auto stream = std::stringstream(iota_str);
    auto stream_view = std::ranges::istream_view<int>(stream);
    // Verify that copies are needed if the range is (only) an input_range.
    BOOST_REQUIRE_GT(do_test_from_range(stream_view, iota_view), 10);
}