//-------------------------------------------------------------
// benchmark.cc: non-local goto test
//-------------------------------------------------------------

#include "context.h"

#include <benchmark/benchmark.h>
#include <cassert>
#include <coroutine>
#include <csetjmp>
#include <generator>
#include <signal.h>
#include <ucontext.h>

//-------------------------------------------------------------

static void empty_test(benchmark::State &state)
{
    for (auto _: state)
        benchmark::DoNotOptimize(_);
}

BENCHMARK(empty_test);

//-------------------------------------------------------------

static void setjmp_test(benchmark::State &state)
{
    jmp_buf env;
    for (auto _: state)
    {
        volatile int phase = 0;
        if (setjmp(env) == 0)
        {
            phase = 1;
            longjmp(env, 1);
        }
        benchmark::DoNotOptimize(phase);
    }
}

BENCHMARK(setjmp_test);

//-------------------------------------------------------------

static void sigsetjmp_test(benchmark::State &state)
{
    sigjmp_buf env;
    for (auto _: state)
    {
        volatile int phase = 0;
        if (sigsetjmp(env, 1) == 0)
        {
            phase = 1;
            siglongjmp(env, 1);
        }
        benchmark::DoNotOptimize(phase);
    }
}

BENCHMARK(sigsetjmp_test);

//-------------------------------------------------------------

static void getcontext_test(benchmark::State &state)
{
    ucontext_t context;
    for (auto _: state)
    {
        volatile int phase = 0;
        if (getcontext(&context) == 0 && phase == 0)
        {
            phase = 1;
            setcontext(&context);
        }
        benchmark::DoNotOptimize(phase);
    }
}

BENCHMARK(getcontext_test);

//-------------------------------------------------------------

static thread_local ucontext_t *main_ctx_ptr = nullptr;
static thread_local ucontext_t *_worker_ctx_ptr = nullptr;
static thread_local volatile int *phase_ptr = nullptr;

static void worker_loop()
{
    do
        *phase_ptr = 1;
    while (!swapcontext(_worker_ctx_ptr, main_ctx_ptr));
}

static void swapcontext_test(benchmark::State &state)
{
    ucontext_t main_ctx{};
    ucontext_t worker_ctx{};
    alignas(16) char worker_stack[1 << 16];

    assert(getcontext(&worker_ctx) == 0);

    worker_ctx.uc_link = nullptr;
    worker_ctx.uc_stack.ss_sp = worker_stack;
    worker_ctx.uc_stack.ss_size = sizeof(worker_stack);
    worker_ctx.uc_stack.ss_flags = 0;
    makecontext(&worker_ctx, worker_loop, 0);

    main_ctx_ptr = &main_ctx;
    _worker_ctx_ptr = &worker_ctx;

    for (auto _: state)
    {
        volatile int phase = 0;
        phase_ptr = &phase;
        assert(swapcontext(&main_ctx, &worker_ctx) == 0);
        benchmark::DoNotOptimize(phase);
    }
}

BENCHMARK(swapcontext_test);

//-------------------------------------------------------------

static void inline_context_test(benchmark::State &state)
{
    as::inline_context context;
    for (auto _: state)
    {
        volatile int phase = 0;
        if (context.get() == as::first)
        {
            phase = 1;
            context.set();
        }
        benchmark::DoNotOptimize(phase);
    }
}

BENCHMARK(inline_context_test);

//-------------------------------------------------------------

static void context_test(benchmark::State &state)
{
    as::context context;
    for (auto _: state)
    {
        volatile int phase = 0;
        if (context.get() == as::first)
        {
            phase = 1;
            context.set();
        }
        benchmark::DoNotOptimize(phase);
    }
}

BENCHMARK(context_test);

//-------------------------------------------------------------

struct fiber
{
    struct promise_type
    {
        fiber get_return_object()
        {
            return fiber{handle_type::from_promise(*this)};
        }
        std::suspend_always initial_suspend() noexcept { return {}; }
        std::suspend_never final_suspend() noexcept { return {}; }
        void return_void() noexcept {}
        void unhandled_exception() noexcept {}
    };

    using handle_type = std::coroutine_handle<promise_type>;
    handle_type handle;

    explicit fiber(handle_type h): handle(h) {}
};

static fiber ping_pong(size_t *counter)
{
    while (true)
    {
        ++(*counter);
        co_await std::suspend_always{};
    }
}

static void coroutine_test(benchmark::State &state)
{
    size_t counter = 0;
    auto f = ping_pong(&counter);

    for (auto _: state)
    {
        f.handle.resume();
        benchmark::DoNotOptimize(counter);
    }
}

BENCHMARK(coroutine_test);

//-------------------------------------------------------------

static std::generator<size_t> counter_generator()
{
    size_t counter = 0;
    while (true)
        co_yield ++counter;
}

static void generator_test(benchmark::State &state)
{
    auto gen = counter_generator();
    auto it = gen.begin();

    for (auto _: state)
    {
        benchmark::DoNotOptimize(++it);
    }
}

BENCHMARK(generator_test);

//-------------------------------------------------------------

BENCHMARK_MAIN();

//-------------------------------------------------------------