const std = @import("std");
const builtin = @import("builtin");

// I'm working on a coroutine library.
// Zig seems to load threadlocals incorrectly when coroutines switch threads.
// I think it loads them too early then returns.

threadlocal var thread_local: usize = 7;

var main_context: [context_size]usize = std.mem.zeroes([context_size]usize);
var repro_context: [context_size]usize = std.mem.zeroes([context_size]usize);

pub fn main() !void {
    // Setup the context on one thread.
    const t = try std.Thread.spawn(.{}, thread, .{});
    t.join();

    // Run the context on another thread.
    switch_context(&main_context, &repro_context);
}

fn thread() void {
    setup_repro_context();
    switch_context(&main_context, &repro_context);
}

fn repro() void {
    // This collects the thread local address on the starting thread.
    const thread1_addr = &thread_local;
    // We context switch to suspend this execution
    switch_context(&repro_context, &main_context);
    // When we get back, we are on a different thread.
    // This should load a different address.
    // That is not happening.
    const thread2_addr = &thread_local;
    std.log.info("thread 1 addr: {*}", .{thread1_addr});
    std.log.info("Error -> thread 2 addr: {*}", .{thread2_addr});
    workaround();

    // Go back to main to exit the program.
    switch_context(&repro_context, &main_context);
}

noinline fn workaround() void {
    // Here it does work due to to separate function causing the address to actually be loaded.
    const thread2_addr = &thread_local;
    std.log.info("Workaround -> thread 2 addr: {*}", .{thread2_addr});
}

// Put everything inline for simplicity.
// This only works on aarch64 and x86_64 unix.
extern fn switch_context(current: [*]u64, target: [*]u64) void;
comptime {
    switch (builtin.cpu.arch) {
        .aarch64 => {
            asm (
                \\ .global _switch_context
                \\ _switch_context:
                \\ .global switch_context
                \\ switch_context:
                \\ 
                \\ # Save all callee saved registers to current context.
                \\ mov x9, sp
                \\ # General
                \\ stp x19, x20, [x0, 0*8]
                \\ stp x21, x22, [x0, 2*8]
                \\ stp x23, x24, [x0, 4*8]
                \\ stp x25, x26, [x0, 6*8]
                \\ stp x27, x28, [x0, 8*8]
                \\ # Double
                \\ stp d8,  d9,  [x0, 10*8]
                \\ stp d10, d11, [x0, 12*8]
                \\ stp d12, d13, [x0, 14*8]
                \\ stp d14, d15, [x0, 16*8]
                \\ # Special
                \\ stp fp,  lr,  [x0, 18*8]
                \\ str x9,       [x0, 20*8]
                \\ 
                \\ 
                \\ # Load target context.
                \\ # General
                \\ ldp x19, x20, [x1, 0*8]
                \\ ldp x21, x22, [x1, 2*8]
                \\ ldp x23, x24, [x1, 4*8]
                \\ ldp x25, x26, [x1, 6*8]
                \\ ldp x27, x28, [x1, 8*8]
                \\ # Double
                \\ ldp d8,  d9,  [x1, 10*8]
                \\ ldp d10, d11, [x1, 12*8]
                \\ ldp d12, d13, [x1, 14*8]
                \\ ldp d14, d15, [x1, 16*8]
                \\ # Special
                \\ ldp fp,  lr,  [x1, 18*8]
                \\ ldr x9,       [x1, 20*8]
                \\ mov sp, x9
                \\ 
                \\ ret
            );
        },
        .x86_64 => {
            if (builtin.os.tag == .windows) {
                @compileError("Unsupported os");
            }
            asm (
                \\ .global _switch_context
                \\ _switch_context:
                \\ .global switch_context
                \\ switch_context:
                \\ 
                \\ # Save all callee saved registers to current context.
                \\ # General
                \\ movq %r12, 0x00(%rdi)
                \\ movq %r13, 0x08(%rdi)
                \\ movq %r14, 0x10(%rdi)
                \\ movq %r15, 0x18(%rdi)
                \\ movq %rbx, 0x20(%rdi)
                \\ # Special
                \\ movq %rbp, 0x28(%rdi)
                \\ movq %rsp, 0x30(%rdi)
                \\ 
                \\ 
                \\ # Load target context.
                \\ # General
                \\ movq 0x00(%rsi), %r12
                \\ movq 0x08(%rsi), %r13
                \\ movq 0x10(%rsi), %r14
                \\ movq 0x18(%rsi), %r15
                \\ movq 0x20(%rsi), %rbx
                \\ # Special
                \\ movq 0x28(%rsi), %rbp
                \\ movq 0x30(%rsi), %rsp
                \\ 
                \\ ret
            );
        },
        else => @compileError("Unsupported cpu architecture"),
    }
}

const context_size =
    switch (builtin.cpu.arch) {
    .aarch64 => 21,
    .x86_64 => 7,
    else => @compileError("Unsupported cpu architecture"),
};

const STACK_ALIGN = 16;
var repro_stack: [16 * 1024 * 1024]u8 = std.mem.zeroes([16 * 1024 * 1024]u8);

fn setup_repro_context() void {
    var sp = @as([*]u8, @ptrCast(&repro_stack)) + repro_stack.len;

    switch (builtin.cpu.arch) {
        .aarch64 => {
            // Ensure stack is aligned.
            sp = @ptrFromInt(@intFromPtr(sp) & ~@as(usize, (STACK_ALIGN - 1)));

            const frame_pointer_index = 18;
            const return_pointer_index = 19;
            const stack_pointer_index = 20;
            repro_context[stack_pointer_index] = @intFromPtr(sp);
            repro_context[frame_pointer_index] = @intFromPtr(sp);
            repro_context[return_pointer_index] = @intFromPtr(&repro);
        },
        .x86_64 => {
            // Makes space to store the return address on the stack.
            sp -= @sizeOf(usize);
            // Ensure stack is aligned.
            sp = @ptrFromInt(@intFromPtr(sp) & ~@as(usize, (STACK_ALIGN - 1)));

            const return_address_ptr = @as(*usize, @alignCast(@ptrCast(sp)));
            return_address_ptr.* = @intFromPtr(&repro);

            const frame_pointer_index = 5;
            const stack_pointer_index = 6;
            repro_context[stack_pointer_index] = @intFromPtr(sp);
            repro_context[frame_pointer_index] = @intFromPtr(sp);
        },
        else => @compileError("Unsupported cpu architecture"),
    }
}