co-engine/co.c at master · Kariiem/co-engine · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
#include <stdio.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <setjmp.h>

#define CO_STACK_SIZE 1024

typedef struct coroutine_context coroutine_context;
struct coroutine_context {
    char *stack_mem;
    void *rip;
    void *rbp;
    void *rsp;
};

#define CO_CAP 10
static coroutine_context contexts[CO_CAP];
static uint32_t count = 0;
static uint32_t current = 0;

void co_exit()
{
    exit(69);
    free(contexts[current].stack_mem - CO_STACK_SIZE);
    /* TODO remove coroutine from queue */
}

__attribute__((naked))
void co_init(void)
{
    __asm__(
        "push %%rbp\n\t"              // Save rbp to maintain stack frame
        "mov %%rsp, %%rbp\n\t"        // Set up new base pointer
        "pop %%rax\n\t"               // Pop return address into rax
        "mov %%rax, %0\n\t"           // Store return address in contexts[0].rip
        "mov %%rbp, %1\n\t"           // Store rbp in contexts[0].rbp
        "mov %%rsp, %2\n\t"           // Store rsp in contexts[0].rsp
        "push %%rax\n\t"              // Push return address back
        "incl %3\n\t"                 // Increment count
        "leave\n\t"                   // Restore rsp from rbp and pop rbp
        "ret\n\t"
        : "=m" (contexts[0].rip),
          "=m" (contexts[0].rbp),
          "=m" (contexts[0].rsp),
          "=m" (count)
        :
        : "rax"
    );
}

void co_go(void (*fn)(void))
{
    coroutine_context *ctx = &contexts[count++];
    ctx->rip = fn;
    ctx->stack_mem = malloc(CO_STACK_SIZE) + CO_STACK_SIZE;
    ctx->rsp = ctx->stack_mem - 8;
    *(uintptr_t *)ctx->rsp = (uintptr_t)co_exit;
}

__attribute__((noreturn, naked))
void co_yield(void)
{
    /* Save current progress */
    __asm__(
        "pop    %%rax\n\t"
        "mov    %%rax, %0\n\t"
        "mov    %%rsp, %1\n\t"
        "mov    %%rbp, %2\n\t"
        : "=m" (contexts[current].rip),
          "=m" (contexts[current].rsp),
          "=m" (contexts[current].rbp)
        :
        : "rax"
    );

    /* Update current: current = (current + 1) % count */
    __asm__(
        "mov    %0, %%rax\n\t"        // Load current into rax
        "inc    %%rax\n\t"            // Increment current
        "xor    %%rdx, %%rdx\n\t"     // Clear rdx for division
        "divl   %1\n\t"               // Divide rax by count, remainder in rdx
        "mov    %%rdx, %0\n\t"        // Store remainder (modulo) back to current
        : "+m" (current)
        : "m" (count)
        : "rax", "rdx"
    );

    /* Restore next coroutine */
    __asm__(
        "mov %0, %%rax\n\t"
        "mov %1, %%rsp\n\t"
        "mov %2, %%rbp\n\t"
        "jmp *%%rax\n\t"
        :
        : "r" (contexts[current].rip),
          "r" (contexts[current].rsp),
          "r" (contexts[current].rbp)
        : "rax"
    );
    __builtin_unreachable();
}

int co_alive(void)
{
    return count > 1;
}

void counter(void)
{
    for (int i = 0; i < 10; ++i) {
        printf("%d\n", i);
        co_yield();
    }
}

int main(void)
{
    co_init();
    co_go(counter);
    co_go(counter);
    co_yield();
    return 0;
}