Add switch to fiber interface, support fiber_switch in asyncify implementation

Makefile

@@ -20,13 +20,19 @@ endif
 all: hello sieve itersum treesum
 
 .PHONY: hello
-hello: hello_asyncify.wasm hello_wasmfx.wasm
+hello: hello_asyncify.wasm hello_asyncify_switch.wasm hello_wasmfx.wasm
 
 hello_asyncify.wasm: inc/fiber.h src/asyncify/asyncify_impl.c examples/hello.c
 	$(WASICC) -DSTACK_POOL_SIZE=$(STACK_POOL_SIZE) -DASYNCIFY_DEFAULT_STACK_SIZE=$(ASYNCIFY_DEFAULT_STACK_SIZE) src/asyncify/asyncify_impl.c $(WASIFLAGS) examples/hello.c -o hello_asyncfiy.pre.wasm
 	$(ASYNCIFY) hello_asyncfiy.pre.wasm -o hello_asyncify.wasm
 	chmod +x hello_asyncify.wasm
 
+# new stuff: build with switch
+hello_asyncify_switch.wasm: inc/fiber-switch.h src/asyncify/asyncify_switch_impl.c examples/hello_switch.c
+	$(WASICC) -DSTACK_POOL_SIZE=$(STACK_POOL_SIZE) -DASYNCIFY_DEFAULT_STACK_SIZE=$(ASYNCIFY_DEFAULT_STACK_SIZE) src/asyncify/asyncify_switch_impl.c $(WASIFLAGS) examples/hello_switch.c -o hello_asyncfiy_switch.pre.wasm
+	$(ASYNCIFY) hello_asyncfiy_switch.pre.wasm -o hello_asyncify_switch.wasm
+	chmod +x hello_asyncify_switch.wasm
+
 hello_wasmfx.wasm: inc/fiber.h src/wasmfx/imports.wat src/wasmfx/wasmfx_impl.c examples/hello.c
 	$(WASICC) $(SHADOW_STACK_FLAG) -DWASMFX_CONT_SHADOW_STACK_SIZE=$(WASMFX_CONT_SHADOW_STACK_SIZE) -DWASMFX_CONT_TABLE_INITIAL_CAPACITY=$(WASMFX_CONT_TABLE_INITIAL_CAPACITY) -Wl,--export-table,--export-memory,--export=__stack_pointer src/wasmfx/wasmfx_impl.c $(WASIFLAGS) examples/hello.c -o hello_wasmfx.pre.wasm
 	$(WASM_INTERP) -d -i src/wasmfx/imports.wat -o fiber_wasmfx_imports.wasm

examples/hello_switch.c

@@ -0,0 +1,70 @@
+// Cooperative printing of "hello world"
+#include <assert.h>
+#include <stdbool.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <fiber-switch.h>
+
+static volatile fiber_t hello_fiber;
+static volatile fiber_t world_fiber;
+
+fiber_result_t status;
+
+bool hello_done;
+bool world_done;
+
+
+void* hello(void *arg) {
+  uint32_t i = (uint32_t)(uintptr_t)arg;
+
+  static const char s[] = "hlowrd";
+
+  if (!hello_done){
+    while (i < strlen(s)) {
+      putc(s[i], stdout);
+      i = (uint32_t)(uintptr_t)fiber_switch(world_fiber, (void*)(uintptr_t)i, &status);
+      if (status == FIBER_OK) hello_done = true;
+    }
+  }
+
+  return NULL;
+}
+
+void* world(void *arg) {
+  uint32_t i = (uint32_t)(uintptr_t)arg;
+  static const char s[] = "el ol";
+
+  if (!world_done) {
+    putc(s[i], stdout);
+    i++;
+    i = (uint32_t)(uintptr_t)fiber_switch(hello_fiber, (void*)(uintptr_t)i, &status);
+    if (status == FIBER_OK) world_done = true;
+  }
+
+  return NULL;
+}
+
+int main(void) {
+  fiber_init();
+
+  hello_fiber = fiber_alloc(hello);
+  world_fiber = fiber_alloc(world);
+
+  hello_done = false;
+  world_done = false;
+
+  uint32_t i = 0;
+
+  (void)fiber_resume(hello_fiber, (void*)(uintptr_t)i, &status);
+
+  putc('\n', stdout);
+  fflush(stdout);
+
+  fiber_free(hello_fiber);
+  fiber_free(world_fiber);
+
+  fiber_finalize();
+  return 0;
+}

inc/fiber-switch.h

