Add userspace threads implementation to toolbox

2 files changed, 279 insertions, 0 deletions

diff --git a/gthread.c b/gthread.c
new file mode 100644
index 0000000..024273c
--- /dev/null
+++ b/gthread.c
@@ -0,0 +1,241 @@
+/*
+ * Copyright (C) 2023 Mikhail Burakov. This file is part of toolbox.
+ *
+ * toolbox 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.
+ *
+ * toolbox 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 toolbox.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+#include "gthread.h"
+
+#include <stddef.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+
+struct GThread {
+  struct GThread* prev;
+
+  struct {
+    uint64_t rbx;  // 0x08
+    uint64_t rbp;  // 0x10
+    uint64_t rsp;  // 0x18
+    uint64_t r12;  // 0x20
+    uint64_t r13;  // 0x28
+    uint64_t r14;  // 0x30
+    uint64_t r15;  // 0x38
+  } registers;
+
+  struct {
+    uint64_t tip;
+    size_t alloc;
+    void* data;
+  } stack;
+};
+
+static _Thread_local struct GThread g_master_thread;
+static _Thread_local struct GThread* __attribute__((used)) g_current_thread;
+static _Thread_local size_t g_threads_count;
+
+static struct GThread* GThreadAlloc(size_t stack_alloc) {
+  struct GThread* thread = malloc(sizeof(struct GThread));
+  if (!thread) {
+    return NULL;
+  }
+  thread->stack.data = malloc(stack_alloc);
+  thread->stack.alloc = stack_alloc;
+  if (!thread->stack.data) {
+    free(thread);
+    return NULL;
+  }
+  return thread;
+}
+
+static void __attribute__((naked)) Trampoline(void) {
+  __asm__(
+      "pop %rdi\n"
+      "ret\n");
+}
+
+static struct GThread* __attribute__((used))
+GThreadCreateImpl(void (*proc)(void*), void* user, uint64_t stack_tip) {
+  void* stack_data[] = {NULL, (void*)Trampoline, user, (void*)proc, NULL};
+  struct GThread* thread = GThreadAlloc(sizeof(stack_data));
+  if (!thread) {
+    return NULL;
+  }
+  if (!g_threads_count++) {
+    g_current_thread = &g_master_thread;
+  }
+  if (stack_tip > g_master_thread.stack.tip) {
+    g_master_thread.stack.tip = stack_tip;
+  }
+  thread->stack.tip = stack_tip;
+  thread->registers.rsp = stack_tip - sizeof(stack_data);
+  memcpy(thread->stack.data, stack_data, sizeof(stack_data));
+  return thread;
+}
+
+static bool __attribute__((used, naked))
+GThreadSaveRegisters(struct GThread* thread) {
+  __asm__(
+      "pop %rax\n"
+      "mov %rbx, 0x08(%rdi)\n"
+      "mov %rbp, 0x10(%rdi)\n"
+      "mov %rsp, 0x18(%rdi)\n"
+      "mov %r12, 0x20(%rdi)\n"
+      "mov %r13, 0x28(%rdi)\n"
+      "mov %r14, 0x30(%rdi)\n"
+      "mov %r15, 0x38(%rdi)\n"
+      "push %rax\n"
+      "ret\n");
+}
+
+static bool __attribute__((used)) GThreadSaveStack(struct GThread* thread) {
+  if (thread->registers.rsp >= thread->stack.tip) {
+    // mburakov: Stack tip is down the stack, no need to save anything.
+    return true;
+  }
+  size_t stack_size = thread->stack.tip - thread->registers.rsp;
+  if (thread->stack.alloc < stack_size) {
+    free(thread->stack.data);
+    thread->stack.data = malloc(stack_size);
+    thread->stack.alloc = stack_size;
+    if (!thread->stack.data) {
+      thread->stack.alloc = 0;
+      return false;
+    }
+  }
+  memcpy(thread->stack.data, (void*)thread->registers.rsp, stack_size);
+  return true;
+}
+
+static void __attribute__((used, naked))
+GThreadRestoreRegisters(struct GThread* thread) {
+  __asm__(
+      "pop %rax\n"
+      "mov 0x08(%rdi), %rbx\n"
+      "mov 0x10(%rdi), %rbp\n"
+      "mov 0x18(%rdi), %rsp\n"
+      "mov 0x20(%rdi), %r12\n"
+      "mov 0x28(%rdi), %r13\n"
+      "mov 0x30(%rdi), %r14\n"
+      "mov 0x38(%rdi), %r15\n"
+      "push %rax\n"
+      "ret\n");
+}
+
+static void __attribute__((used)) GThreadRestoreStack(struct GThread* thread) {
+  if (thread->registers.rsp >= thread->stack.tip) {
+    // mburakov: Stack tip is down the stack, no need to restore anything.
+    return;
+  }
+  size_t stack_size = thread->stack.tip - thread->registers.rsp;
+  memcpy((void*)thread->registers.rsp, thread->stack.data, stack_size);
+}
+
+static void GThreadFree(struct GThread* thread) {
+  free(thread->stack.data);
+  free(thread);
+}
+
+struct GThread* __attribute__((naked))
+GThreadCreate(void (*proc)(void*), void* user) {
+  __asm__(
+      // mburakov: Save stack tip and pass it to the impl. Tip is measured from
+      // outside of spawn function (8 bytes offset). On threads switching, stack
+      // would be preserved starting from this point till stack pointer.
+      "lea 8(%rsp), %rdx\n"
+      "jmp GThreadCreateImpl\n");
+}
+
+bool __attribute__((naked)) GThreadWake(struct GThread* thread) {
+  __asm__(
+      // mburakov: Argument would be overwritten by upcoming function calls.
+      "push %rdi\n"
+
+      // mburakov: Preserve registers of the current thread.
+      "mov %fs:g_current_thread@tpoff, %rdi\n"
+      "call GThreadSaveRegisters\n"
+
+      // mburakov: Preserve stack of the current thread.
+      "call GThreadSaveStack\n"
+      "test %al, %al\n"
+      "je escape_wake\n"
+
+      // mburakov: Recover registers of the new thread.
+      "mov (%rsp), %rdi\n"
+      "call GThreadRestoreRegisters\n"
+
+      // mburakov: Update prev and current thread pointers.
+      "mov %fs:g_current_thread@tpoff, %rsi\n"
+      "mov %rdi, %fs:g_current_thread@tpoff\n"
+      "mov %rsi, (%rdi)\n"
+
+      // mburakov: Recover stack of the new thread.
+      "call GThreadRestoreStack\n"
+      "mov $1, %eax\n"
+
+      "escape_wake:\n"
+      "add $8, %rsp\n"
+      "ret\n");
+}
+
+bool __attribute__((naked)) GThreadYield(void) {
+  __asm__(
+      // mburakov: Yielding on master thread has no effect.
+      "mov %fs:g_current_thread@tpoff, %rdi\n"
+      "mov (%rdi), %rax\n"
+      "test %rax, %rax\n"
+      "jne proceed\n"
+      "ret\n"
+      "proceed:\n"
+      "push %rax\n"
+
+      // mburakov: Preserve registers of the current thread.
+      "mov %fs:g_current_thread@tpoff, %rdi\n"
+      "call GThreadSaveRegisters\n"
+
+      // mburakov: Preserve stack of the current thread.
+      "call GThreadSaveStack\n"
+      "test %al, %al\n"
+      "je escape_yield\n"
+
+      // mburakov: Recover registers of the new thread.
+      "mov (%rsp), %rdi\n"
+      "call GThreadRestoreRegisters\n"
+
+      // mburakov: Update current thread pointer.
+      "mov %rdi, %fs:g_current_thread@tpoff\n"
+
+      // mburakov: Recover stack of the new thread.
+      "call GThreadRestoreStack\n"
+      "mov $1, %eax\n"
+
+      "escape_yield:\n"
+      "add $8, %rsp\n"
+      "ret\n");
+}
+
+void GThreadDestroy(struct GThread* thread) {
+  if (thread == &g_master_thread || thread == g_current_thread) {
+    // mburakov: It is prohibited to cancel either master or current thread.
+    abort();
+  }
+  GThreadFree(thread);
+  if (!--g_threads_count) {
+    // mburakov: No more non-master threads. Cleanup master thread stack data.
+    g_master_thread.stack.tip = 0;
+    free(g_master_thread.stack.data);
+    g_master_thread.stack.alloc = 0;
+  }
+}
diff --git a/gthread.h b/gthread.h
new file mode 100644
index 0000000..3cacd0b
--- /dev/null
+++ b/gthread.h
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2023 Mikhail Burakov. This file is part of toolbox.
+ *
+ * toolbox 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.
+ *
+ * toolbox 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 toolbox.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+#ifndef TOOLBOX_GTHREAD_H_
+#define TOOLBOX_GTHREAD_H_
+
+#include <stdbool.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif  // __cplusplus
+
+struct GThread;
+
+struct GThread* GThreadCreate(void (*proc)(void*), void* user);
+bool GThreadWake(struct GThread* thread);
+bool GThreadYield(void);
+void GThreadDestroy(struct GThread* thread);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif  // __cplusplus
+
+#endif  // TOOLBOX_GTHREAD_H_