/*
 * This file is part of the Flowee project
 * Copyright (C) 2021-2016 The Bitcoin Core developers
 * Copyright (C) 2022 Tom Zander <tom@flowee.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "allocator_tests.h"

#include <support/allocators/secure.h>

void TestAllocator::testArena()
{
    // Fake memory base address for testing
    // without actually using memory.
    void *synth_base = reinterpret_cast<void *>(0x08000000);
    const size_t synth_size = 1024 * 1024;
    Arena b(synth_base, synth_size, 16);
    void *chunk = b.alloc(1000);
#ifdef ARENA_DEBUG
    b.walk();
#endif
    QVERIFY(chunk != nullptr);
    // Aligned to 16
    QVERIFY(b.stats().used == 1008);
    // Nothing has disappeared?
    QVERIFY(b.stats().total == synth_size);
    b.free(chunk);
#ifdef ARENA_DEBUG
    b.walk();
#endif
    QVERIFY(b.stats().used == 0);
    QVERIFY(b.stats().free == synth_size);
    try {
        // Test exception on double-free
        b.free(chunk);
        QVERIFY(0);
    } catch (std::runtime_error &) {
    }

    void *a0 = b.alloc(128);
    void *a1 = b.alloc(256);
    void *a2 = b.alloc(512);
    QVERIFY(b.stats().used == 896);
    QVERIFY(b.stats().total == synth_size);
#ifdef ARENA_DEBUG
    b.walk();
#endif
    b.free(a0);
#ifdef ARENA_DEBUG
    b.walk();
#endif
    QVERIFY(b.stats().used == 768);
    b.free(a1);
    QVERIFY(b.stats().used == 512);
    void *a3 = b.alloc(128);
#ifdef ARENA_DEBUG
    b.walk();
#endif
    QVERIFY(b.stats().used == 640);
    b.free(a2);
    QVERIFY(b.stats().used == 128);
    b.free(a3);
    QVERIFY(b.stats().used == 0);
    QCOMPARE(b.stats().chunks_used, 0U);
    QVERIFY(b.stats().total == synth_size);
    QVERIFY(b.stats().free == synth_size);
    QCOMPARE(b.stats().chunks_free, 1U);

    std::vector<void *> addr;
    // allocating 0 always returns nullptr
    QVERIFY(b.alloc(0) == nullptr);
#ifdef ARENA_DEBUG
    b.walk();
#endif
    // Sweeping allocate all memory
    for (int x = 0; x < 1024; ++x) {
        addr.push_back(b.alloc(1024));
    }
    QVERIFY(b.stats().free == 0);
    // memory is full, this must return nullptr
    QVERIFY(b.alloc(1024) == nullptr);
    QVERIFY(b.alloc(0) == nullptr);
    for (int x = 0; x < 1024; ++x) {
        b.free(addr[x]);
    }
    addr.clear();
    QVERIFY(b.stats().total == synth_size);
    QVERIFY(b.stats().free == synth_size);

    // Now in the other direction...
    for (int x = 0; x < 1024; ++x) {
        addr.push_back(b.alloc(1024));
    }
    for (int x = 0; x < 1024; ++x) {
        b.free(addr[1023 - x]);
    }
    addr.clear();

    // Now allocate in smaller unequal chunks, then deallocate haphazardly
    // Not all the chunks will succeed allocating, but freeing nullptr is
    // allowed so that is no problem.
    for (int x = 0; x < 2048; ++x) {
        addr.push_back(b.alloc(x + 1));
    }
    for (int x = 0; x < 2048; ++x) {
        b.free(addr[((x * 23) % 2048) ^ 242]);
    }
    addr.clear();

    // Go entirely wild: free and alloc interleaved, generate targets and sizes
    // using pseudo-randomness.
    for (int x = 0; x < 2048; ++x) {
        addr.push_back(nullptr);
    }
    uint32_t s = 0x12345678;
    for (int x = 0; x < 5000; ++x) {
        int idx = s & (addr.size() - 1);
        if (s & 0x80000000) {
            b.free(addr[idx]);
            addr[idx] = nullptr;
        } else if (!addr[idx]) {
            addr[idx] = b.alloc((s >> 16) & 2047);
        }
        bool lsb = s & 1;
        s >>= 1;
        // LFSR period 0xf7ffffe0
        if (lsb) {
            s ^= 0xf00f00f0;
        }
    }
    for (void *ptr : addr) {
        b.free(ptr);
    }
    addr.clear();

    QVERIFY(b.stats().total == synth_size);
    QVERIFY(b.stats().free == synth_size);
}

/** Mock LockedPageAllocator for testing */
class TestLockedPageAllocator : public LockedPageAllocator {
public:
    TestLockedPageAllocator(int count_in, int lockedcount_in)
        : count(count_in), lockedcount(lockedcount_in) {}
    void *AllocateLocked(size_t len, bool *lockingSuccess) override {
        *lockingSuccess = false;
        if (count > 0) {
            --count;

            if (lockedcount > 0) {
                --lockedcount;
                *lockingSuccess = true;
            }

            // Fake address, do not actually use this memory
            return reinterpret_cast<void *>(0x08000000 + (count << 24));
        }
        return nullptr;
    }
    void FreeLocked(void *addr, size_t len) override {}
    size_t GetLimit() override { return std::numeric_limits<size_t>::max(); }

private:
    int count;
    int lockedcount;
};

