co-engine/co.asm 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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
format elf64 executable
include "syscall.inc"

macro enter {
    push rbp
    mov rbp, rsp
}

macro leave {
    mov rsp, rbp
    pop rbp
}

counter:
    enter
    sub rsp, 8
    mov QWORD [rsp], 0
.again:
    cmp QWORD [rsp], 10
    jge .over
    mov rdi, [rsp]
    call print
    call co_yield
    inc QWORD [rsp]
    jmp .again
.over:
    leave
    ret

entry _start
_start:
    ;call co_init

    mov rdi, counter
    call coroutine_go

    mov rdi, counter
    call coroutine_go

    mov rdi, counter
    call coroutine_go

    mov rdi, counter
    call coroutine_go

    mov rdi, counter
    call coroutine_go

.forever:
    call co_yield
    jmp .forever
    exit 101

co_exit:
    exit 69

co_init:
    pop rax                     ;; after call co_init, the stack top has rip
    enter
    mov QWORD [co_rip], rax
    mov QWORD [co_rbp], rbp
    mov QWORD [co_rsp], rsp
    inc QWORD [coroutines_count]
    leave
    push rax
    ret

coroutine_go:
    enter
    ;; ctx->rip = rip, the address of the parameter function pointer
    mov rsi, [coroutines_count]       ; load index
    mov QWORD [co_rip + rsi*8], rdi   ; store rdi into co_rip[rsi]

    ;; ctx->rbp = stack_base
    mov rbx, [stack_end]
    mov QWORD [co_rbp + rsi*8], rbx

    ;; since we push the return address on the stack immediatley,
    ;; ctx->rsp = ctx->rbp - 8
    sub rbx, 8
    mov QWORD [co_rsp + rsi*8], rbx
    mov QWORD [rbx], co_exit

    add rbx, 8
    sub rbx, CO_STACK_SIZE
    mov [stack_end], rbx

    ;; increment count
    inc QWORD [coroutines_count]
    leave
    ret


co_yield:
    ; save current progress
    pop rax
    mov rsi, [coroutines_current]       ; load _currnt_ index
    mov [co_rip + rsi*8], rax
    mov [co_rsp + rsi*8], rsp
    mov [co_rbp + rsi*8], rbp

    ; advance coroutines_current
    xor rcx, rcx
    inc rsi
    cmp rsi, [coroutines_count]
    cmovge rsi, rcx
    mov [coroutines_current], rsi

    mov rbp, [co_rbp + rsi*8]
    mov rsp, [co_rsp + rsi*8]
    mov rax, [co_rip + rsi*8]
    jmp rax

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

    ;; They are declared here, but accessed from the above code, this is
    ;; allowed by fasm, we can forward reference numerical constants

    ;; Assumption: we have at most 10 coroutines, each having a stack with
    ;; size CO_STACK_SIZE
    ;; coroutine count holds the number of allocated coroutines
CO_CAP        = 10
CO_STACK_SIZE = 1024
coroutines_count: dq 0

    ;; the current active coroutine is the one having control (rsp and rip
    ;; are set to it's values)
coroutines_current: dq 0

    ;; stack end points to the end of the allocation arena, and simplifies
    ;; stack memory allocation, so instead of doing
    ;; co_rsp + 8*count = coroutines_data + CO_STACK_SIZE * (co_count + 1)
    ;; we do
    ;; co_rsp + 8*count = stack_end
    ;; stack_end = stack_end - CO_STACK_SIZE
stack_end: dq (coroutines_data + CO_CAP * CO_STACK_SIZE)

    ;; this is the coroutine queue representation as an struct of arrays,
    ;; instead of array of structs (as I would have done it in C).
co_rbp:          rq CO_CAP
co_rsp:          rq CO_CAP
co_rip:          rq CO_CAP
coroutines_data: rq CO_CAP * CO_STACK_SIZE