@@ -0,0 +1,45 @@
+/** A basic fiber interface. **/
+#ifndef WASMFX_FIBER_C_H
+#define WASMFX_FIBER_C_H
+
+#define export(NAME) __attribute__((export_name(NAME)))
+
+/** The signature of a fiber entry point. **/
+typedef void* (*fiber_entry_point_t)(void*);
+
+/** The abstract type of a fiber object. **/
+typedef struct fiber* fiber_t;
+
+/** Allocates a new fiber with the default stack size. **/
+export("fiber_alloc")
+fiber_t fiber_alloc(fiber_entry_point_t entry);
+
+/** Reclaims the memory occupied by the fiber object. **/
+export("fiber_free")
+void fiber_free(fiber_t fiber);
+
+
+/** Possible status codes for `fiber_switch`. **/
+typedef enum { FIBER_OK = 0, FIBER_SWITCH = 1, FIBER_ERROR = 2 } fiber_result_t;
+
+
+/** Switches to a given `fiber` with argument `arg`, returning some value
+    of type `void*`. The output parameter `status` indicates whether
+    the fiber ran to completion (`FIBER_OK`), yielded control
+    (`FIBER_YIELD`), or failed (`FIBER_ERROR`), in the latter case the
+    return value will be undefined. **/
+export("fiber_switch")
+void* fiber_switch(fiber_t fiber, void *arg, fiber_result_t *status);
+
+/** Initialises the fiber runtime. It must be called exactly once
+    before using any of the other fiber_* functions. **/
+export("fiber_init")
+void fiber_init(void);
+
+/** Tears down the fiber runtime. It must be called after the final
+    use of any of the other fiber_* functions. **/
+export("fiber_finalize")
+void fiber_finalize(void);
+
+#undef export
+#endif

src/asyncify/asyncify_switch_impl.c