void TestAllocator::testLockedPoolMock()
{
    // Test over three virtual arenas, of which one will succeed being locked
    auto x = std::make_unique<TestLockedPageAllocator>(3, 1);
    LockedPool pool(std::move(x));
    QVERIFY(pool.stats().total == 0);
    QVERIFY(pool.stats().locked == 0);

    // Ensure unreasonable requests are refused without allocating anything
    void *invalid_toosmall = pool.alloc(0);
    QVERIFY(invalid_toosmall == nullptr);
    QCOMPARE(pool.stats().used, 0);
    QCOMPARE(pool.stats().free, 0);
    void *invalid_toobig = pool.alloc(LockedPool::ARENA_SIZE + 1);
    QVERIFY(invalid_toobig == nullptr);
    QCOMPARE(pool.stats().used, 0);
    QCOMPARE(pool.stats().free, 0);

    void *a0 = pool.alloc(LockedPool::ARENA_SIZE / 2);
    QVERIFY(a0);
    QVERIFY(pool.stats().locked == LockedPool::ARENA_SIZE);
    void *a1 = pool.alloc(LockedPool::ARENA_SIZE / 2);
    QVERIFY(a1);
    void *a2 = pool.alloc(LockedPool::ARENA_SIZE / 2);
    QVERIFY(a2);
    void *a3 = pool.alloc(LockedPool::ARENA_SIZE / 2);
    QVERIFY(a3);
    void *a4 = pool.alloc(LockedPool::ARENA_SIZE / 2);
    QVERIFY(a4);
    void *a5 = pool.alloc(LockedPool::ARENA_SIZE / 2);
    QVERIFY(a5);
    // We've passed a count of three arenas, so this allocation should fail
    void *a6 = pool.alloc(16);
    QVERIFY(!a6);

    pool.free(a0);
    pool.free(a2);
    pool.free(a4);
    pool.free(a1);
    pool.free(a3);
    pool.free(a5);
    QVERIFY(pool.stats().total == 3 * LockedPool::ARENA_SIZE);
    QVERIFY(pool.stats().locked == LockedPool::ARENA_SIZE);
    QVERIFY(pool.stats().used == 0);
}

// These tests used the live LockedPoolManager object, this is also used by
// other tests so the conditions are somewhat less controllable and thus the
// tests are somewhat more error-prone.
void TestAllocator::testLockedPool()
{
    LockedPoolManager &pool = LockedPoolManager::instance();
    LockedPool::Stats initial = pool.stats();

    void *a0 = pool.alloc(16);
    QVERIFY(a0);
    // Test reading and writing the allocated memory
    *((uint32_t *)a0) = 0x1234;
    QVERIFY(*((uint32_t *)a0) == 0x1234);

    pool.free(a0);
    try {
        // Test exception on double-free
        pool.free(a0);
        QVERIFY(0);
    } catch (std::runtime_error &) {
    }
    // If more than one new arena was allocated for the above tests, something
    // is wrong
    QVERIFY(pool.stats().total <= (initial.total + LockedPool::ARENA_SIZE));
    // Usage must be back to where it started
    QVERIFY(pool.stats().used == initial.used);
}