From 0ed31de7833d16e03846ad00f127d2c7f4c8f2a6 Mon Sep 17 00:00:00 2001
From: Aiden Gall <aiden@aidengall.xyz>
Date: Tue, 11 Feb 2025 12:45:45 +0000
Subject: (tco_args) add support for coroutine arguments

---
 tco.asm | 42 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

(limited to 'tco.asm')

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 ?
+	}
-- 
cgit v1.2.3