path: root/gthread.c

/*
 * 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;
  }
}