(tco_args) add support for coroutine arguments

1 files changed, 41 insertions, 1 deletions

diff --git a/tco.asm b/tco.asm
index 082d353..a26de62 100644
--- a/tco.asm
+++ b/tco.asm
@@ -16,6 +16,7 @@
 format ELF64
 
 public tco_go
+public tco_args
 public tco_yield
 
 ; assembly-time configuration options
@@ -57,7 +58,7 @@ else if MALLOC eq mmap
 end if
 
 section '.text' executable
-; int tco_go(void (*f)(void))
+; int tco_go(void (*f)(...))
 ; spawns a coroutine
 tco_go:
 	call stash
@@ -127,6 +128,18 @@ end if
 	lea rdx, [deinit]
 	push rdx
 
+	; restore coroutine arguments
+	irps reg, rax rdi rsi rdx r10 r8 r9 {
+		mov reg, [coroutine_args.#reg]
+	}
+
+	test rax, rax
+	jz @f
+
+	rept 8 n:0 {
+		movaps xmm#n, [coroutine_args.xmm#n]
+	}
+@@:
 	jmp rcx
 
 .oom:
@@ -139,6 +152,22 @@ else if MALLOC in <aligned_alloc,mmap>
 end if
 	ret
 
+; void tco_args(...)
+; sets the arguments of the next coroutine
+tco_args:
+	irps reg, rax rdi rsi rdx r10 r8 r9 {
+		mov [coroutine_args.#reg], reg
+	}
+
+	test rax, rax
+	jz @f
+
+	rept 8 n:0 {
+		movaps [coroutine_args.xmm#n], xmm#n
+	}
+@@:
+	ret
+
 ; void tco_yield(void)
 ; yield to next coroutine
 tco_yield:
@@ -229,3 +258,14 @@ section '.data' writeable
 
 	current_ctx_ptr dq root_ctx
 	prev_ctx_ptr dq root_ctx
+
+section '.bss' writeable align 16
+coroutine_args:
+	rept 8 n:0 {
+		label .xmm#n dqword
+		rq 2
+	}
+
+	irps reg, rax rdi rsi rdx r10 r8 r9 {
+		.#reg dq ?
+	}