@@ -0,0 +1,231 @@
+// An asyncify implementation of the basic fiber interface.
+
+#include <stdlib.h>
+#include <stdint.h>
+// for printing things, remove later
+#include <stdio.h>
+#include <assert.h>
+
+#include "fiber-switch.h"
+#define import(NAME) __attribute__((import_module("asyncify"),import_name(NAME)))
+
+
+/** Asyncify imports **/
+// The following functions are asyncify primitives:
+// * asyncify_start_unwind(iptr): initiates a continuation
+//   capture. The argument `iptr` is a pointer to an asyncify stack.
+// * asyncify_stop_unwind(): delimits a capture continuations.
+// * asyncify_start_rewind(iptr): initiates a continuation
+//   reinstatement. The argument `iptr` is a pointer to an asyncify
+//   stack.
+// * asyncfiy_stop_rewind(): delimits the extent of a continuation
+//  reinstatement.
+extern
+import("start_unwind")
+void asyncify_start_unwind(void*);
+
+extern
+import("stop_unwind")
+void asyncify_stop_unwind(void);
+
+extern
+import("start_rewind")
+void asyncify_start_rewind(void*);
+
+extern
+import("stop_rewind")
+void asyncify_stop_rewind(void);
+
+// The default stack size is 2MB.
+static const size_t default_stack_size = ASYNCIFY_DEFAULT_STACK_SIZE;
+
+// We track the currently active fiber via this global variable.
+static volatile fiber_t active_fiber = NULL;
+
+// Fiber states:
+// * ACTIVE: the fiber is actively executing.
+// * YIELDING: the fiber is suspended.
+// * DONE: the fiber is finished (i.e. run to completion).
+typedef enum { ACTIVE, YIELDING, DONE } fiber_state_t;
+
+// A fiber stack is an asyncify stack, i.e. a reserved area of memory
+// for asyncify to store the call chain and locals. Note: asyncify
+// assumes `end` is at offset 4. Moreover, asyncify stacks grow
+// upwards, so it must be that top <= end.
+struct  __attribute__((packed)) fiber_stack {
+  uint8_t *top;
+  uint8_t *end;
+  uint8_t *buffer;
+};
+static_assert(sizeof(uint8_t*) == 4, "sizeof(uint8_t*) != 4");
+static_assert(sizeof(struct fiber_stack) == 12, "struct fiber_stack: No padding allowed");
+
+// The fiber structure embeds the asyncify stack (struct fiber_stack),
+// its state, an entry point, and two buffers for communicating
+// payloads and managing fiber local data, respectively.
+struct fiber {
+  /** The underlying asyncify stack. */
+  struct fiber_stack stack;
+  // Fiber state.
+  fiber_state_t state;
+  // Initial function to run on the fiber.
+  fiber_entry_point_t entry;
+  // Payload buffer.
+  void *arg;
+};
+
+// Allocates a fiber stack of size stack_size.
+static struct fiber_stack fiber_stack_alloc(size_t stack_size) {
+  uint8_t *buffer = malloc(sizeof(uint8_t) * stack_size);
+  uint8_t *top = buffer;
+  uint8_t *end = buffer + stack_size;
+  struct fiber_stack stack = (struct fiber_stack) { top, end, /* NULL, */ buffer };
+  return stack;
+}
+
+// Frees an allocated fiber_stack.
+static void fiber_stack_free(struct fiber_stack fiber_stack) {
+  free(fiber_stack.buffer);
+}
+
+#if defined STACK_POOL_SIZE && STACK_POOL_SIZE > 0
+// Fiber stack pool
+struct stack_pool {
+  int32_t next;
+  struct fiber_stack stacks[STACK_POOL_SIZE];
+};
+
+static struct fiber_stack stack_pool_next(volatile struct stack_pool *pool) {
+  assert(pool->next >= 0 && pool->next < STACK_POOL_SIZE);
+  return pool->stacks[pool->next++];
+}
+
+static void stack_pool_reclaim(volatile struct stack_pool *pool, struct fiber_stack stack) {
+  assert(pool->next > 0 && pool->next <= STACK_POOL_SIZE);
+  pool->stacks[--pool->next] = stack;
+  return;
+}
+
+static volatile struct stack_pool pool;
+#endif
+
+// Allocates a fiber object.
+// NOTE: the entry point `fn` should be careful about uses of `printf`
+// and related functions, as they can cause asyncify to corrupt its
+// own state. See `wasi-io.h` for asyncify-safe printing functions.
+fiber_t fiber_sized_alloc(size_t stack_size, fiber_entry_point_t entry) {
+  fiber_t fiber = (fiber_t)malloc(sizeof(struct fiber));
+#if defined STACK_POOL_SIZE && STACK_POOL_SIZE > 0
+  (void)stack_size;
+  fiber->stack = stack_pool_next(&pool);
+  // TODO(dhil): It may be necessary to reset the top pointer.  I'd
+  // need to test on a larger example.
+  fiber->stack.top = fiber->stack.buffer;
+#else
+  fiber->stack = fiber_stack_alloc(stack_size);
+#endif
+  fiber->state = ACTIVE;
+  fiber->entry = entry;
+  fiber->arg = NULL;
+  return fiber;
+}
+
+// Allocates a fiber object with the default stack size.
+__attribute__((noinline))
+fiber_t fiber_alloc(fiber_entry_point_t entry) {
+  //printf("fiber_alloc called\n");
+  return fiber_sized_alloc(default_stack_size, entry);
+}
+
+// Frees a fiber object.
+__attribute__((noinline))
+void fiber_free(fiber_t fiber) {
+#if defined STACK_POOL_SIZE && STACK_POOL_SIZE > 0
+  stack_pool_reclaim(&pool, fiber->stack);
+#else
+  fiber_stack_free(fiber->stack);
+#endif
+  free(fiber);
+}
+
+
+// Switches to a given fiber, transferring control to it.
+__attribute__((noinline))
+void* fiber_switch(fiber_t fiber, void *arg, fiber_result_t *result) {
+
+  // Sanity check: we are done, signal error and return.
+  if (fiber->state == DONE) {
+    printf("fiber_switch: fiber is DONE\n");
+    *result = FIBER_ERROR;
+    return NULL;
+  }
+
+  // Otherwise, set the argument buffer, yield from current fiber, and start unwinding.
+  active_fiber->arg = arg;
+  active_fiber->state = YIELDING;
+
+  // Note: it appears that every start_unwind needs to be immediately followed by a stop_unwind,
+  //       otherwise it crashes. This behaviour wasn't obvious to me and isn't documented anywhere.
+  //       Not sure if the following lines are correct, but it gets my test program working.
+  asyncify_start_unwind(&active_fiber->stack);
+  asyncify_stop_unwind();
+
+  // Remember the currently executing fiber.
+  volatile fiber_t prev = active_fiber;
+
+  // Set the given fiber as the actively executing fiber.
+  active_fiber = fiber;
+
+  // If we are resuming a suspended fiber...
+  if (fiber->state == YIELDING) {
+    // ... then update the argument buffer.
+    fiber->arg = arg;
+    // ... and initiate the stack rewind.
+    // Note: like above, it appears that every start_rewind needs to be immediately followed by a stop_rewind,
+    //       No idea if this is correct either.
+    asyncify_start_rewind(&fiber->stack);
+    asyncify_stop_rewind();
+    active_fiber->state = ACTIVE;
+  }
+
+  void *fiber_result = fiber->entry(arg);
+
+  // Check whether the fiber finished or suspended.
+  if (fiber->state != YIELDING) fiber->state = DONE;
+
+  // Restore the previously executing fiber.
+  active_fiber = prev;
+
+  // Signal success.
+  if (fiber->state == YIELDING) {
+    *result = FIBER_YIELD;
+    return fiber->arg;
+  } else {
+    *result = FIBER_OK;
+    return fiber_result;
+  }
+}
+
+// Noop when stack pooling is disabled.
+void fiber_init(void) {
+#if defined STACK_POOL_SIZE && STACK_POOL_SIZE > 0
+  pool.next = STACK_POOL_SIZE;
+  for (uint32_t i = 0; i < STACK_POOL_SIZE; i++) {
+    stack_pool_reclaim(&pool, fiber_stack_alloc(default_stack_size));
+  }
+  assert(pool.next == 0);
+#endif
+}
+
+// Noop when stack pooling is disabled.
+void fiber_finalize(void) {
+#if defined STACK_POOL_SIZE && STACK_POOL_SIZE > 0
+  assert(pool.next == 0);
+  for (uint32_t i = 0; i < STACK_POOL_SIZE; i++) {
+    fiber_stack_free(stack_pool_next(&pool));
+  }
+  assert(pool.next == STACK_POOL_SIZE);
+#endif
+}
+
+#undef import