libcr: Faster context switches, wait for interrupt

2 files changed, 151 insertions, 102 deletions

diff --git a/libcr/coroutine.c b/libcr/coroutine.c
index 9ddfa00..d221ca7 100644
--- a/libcr/coroutine.c
+++ b/libcr/coroutine.c
@@ -4,13 +4,20 @@
  * SPDX-Licence-Identifier: AGPL-3.0-or-later
  */
 
+#include <assert.h>
+#include <setjmp.h> /* for setjmp(), longjmp(), jmp_buf */
 #include <stdint.h> /* for uint8_t */
 #include <stdio.h>  /* for printf(), fprintf(), stderr */
 #include <stdlib.h> /* for aligned_alloc(), free() */
-#include <assert.h>
-#include <setjmp.h>
 
-#include "libcr/coroutine.h"
+#if __x86_64__
+#  include <unistd.h> /* for pause() */
+#elif __arm__
+#  include "hardware/sync.h" /* for __wfi(); */
+#  define pause() __wfi()
+#endif
+
+#include <libcr/coroutine.h>
 
 /* Configuration **************************************************************/
 
@@ -110,23 +117,9 @@
  *   coroutine_main()), and besides being
  *   wrong-for-our-desired-flow-control, that's a stack location that
  *   no longer exists.
- *
- * Things to consider changing:
- *
- * - Consider having _cr_transition() go ahead and find the next
- *   coroutine to run and longjmp() direcly to it, instead of first
- *   jumping back to coroutine_main().  This could save a few cycles
- *   and a few bytes.
  */
 
