wip

1 files changed, 40 insertions, 47 deletions

diff --git a/coroutine.c b/coroutine.c
index 65971d4..fe6f5be 100644
--- a/coroutine.c
+++ b/coroutine.c
@@ -5,8 +5,8 @@
  */
 
 #include <stdint.h> /* for uint8_t */
-#include <string.h> /* for memset() */
 #include <stdio.h>  /* for fprintf(), stderr */
+#include <stdlib.h> /* for calloc(), free() */
 #include <assert.h>
 #include <setjmp.h>
 
@@ -14,8 +14,7 @@
 #include "coroutine.h"
 
 enum coroutine_state {
-	CR_ZERO = 0,
-	CR_GARBAGE,
+	CR_NONE = 0,
 	CR_INITIALIZING,
 	CR_RUNNING,
 	CR_RUNNABLE,
@@ -25,48 +24,31 @@ enum coroutine_state {
 struct coroutine {
 	enum coroutine_state state;
 	jmp_buf              env;
-	uint8_t              stack[COROUTINE_STACK_SIZE];
+	size_t               stack_size;
+	void                *stack;
 };
 
 static struct coroutine coroutine_table[COROUTINE_NUM] = {0};
 static cid_t            coroutine_running              = 0;
-static jmp_buf          scheduler_env;
+static jmp_buf          coroutine_add_env;
+static jmp_buf          coroutine_main_env;
 
 void coroutine_init(void) {}
 
-void *__sp;
-#define SP ({ asm("movq %%rsp,%0" : "=r"(__sp)); __sp; })
-
-static void print_addr(char *name, void *addr) {
-	if ((void*)&coroutine_table <= addr && addr < ((void*)&coroutine_table) + sizeof(coroutine_table)) {
-		for (size_t i = 0; i < COROUTINE_NUM; i++) {
-			if ((void*)&coroutine_table[i] <= addr && addr < ((void*)&coroutine_table[i]) + sizeof(coroutine_table[i])) {
-				printf("%s=%p (in coroutine_table[%zu] [%p, %p])\n", name, addr, i,
-				       (void*)&coroutine_table[i], ((void*)&coroutine_table[i]) + sizeof(coroutine_table[i]));
-				break;
-			}
-		}
-	} else {
-		printf("%s=%p (outside coroutine_table [%p, %p])\n", name, addr,
-		       (void*)&coroutine_table, ((void*)&coroutine_table) + sizeof(coroutine_table));
-	}
-}
-
 static void call_with_stack(void *stack, cr_fn_t fn, void *args) {
-	static void  *scheduler_sp = NULL;
+	static void *saved_sp = NULL;
 
-	/* I don't intend to run nThis only really exists for running on ARM-32, but being
+	/* This only really exists for running on ARM-32, but being
 	 * able to run it on x86-64 is useful for debugging.  */
 #if __x86_64__
 #define STACK_GROWS_DOWNWARD 1
-	print_addr("sp", SP);
-	asm volatile ("movq %%rsp , %0\n\t"    /* scheduler_sp = sp */
+	asm volatile ("movq %%rsp , %0\n\t"    /* saved_sp = sp */
 	              "movq %1    , %%rsp\n\t" /* sp = stack */
 	              "movq %3    , %%rdi\n\t" /* arg0 = args */
 	              "call *%2\n\t"           /* fn() */
-	              "movq %0    , %%rsp"     /* sp = scheduler_sp */
+	              "movq %0    , %%rsp"     /* sp = saved_sp */
 	              :
-	              : /* %0 */"m"(scheduler_sp),
+	              : /* %0 */"m"(saved_sp),
 	                /* %1 */"r"(stack),
 	                /* %2 */"r"(fn),
 	                /* %3 */"r"(args)
@@ -87,35 +69,39 @@ static void call_with_stack(void *stack, cr_fn_t fn, void *args) {
 cid_t coroutine_add(cr_fn_t fn, void *args) {
 	static cid_t last_created = 0;
 
+	assert(coroutine_running == 0 || coroutine_table[coroutine_running-1].state == CR_RUNNING);
+
 	cid_t child;
 	{
 		size_t idx_base = last_created;
 		for (size_t idx_shift = 0; idx_shift < COROUTINE_NUM; idx_shift++) {
 			child = ((idx_base + idx_shift) % COROUTINE_NUM) + 1;
-			if (coroutine_table[child-1].state == CR_ZERO || coroutine_table[child-1].state == CR_GARBAGE)
+			if (coroutine_table[child-1].state == CR_NONE)
 				goto found;
 		}
 		return 0;
 	found:
 	}
 
-	if (coroutine_table[child-1].state == CR_GARBAGE)
-		memset(&coroutine_table[child-1], 0, sizeof(struct coroutine));
 	last_created = child;
 
+	coroutine_table[child-1].stack_size = COROUTINE_STACK_SIZE;
+	coroutine_table[child-1].stack = calloc(1, coroutine_table[child-1].stack_size);
+
 	cid_t parent = coroutine_running;
 	assert(parent == 0 || coroutine_table[parent-1].state == CR_RUNNING);
 	coroutine_running = child; 
 	coroutine_table[child-1].state = CR_INITIALIZING;
-	if (child == 2) {
-		print_addr("&child", &child);
-	}
-	if (!setjmp(scheduler_env)) {
-		call_with_stack(coroutine_table[child-1].stack + (STACK_GROWS_DOWNWARD ? COROUTINE_STACK_SIZE : 0), fn, args);
-		assert(coroutine_table[child-1].state == CR_INITIALIZING);
-		coroutine_table[child-1].state = CR_RUNNABLE;
+	if (!setjmp(coroutine_add_env)) { /* point=a */
+		/* run until cr_begin() */
+		call_with_stack(coroutine_table[child-1].stack + (STACK_GROWS_DOWNWARD ? coroutine_table[child-1].stack_size : 0), fn, args);
+		/* cr_begin() calls longjmp(point=a); if fn returns
+		 * then that means it didn't call cr_begin(), which is
+		 * wrong.  */
+		assert(false);
 	}
 	assert(coroutine_table[child-1].state == CR_RUNNABLE);
+	assert(parent == 0 || coroutine_table[parent-1].state == CR_RUNNING);
 	coroutine_running = parent;
 
 	return child;
@@ -129,9 +115,14 @@ void coroutine_main(void) {
 		if (coroutine_table[coroutine_running-1].state == CR_RUNNABLE) {
 			ran = true;
 			coroutine_table[coroutine_running-1].state = CR_RUNNING;
-			if (!setjmp(scheduler_env)) {
-				longjmp(coroutine_table[coroutine_running-1].env, 1);
-				coroutine_table[coroutine_running-1].state = CR_GARBAGE;
+			if (!setjmp(coroutine_main_env)) { /* point=b */
+				longjmp(coroutine_table[coroutine_running-1].env, 1); /* jump to point=c */
+				/* Consider returning to be the same as cr_exit(). */
+				coroutine_table[coroutine_running-1].state = CR_NONE;
+			}
+			if (coroutine_table[coroutine_running-1].state == CR_NONE) {
+				free(coroutine_table[coroutine_running-1].stack);
+				coroutine_table[coroutine_running-1] = (struct coroutine){0};
 			}
 		}
 		if (coroutine_running == COROUTINE_NUM+1 && !ran) {
@@ -141,19 +132,21 @@ void coroutine_main(void) {
 	}
 }
 
-bool _cr_begin(void) {
+bool cr_begin(void) {
 	assert(coroutine_table[coroutine_running-1].state == CR_INITIALIZING);
-	return setjmp(coroutine_table[coroutine_running-1].env);
+	coroutine_table[coroutine_running-1].state = CR_RUNNABLE;
+	if (!setjmp(coroutine_table[coroutine_running-1].env)) /* point=c1 */
+		longjmp(coroutine_add_env, 1); /* jump to point=a */
 }
 
 static inline void _cr_transition(enum coroutine_state state) {
 	assert(coroutine_table[coroutine_running-1].state == CR_RUNNING);
 	coroutine_table[coroutine_running-1].state = state;
-	if (state == CR_GARBAGE || !setjmp(coroutine_table[coroutine_running-1].env))
-		longjmp(scheduler_env, 1);
+	if (state == CR_NONE || !setjmp(coroutine_table[coroutine_running-1].env))  /* point=c2 */
+		longjmp(coroutine_main_env, 1); /* jump to point=b */
 }
 
-void cr_exit(void)            { _cr_transition(CR_GARBAGE); }
+void cr_exit(void)            { _cr_transition(CR_NONE); }
 void cr_yield(void)           { _cr_transition(CR_RUNNABLE); }
 void cr_pause_and_yield(void) { _cr_transition(CR_PAUSED); }