PublicBin/minimal_scheduler.cpp

// Build with: clang++ --std=c++20 -g -O1 schedtest.cpp

#include <array>
#include <span>
#include <memory>
#include <list>
#include <functional>
#include <iostream>

#if defined(__x86_64__)
    #if defined(__GNUG__)
        constexpr bool is_x86_64_linux = true;
    #else
        #if defined(__clang__)
            constexpr bool is_x86_64_linux = true;
        #else
            constexpr bool is_x86_64_linux = false;
        #endif
    #endif
#else
    constexpr bool is_x86_64_linux = false;
#endif

static_assert(is_x86_64_linux, "This code probably only works on x86_64 GNU Extensions C++");

struct platform_data {
    platform_data() = default;
    
    platform_data(std::span<char> stack_str)
    : stack_ptr(&*(stack_str.end()-16))
    , base_ptr(stack_ptr)
    {}

    uint64_t rbx, r12, r13, r14, r15;

    void* stack_ptr;
    void* base_ptr;

    void pull() __attribute__((always_inline))
    {
        __asm__ __volatile__(
            "movq %%rsp, %0\n"
            "movq %%rbp, %1\n" 
            "movq %%rbx, %2\n" 
            "movq %%r12, %3\n" 
            "movq %%r13, %4\n" 
            "movq %%r14, %5\n" 
            "movq %%r15, %6\n" 
            : "=m"(stack_ptr)
            , "=m"(base_ptr)
            , "=m"(rbx)
            , "=m"(r12)
            , "=m"(r13)
            , "=m"(r14)
            , "=m"(r15)
        );
    }

    void* push(void* location)  __attribute__((always_inline))
    {
        volatile void* volatile tmp = static_cast<char*>(stack_ptr) - sizeof(void*);
        *static_cast<volatile void* volatile * volatile>(tmp) = location;
        __asm__ __volatile__(
            "movq %1, %%rsp\n"
            "movq %2, %%rbp\n"
            "movq %3, %%rbx\n"
            "movq %4, %%r12\n"
            "movq %5, %%r13\n"
            "movq %6, %%r14\n"
            "movq %7, %%r15\n"
            "popq %0\n"
            : "+r"(location)
            : "m"(tmp)
            , "m"(base_ptr)
            , "m"(rbx)
            , "m"(r12)
            , "m"(r13)
            , "m"(r14)
            , "m"(r15)
            : "memory"
        );
        return location;
    }
};

enum class process_status {
    inactive = 0,
    running = 1,
    waiting = 2,
    finished = 3,
    zombie = 4
};

struct process {
    char* stack;
    size_t sz;
    platform_data scheduling_swapper;
    process_status state = process_status::inactive;
    std::function<void()> fn;

    process(std::function<void()> _fn, size_t _sz = 16384)
    : stack(new char[_sz])
    , sz(_sz)
    , scheduling_swapper(std::span<char>(stack, sz))
    , fn(_fn)
    {}

    process(char* stack_base)
    : stack(stack_base)
    , sz(0)
    {}

    process(const process&) = delete;

    ~process() {
        if(sz) delete[] stack;
    }
};


__attribute__((noinline)) struct system* spawner (struct system* sys);

struct system {
    static system sys;

    std::list<std::unique_ptr<process>>
        running,
        waiting,
        naughty;
    
    std::unique_ptr<process> previous;
    std::unique_ptr<process> current;
    
    std::unique_ptr<process> one() {
        auto v = std::move(running.back());
        running.pop_back();
        return v;
    }

    void rid(std::unique_ptr<process> current) {
        switch(current->state) {
            case process_status::inactive:
            case process_status::running:
                running.push_front(std::move(current));
                break;
            case process_status::finished:
                clean(std::move(current));
                break;
            case process_status::zombie:
                naughty.push_front(std::move(current));
                break;
            case process_status::waiting:
                waiting.push_front(std::move(current));
                break;
        }
    }

    void clean(std::unique_ptr<process>) {}

    void yield_to(std::unique_ptr<process> target) noexcept {
        current->scheduling_swapper.pull();
        sys.rid(std::move(current));
        current = std::move(target);
        current->scheduling_swapper.push(this);
        spawner(&sys);
    }


    void yield() noexcept {
        current->scheduling_swapper.pull();
        sys.rid(std::move(current));
        current = one();
        current->scheduling_swapper.push(this);
        spawner(&sys);
    }
};


// Needs to return the system one way or another
__attribute__((noinline)) struct system* spawner (struct system* sys) {
    auto& proc = *system::sys.current;
    if(proc.state == process_status::inactive) {
        proc.state = process_status::running;
        proc.fn();
        proc.state = process_status::finished;
        sys->current->scheduling_swapper.pull();
        sys->yield();
    }
    return sys;
}

struct system system::sys;

int main() {
    char c;
    system::sys.current = std::make_unique<process>(&c);
    std::cout << "1";
    system::sys.current->state = process_status::running;
    system::sys.yield_to(std::make_unique<process>([](){
        std::cout << "2";
    }));
    std::cout << "3\n";
    return 0;
}