#include "context.h"

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

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

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

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

static void getcontext_test(benchmark::State &state)
{
    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;

    auto worker = +[]()
    {
        while (true)
        {
            *phase_ptr = 1;
            swapcontext(worker_ctx_ptr, main_ctx_ptr);
        }
    };

    ucontext_t main_ctx{};
    ucontext_t worker_ctx{};
    alignas(16) char worker_stack[1 << 16];

    if (getcontext(&worker_ctx) != 0)
    {
        state.SkipWithError(
            "getcontext failed while preparing worker context");
        return;
    }

    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, 0);

    main_ctx_ptr = &main_ctx;
    worker_ctx_ptr = &worker_ctx;

    for (auto _: state)
    {
        volatile int phase = 0;

        phase_ptr = &phase;
        swapcontext(&main_ctx, &worker_ctx);

        benchmark::DoNotOptimize(phase);
        benchmark::ClobberMemory();
    }
}
BENCHMARK(getcontext_test);

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

static void context_test(benchmark::State &state)
{
    for (auto _: state)
    {
        as::context context;
        volatile int phase = 0;
        if (context.get() == as::first)
        {
            phase = 1;
            context.set();
        }
        benchmark::DoNotOptimize(phase);
        benchmark::ClobberMemory();
    }
}
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_always 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) {}
    ~fiber()
    {
        if (handle)
            handle.destroy();
    }
};

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

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

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

    benchmark::DoNotOptimize(counter);
}
BENCHMARK(coroutine_test);

BENCHMARK_MAIN();