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

---
 Makefile   |  1 +
 tco.asm    | 42 +++++++++++++++++++++++++++++++++++++++++-
 tco.h      |  3 ++-
 tco_args.3 |  1 +
 tco_go.3   | 36 +++++++++++++++++++++++++++---------
 5 files changed, 72 insertions(+), 11 deletions(-)
 create mode 100644 tco_args.3

diff --git a/Makefile b/Makefile
index 773c346..084c41b 100644
--- a/Makefile
+++ b/Makefile
@@ -41,4 +41,5 @@ install_hdr:
 
 install_man:
 	install -Dm644 tco_go.3 ${DESTDIR}${MANPREFIX}/man3/tco_go.3
+	install -Dm644 tco_args.3 ${DESTDIR}${MANPREFIX}/man3/tco_args.3
 	install -Dm644 tco_yield.3 ${DESTDIR}${MANPREFIX}/man3/tco_yield.3
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 ?
+	}
diff --git a/tco.h b/tco.h
index c8a2f3d..897f2fa 100644
--- a/tco.h
+++ b/tco.h
@@ -19,5 +19,6 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>. */
 #define coroutine
 #endif
 
-int tco_go(void (*f)(void));
+int tco_go(void (*f)());
+void tco_args();
 void tco_yield(void);
diff --git a/tco_args.3 b/tco_args.3
new file mode 100644
index 0000000..344116c
--- /dev/null
+++ b/tco_args.3
@@ -0,0 +1 @@
+.so man3/tco_go.3
diff --git a/tco_go.3 b/tco_go.3
index 531879b..ce7355d 100644
--- a/tco_go.3
+++ b/tco_go.3
@@ -1,6 +1,6 @@
 .TH tco_go 3 2025-02-06
 .SH NAME
-tco_go, tco_yield
+tco_go, tco_args, tco_yield
 .SH LIBRARY
 Tiny Coroutines
 .RI ( libtco ", " \-ltco )
@@ -8,7 +8,8 @@ Tiny Coroutines
 .nf
 .B #include <tco.h>
 .P
-.BI "int tco_go(void (*f)(void));"
+.BI "int tco_go(void (*" f ")(...));"
+.B "int tco_args(...);"
 .BI "void tco_yield(void);"
 .fi
 .SH DESCRIPTION
@@ -19,17 +20,19 @@ The coroutine function pointer
 .I f
 is required to not be inlined
 (the
-.BR coroutine
+.B coroutine
 macro is provided as a convienient way to use the
-.BR gcc (1)
-/
-.BR clang (1)
-.BR noinline
+.BR gcc (1)/ clang (1)
+.B noinline
 attribute, and to communicate the intented use of a function as a coroutine).
 When the scheduler switches back to a context that has called
 .BR tco_go (),
 it resumes execution from after that call.
 .P
+.BR tco_args ()
+sets the function arguments given to the coroutine spawned by
+.BR tco_go ().
+.P
 .BR tco_yield ()
 saves the current execution context and switches to the next scheduled context,
 pausing execution of the current coroutine.
@@ -43,8 +46,23 @@ returns 0 on success, and a non-zero error code when the coroutine fails to
 initialise.
 .SH CAVEATS
 This library works by preserving the state of the callee-saved registers and
-restoring them later when the coroutine is resumed. This approach may be
-incompatible with other forms of black magic voodoo, your mileage may vary.
+restoring them later when the coroutine is resumed.
+This approach may be incompatible with other forms of black magic voodoo, your
+mileage may vary.
+.P
+.BR tco_args ()
+currently only supports arguments passed by registers
+.RI ( %rdi ,
+.IR %rsi ,
+.IR %rdx ,
+.IR %r10 ,
+.IR %r8 ,
+.I %r9
+and
+.IR %xmm0\-8 ).
+Spawning coroutines that take arguments on the stack will result in stack
+corruption.
+Refer to the System V ABI for more details.
 .SH AUTHOR AND LICENSE
 Copyright (C) 2025 Aiden Gall
 .P
-- 
cgit v1.2.3