-#if CONFIG_COROUTINE_DEBUG
-#	define debugf(...) printf("dbg: " __VA_ARGS__)
-#else
-#	define debugf(...)
-#endif
-
-static jmp_buf coroutine_add_env;
-static jmp_buf coroutine_main_env;
+/* types **********************************************************************/
 
 enum coroutine_state {
 	CR_NONE = 0,     /* this slot in the table is empty */
@@ -136,6 +129,29 @@ enum coroutine_state {
 	CR_PAUSED,       /* not running, and not runnable */ 
 };
 
+struct coroutine {
+	volatile enum coroutine_state  state;
+	volatile bool                  sig_unpause;
+	jmp_buf                        env;
+	size_t                         stack_size;
+	void                          *stack;
+};
+
+/* constants ******************************************************************/
+
+#define STACK_ALIGNMENT ({ __attribute__((aligned)) void fn(void) {}; __alignof__(fn); })
+
+#if CONFIG_COROUTINE_MEASURE_STACK || CONFIG_COROUTINE_PROTECT_STACK
+/* We just need a pattern that is unlikely to occur naturaly; this is
+ * just a few bytes that I read from /dev/random.  */
+static const uint8_t stack_pattern[] = {0x1e, 0x15, 0x16, 0x0a, 0xcc, 0x52, 0x7e, 0xb7};
+#endif
+
+/* global variables ***********************************************************/
+
+static jmp_buf coroutine_add_env;
+static jmp_buf coroutine_main_env;
+
 /*
  * Invariants (and non-invariants):
  *
@@ -153,17 +169,42 @@ enum coroutine_state {
  *   coroutine as CR_RUNNING.
  */
 
-struct coroutine {
-	volatile enum coroutine_state  state;
-	volatile bool                  sig_unpause;
-	jmp_buf                        env;
-	size_t                         stack_size;
-	void                          *stack;
-};
-
 static struct coroutine coroutine_table[CONFIG_COROUTINE_NUM] = {0};
 static cid_t            coroutine_running                     = 0;
 
+/* utility functions **********************************************************/
+
+#if CONFIG_COROUTINE_DEBUG
+#	define debugf(...) printf("dbg: " __VA_ARGS__)
+#else
+#	define debugf(...)
+#endif
+
+static inline void assert_cid(cid_t cid) {
+	assert(cid > 0);
+	assert(cid <= CONFIG_COROUTINE_NUM);
+#if CONFIG_COROUTINE_PROTECT_STACK
+	assert(coroutine_table[cid-1].stack_size);
+	assert(coroutine_table[cid-1].stack);
+	for (size_t i = 0; i < sizeof(stack_pattern); i++) {
+		size_t j = coroutine_table[cid-1].stack_size - (i+1);
+		assert(((uint8_t*)coroutine_table[cid-1].stack)[i] == stack_pattern[i]);
+		assert(((uint8_t*)coroutine_table[cid-1].stack)[j] == stack_pattern[j%sizeof(stack_pattern)]);
+	}
+#endif
+}
+
+#define assert_cid_state(cid, opstate) do {                     \
+		assert_cid(cid);                                \
+		assert(coroutine_table[(cid)-1].state opstate); \
+	} while (0)
+
+
+/* Return `n` rounded up to the nearest multiple of `d` */
+#define round_up(n, d) ( ( ((n)+(d)-1) / (d) ) * (d) )
+
+/* call_with_stack() **********************************************************/
+
 static void call_with_stack(void *stack, cr_fn_t fn, void *args) {
 	static void *saved_sp = NULL;
 
@@ -207,40 +248,7 @@ static void call_with_stack(void *stack, cr_fn_t fn, void *args) {
 #  error unsupported architecture
 #endif
 }
-
-#if CONFIG_COROUTINE_MEASURE_STACK || CONFIG_COROUTINE_PROTECT_STACK
-/* We just need a pattern that is unlikely to occur naturaly; this is
- * just a few bytes that I read from /dev/random.  */
-static const uint8_t stack_pattern[] = {0x1e, 0x15, 0x16, 0x0a, 0xcc, 0x52, 0x7e, 0xb7};
-#endif
-
-static inline void assert_cid(cid_t cid) {
-	assert(cid > 0);
-	assert(cid <= CONFIG_COROUTINE_NUM);
-#if CONFIG_COROUTINE_PROTECT_STACK
-	assert(coroutine_table[cid-1].stack_size);
-	assert(coroutine_table[cid-1].stack);
-	for (size_t i = 0; i < sizeof(stack_pattern); i++) {
-		size_t j = coroutine_table[cid-1].stack_size - (i+1);
-		assert(((uint8_t*)coroutine_table[cid-1].stack)[i] == stack_pattern[i]);
-		assert(((uint8_t*)coroutine_table[cid-1].stack)[j] == stack_pattern[j%sizeof(stack_pattern)]);
-	}
-#endif
-}
-
-#define assert_cid_state(cid, opstate) do {                     \
-		assert_cid(cid);                                \
-		assert(coroutine_table[(cid)-1].state opstate); \
-	} while (0)
-
-cid_t coroutine_add(cr_fn_t fn, void *args) {
-	return coroutine_add_with_stack_size(CONFIG_COROUTINE_DEFAULT_STACK_SIZE, fn, args);
-}
-
-#define STACK_ALIGNMENT ({ __attribute__((aligned)) void fn(void) {}; __alignof__(fn); })
-
-/* Return `n` rounded up to the nearest multiple of `d` */
-#define round_up(n, d) ( ( ((n)+(d)-1) / (d) ) * (d) )
+/* coroutine_add() ************************************************************/
 
 cid_t coroutine_add_with_stack_size(size_t stack_size, cr_fn_t fn, void *args) {
 	static cid_t last_created = 0;
@@ -250,7 +258,7 @@ cid_t coroutine_add_with_stack_size(size_t stack_size, cr_fn_t fn, void *args) {
 		assert_cid_state(parent, == CR_RUNNING);
 	assert(stack_size);
 	assert(fn);
-	debugf("coroutine_add_with_stack_size(%zu, %#p, %#p)...\n", stack_size, fn, args);
+	debugf("coroutine_add_with_stack_size(%zu, %p, %p)...\n", stack_size, fn, args);
 
 	cid_t child;
 	{
@@ -285,8 +293,8 @@ cid_t coroutine_add_with_stack_size(size_t stack_size, cr_fn_t fn, void *args) {
 		stack_base += round_up(sizeof(stack_pattern), STACK_ALIGNMENT);
 #  endif
 #endif
-		debugf("...stack     =%#p\n", coroutine_table[child-1].stack);
-		debugf("...stack_base=%#p\n", stack_base);
+		debugf("...stack     =%p\n", coroutine_table[child-1].stack);
+		debugf("...stack_base=%p\n", stack_base);
 		/* run until cr_begin() */
 		call_with_stack(stack_base, fn, args);
 		assert(false); /* should cr_begin() instead of returning */
@@ -299,47 +307,71 @@ cid_t coroutine_add_with_stack_size(size_t stack_size, cr_fn_t fn, void *args) {
 	return child;
 }
 
-void coroutine_main(void) {
-	debugf("coroutine_main()\n");
-	bool ran = false;
-	for (coroutine_running = 1;; coroutine_running = (coroutine_running%CONFIG_COROUTINE_NUM)+1) {
-		if (coroutine_running == 1)
-			ran = false;
-		struct coroutine *cr = &coroutine_table[coroutine_running-1];
-		if (cr->state == CR_RUNNABLE) {
-			debugf("running cid=%zu...\n", coroutine_running);
-			cr->state = CR_RUNNING;
-			if (!setjmp(coroutine_main_env)) { /* point=b */
-				longjmp(cr->env, 1); /* jump to point=c */
-				assert(false); /* should cr_exit() instead of returning */
-			}
-			ran = true;
-			assert_cid_state(coroutine_running, != CR_RUNNING);
-			if (cr->state == CR_NONE) {
+cid_t coroutine_add(cr_fn_t fn, void *args) {
+	return coroutine_add_with_stack_size(CONFIG_COROUTINE_DEFAULT_STACK_SIZE, fn, args);
+}
+
+/* coroutine_main() ***********************************************************/
+
 #if CONFIG_COROUTINE_MEASURE_STACK
-				size_t stack_size = cr->stack_size - (CONFIG_COROUTINE_PROTECT_STACK ? 2*sizeof(stack_pattern) : 0);
-				size_t stack_used = stack_size;
-				for (;;) {
-					size_t i = STACK_GROWS_DOWNWARD
-						? (CONFIG_COROUTINE_PROTECT_STACK ? sizeof(stack_pattern) : 0) + stack_size - stack_used
-						: stack_used - 1 - (CONFIG_COROUTINE_PROTECT_STACK ? sizeof(stack_pattern) : 0);
-					if (stack_used == 0 || ((uint8_t*)cr->stack)[i] != stack_pattern[i%sizeof(stack_pattern)])
-						break;
-					stack_used--;
-				}
-				printf("info: cid=%zu: exited having used %zu B stack space\n", coroutine_running, stack_used);
+struct stack_stats {
+	size_t cap;
+	size_t max;
+	//size_t cur;
+};
+
+static void measure_stack(cid_t cid, struct stack_stats *ret) {
+	ret->cap = coroutine_table[cid-1].stack_size - (CONFIG_COROUTINE_PROTECT_STACK ? 2*round_up(sizeof(stack_pattern), STACK_ALIGNMENT) : 0);
+
+	ret->max = ret->cap;
+	for (;;) {
+		size_t i = STACK_GROWS_DOWNWARD
+			? (CONFIG_COROUTINE_PROTECT_STACK ? round_up(sizeof(stack_pattern), STACK_ALIGNMENT) : 0) + ret->cap - ret->max
+			: ret->max - 1 - (CONFIG_COROUTINE_PROTECT_STACK ? round_up(sizeof(stack_pattern), STACK_ALIGNMENT) : 0);
+		if (ret->max == 0 || ((uint8_t*)coroutine_table[cid-1].stack)[i] != stack_pattern[i%sizeof(stack_pattern)])
+			break;
+		ret->max--;
+	}
+}
 #endif
-				free(cr->stack);
-				coroutine_table[coroutine_running-1] = (struct coroutine){0};
-			}
-		}
-		if (coroutine_running == CONFIG_COROUTINE_NUM && !ran) {
-			fprintf(stderr, "error: no runnable coroutines\n");
+
+static inline cid_t next_coroutine() {
+	for (cid_t next = (coroutine_running%CONFIG_COROUTINE_NUM)+1;
+	     next != coroutine_running;
+	     next = (next%CONFIG_COROUTINE_NUM)+1) {
+		if (coroutine_table[next-1].state == CR_RUNNABLE)
+			return next;
+	}
+	return 0;
+}
+
+void coroutine_main(void) {
+	debugf("coroutine_main()\n");
+	coroutine_running = 0;
+	for (;;) {
+		cid_t next = next_coroutine();
+		if (!next) {
+			fprintf(stderr, "error: no coroutines\n");
 			return;
 		}
+		if (!setjmp(coroutine_main_env)) { /* point=b */
+			coroutine_running = next;
+			coroutine_table[coroutine_running-1].state = CR_RUNNING;
+			longjmp(coroutine_table[coroutine_running-1].env, 1); /* jump to point=c */
+		}
+		assert_cid_state(coroutine_running, == CR_NONE);
+#if CONFIG_COROUTINE_MEASURE_STACK
+		struct stack_stats sizes;
+		measure_stack(coroutine_running, &sizes);
+		printf("info: cid=%zu: exited having used %zu B stack space\n", coroutine_running, sizes.max);
+#endif
+		free(coroutine_table[coroutine_running-1].stack);
+		coroutine_table[coroutine_running-1] = (struct coroutine){0};
 	}
 }
 
+/* cr_*() *********************************************************************/
+
 void cr_begin(void) {
 	assert_cid_state(coroutine_running, == CR_INITIALIZING);
 
@@ -352,8 +384,26 @@ static inline void _cr_transition(enum coroutine_state state) {
 	debugf("cid=%zu: transition %i->%i\n", coroutine_running, coroutine_table[coroutine_running-1].state, state);
 
 	coroutine_table[coroutine_running-1].state = state;
-	if (!setjmp(coroutine_table[coroutine_running-1].env)) /* point=c2 */
-		longjmp(coroutine_main_env, 1); /* jump to point=b */
+
+	cid_t next;
+	for (;;) {
+		next = next_coroutine();
+		if (next)
+			break;
+		if (state == CR_RUNNABLE) {
+			debugf("cid=%zu: no other runnable coroutines, not yielding\n", coroutine_running);
+			coroutine_table[coroutine_running-1].state = CR_RUNNING;
+			return;
+		} else {
+			debugf("cid=%zu: no runnable coroutines, sleeping\n", coroutine_running);
+			pause();
+		}
+	}
+	if (!setjmp(coroutine_table[coroutine_running-1].env)) { /* point=c2 */
+		coroutine_running = next;
+		coroutine_table[coroutine_running-1].state = CR_RUNNING;
+		longjmp(coroutine_table[coroutine_running-1].env, 1); /* jump to point=b */
+	}
 }
 
 void cr_yield(void) {
diff --git a/libcr/include/libcr/coroutine.h b/libcr/include/libcr/coroutine.h
index 47467bd..faaf6c3 100644
--- a/libcr/include/libcr/coroutine.h
+++ b/libcr/include/libcr/coroutine.h
@@ -96,9 +96,8 @@ cid_t coroutine_add(cr_fn_t fn, void *args);
  * The main scheduler loop.
  *
  * "Should" never return, but will print a message to stderr and
- * return if the program has deadlocked and there are no runnable
- * coroutines.  So be sure to call coroutine_add() at least once
- * before calling this.
+ * return if there are no coroutines (there were no calls to
+ * coroutine_add(), or all coroutines cr_exit()).
  */
 void coroutine_main(void);