Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
#include <algorithm>
#include <array>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
static const size_t STACK_SIZE = 16777216;
static const size_t DEFAULT_THREADS = 4;
struct Runtime;
static Runtime* runtime;
enum class ThreadState {
Available,
Ready,
Running,
};
struct ThreadContext {
uint64_t rsp, r15, r14, r13, r12, rbx, rbp;
ThreadContext() : rsp(), r15(), r14(), r13(), r12(), rbx(), rbp() {}
static __attribute__((noinline)) __attribute__((naked)) void
switch_to(ThreadContext&, ThreadContext&) {
__asm__ volatile(
"movq %rsp, (%rdi)\n\t"
"movq %r15, 0x8(%rdi)\n\t"
"movq %r14, 0x10(%rdi)\n\t"
"movq %r13, 0x18(%rdi)\n\t"
"movq %r12, 0x20(%rdi)\n\t"
"movq %rbx, 0x28(%rdi)\n\t"
"movq %rbp, 0x30(%rdi)\n\t"
"movq (%rsi), %rsp\n\t"
"movq 0x8(%rsi), %r15\n\t"
"movq 0x10(%rsi), %r14\n\t"
"movq 0x18(%rsi), %r13\n\t"
"movq 0x20(%rsi), %r12\n\t"
"movq 0x28(%rsi), %rbx\n\t"
"movq 0x30(%rsi), %rbp\n\t"
"ret");
}
};
static_assert(offsetof(ThreadContext, rsp) == 0, "unexpected offset");
static_assert(offsetof(ThreadContext, r15) == 010, "unexpected offset");
static_assert(offsetof(ThreadContext, r14) == 020, "unexpected offset");
static_assert(offsetof(ThreadContext, r13) == 030, "unexpected offset");
static_assert(offsetof(ThreadContext, r12) == 040, "unexpected offset");
static_assert(offsetof(ThreadContext, rbx) == 050, "unexpected offset");
static_assert(offsetof(ThreadContext, rbp) == 060, "unexpected offset");
struct Thread {
size_t id;
void* stack;
ThreadContext ctx;
ThreadState state;
Thread() : id(0), stack(), state(ThreadState::Available) {
stack = malloc(STACK_SIZE);
if (!stack) {
fputs("Could not allocate memory for new thread.\n", stderr);
exit(1);
}
}
};
struct Runtime {
std::array<Thread, 1 + DEFAULT_THREADS> threads;
size_t current;
Runtime() : threads(), current(0) {
threads.front().state = ThreadState::Running;
for (size_t i = 1; i <= DEFAULT_THREADS; ++i) { threads[i].id = i; }
}
void init() { runtime = this; }
void run() {
while (yield())
;
exit(0);
}
void ret() {
if (current) {
threads[current].state = ThreadState::Available;
yield();
}
}
bool yield() {
auto next_ready = std::find_if(
threads.begin() + current, threads.end(),
[](Thread const& th) { return th.state == ThreadState::Ready; });
if (next_ready == threads.end()) {
next_ready = std::find_if(
threads.begin(), threads.begin() + current,
[](Thread const& th) { return th.state == ThreadState::Ready; });
}
if (next_ready == threads.begin() + current) { return false; }
if (threads[current].state != ThreadState::Available)
threads[current].state = ThreadState::Ready;
next_ready->state = ThreadState::Running;
size_t old = current;
current = next_ready - threads.begin();
ThreadContext::switch_to(threads[old].ctx, next_ready->ctx);
return true;
}
static void finish() { runtime->ret(); }
void spawn(void (*func)()) {
auto avail =
std::find_if(threads.begin(), threads.end(), [](Thread const& th) {
return th.state == ThreadState::Available;
});
if (avail == threads.end()) {
fputs("Cannot spawn task. No available threads.\n", stderr);
exit(2);
}
void (*guard)() = Runtime::finish;
memcpy((unsigned char*) avail->stack + STACK_SIZE - 8, &guard, 8);
memcpy((unsigned char*) avail->stack + STACK_SIZE - 16, &func, 8);
avail->ctx.rsp =
(uint64_t)((unsigned char*) avail->stack + STACK_SIZE - 16);
avail->state = ThreadState::Ready;
}
};
static void thread1() {
puts("thread 1 started");
for (int i = 0; i < 10; ++i) {
printf("thread 1 counting %d\n", i);
runtime->yield();
}
}
static void thread2() {
puts("thread 2 started");
for (int i = 0; i < 10; ++i) {
printf("thread 2 counting %d\n", i);
runtime->yield();
}
}
int main() {
Runtime().init();
runtime->spawn(thread1);
runtime->spawn(thread2);
runtime->run();
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.