MontaukOS/programs/include/montauk/thread.h

/*
    * thread.h
    * Userspace threading primitives for MontaukOS
    * Copyright (c) 2026 Daniel Hammer
    *
    * Threads share their parent's address space, IPC handles, cwd, and
    * heap allocator. The caller owns the thread stack: thread_spawn
    * allocates one out of the user heap and the kernel uses it verbatim.
    *
    * Lifetimes:
    *   - thread_spawn returns a positive TID on success.
    *   - thread_join blocks until the target TID exits, frees its kernel
    *     stack, then returns its exit code via *out_code.
    *   - thread_exit terminates only the current thread. If the main thread
    *     calls exit() (SYS_EXIT) the whole process tears down, killing any
    *     surviving sibling threads.
    *   - thread_self returns the calling thread's TID.
*/

#pragma once
#include <cstdint>
#include <Api/Syscall.hpp>
#include <montauk/syscall.h>
#include <montauk/heap.h>

namespace montauk {

    // Default per-thread stack size (64 KiB). Enough for typical app work;
    // matches the main thread's 32 KiB lower bound with headroom for
    // TrueType rendering call chains.
    static constexpr uint64_t DEFAULT_THREAD_STACK_BYTES = 64 * 1024;

    using ThreadEntry = int (*)(void* arg);

    // Terminate the calling thread. Never returns. If called by the main
    // thread this is equivalent to exit() (the whole process exits).
    [[noreturn]] inline void thread_exit(int code = 0) {
        syscall1(montauk::abi::SYS_THREAD_EXIT, (uint64_t)code);
        __builtin_unreachable();
    }

    namespace detail {
        struct ThreadCtx {
            ThreadEntry user_entry;
            void* user_arg;
            void* stack_base;
        };

        // Userspace trampoline: bridges from the raw entry the kernel jumps
        // to into the typed entry, then funnels into SYS_THREAD_EXIT. We
        // route the exit through libc rather than relying on a kernel-side
        // exit stub on the user stack -- the kernel never writes to user
        // memory on this path.
        //
        // The thread's stack itself is intentionally not freed here: we are
        // still running on it. It is reclaimed when the process exits, or
        // the joiner may free it explicitly after thread_join.
        [[noreturn]] inline void thread_trampoline(detail::ThreadCtx* ctx) {
            int code = ctx->user_entry(ctx->user_arg);
            montauk::mfree(ctx);
            thread_exit(code);
        }
    }

    // Spawn a new thread that begins executing `entry(arg)`. Returns the
    // new TID on success, or -1 on failure. The thread's stack is
    // allocated from the user heap; it is leaked on thread exit (the
    // thread itself cannot free the stack it is running on). The kernel
    // reclaims it on process exit. Callers that need to spawn many short-
    // lived threads should pool stacks themselves.
    inline int thread_spawn(ThreadEntry entry, void* arg,
                            uint64_t stack_bytes = 0) {
        if (entry == nullptr) return -1;
        if (stack_bytes == 0) stack_bytes = DEFAULT_THREAD_STACK_BYTES;
        stack_bytes = (stack_bytes + 15) & ~15ULL;

        void* stack = montauk::malloc(stack_bytes);
        if (stack == nullptr) return -1;
        auto* ctx = (detail::ThreadCtx*)montauk::malloc(sizeof(detail::ThreadCtx));
        if (ctx == nullptr) {
            montauk::mfree(stack);
            return -1;
        }
        ctx->user_entry = entry;
        ctx->user_arg = arg;
        ctx->stack_base = stack;

        uint64_t stack_top = ((uint64_t)stack + stack_bytes) & ~0xFULL;
        int tid = (int)syscall3(montauk::abi::SYS_THREAD_SPAWN,
                                (uint64_t)&detail::thread_trampoline,
                                (uint64_t)ctx, stack_top);
        if (tid < 0) {
            montauk::mfree(ctx);
            montauk::mfree(stack);
            return -1;
        }
        return tid;
    }

    // Block until the thread identified by `tid` terminates. Returns 0 on
    // success (with the thread's exit code in *out_code if non-null) or
    // -1 if `tid` is not a joinable sibling.
    inline int thread_join(int tid, int* out_code = nullptr) {
        return (int)syscall2(montauk::abi::SYS_THREAD_JOIN,
                             (uint64_t)tid, (uint64_t)out_code);
    }

    // Return the calling thread's TID (== getpid() for the main thread).
    inline int thread_self() {
        return (int)syscall0(montauk::abi::SYS_THREAD_SELF);
    }

    // Lightweight mutex backed by a single atomic word + the kernel
    // yield syscall. Adequate for short critical sections; a heavier
    // primitive can wrap mailbox_recv when blocking semantics are needed.
    struct Mutex {
        volatile uint32_t locked = 0;

        void lock() {
            while (__atomic_exchange_n(&locked, 1, __ATOMIC_ACQUIRE) != 0) {
                montauk::yield();
            }
        }

        bool try_lock() {
            return __atomic_exchange_n(&locked, 1, __ATOMIC_ACQUIRE) == 0;
        }

        void unlock() {
            __atomic_store_n(&locked, 0, __ATOMIC_RELEASE);
        }
    };

}