Changing the API required so that we do not have to call jump_fcontext() at the end of the coroutine function.

Author: kyle-github

include/fcontext.h

@@ -88,7 +88,7 @@ typedef struct {
  * Called when context is first entered via jump_fcontext.
  * The transfer_t contains the previous context and initial data.
  */
-typedef void (*fcontext_fn_t)(fcontext_transfer_t);
+typedef fcontext_transfer_t (*fcontext_fn_t)(fcontext_transfer_t);
 
 /**
  * Function run "on top" of an existing context.

src/fcontext_wrapper.c

@@ -44,6 +44,35 @@ static size_t get_page_size(void) {
 #endif
 }
 
+/* Trampoline function that wraps the user's entry point */
+static fcontext_transfer_t context_trampoline(fcontext_transfer_t t) {
+    /* 1. Retrieve the user function passed from fcontext_init_trampoline */
+    fcontext_fn_t fn = (fcontext_fn_t)t.data;
+
+    /* 2. Yield back to fcontext_create immediately.
+     *    We return control to the creator so it can return the context to the user.
+     */
+    t = jump_fcontext(t.prev_context, NULL);
+
+    /* 3. We are now resumed by the user (via fcontext_swap/jump_fcontext).
+     *    t.prev_context is the user's context.
+     *    t.data is the user's data.
+     */
+    t = fn(t);
+
+    /* 4. Automatic cleanup: Jump back to the last known context */
+    return jump_fcontext(t.prev_context, t.data);
+}
+
+/* Helper to initialize the trampoline. Call this inside fcontext_create* functions. */
+static fcontext_t fcontext_init_trampoline(fcontext_t ctx, fcontext_fn_t fn) {
+    /* Perform a context switch to the new context to pass the function pointer.
+     * The trampoline will read 'fn' from 'data' and immediately yield back.
+     */
+    fcontext_transfer_t t = jump_fcontext(ctx, (void *)fn);
+    return t.prev_context;
+}
+
 /**
  * Create a new context using malloc with software guard zones.
  */
@@ -92,7 +121,9 @@ fcontext_stack_t *fcontext_create_malloc(size_t stack_size, size_t guard_size, f
     /* The stack pointer passed to make_fcontext must be aligned */
     sp = fcontext_align_stack_pointer(sp);
 
-    ctx->context = make_fcontext(sp, stack_size, entry_fn);
+    /* Create context pointing to trampoline, then init it */
+    ctx->context = make_fcontext(sp, stack_size, (fcontext_fn_t)context_trampoline);
+    ctx->context = fcontext_init_trampoline(ctx->context, entry_fn);
     ctx->alloc_type = FCONTEXT_ALLOC_MALLOC;
     ctx->alloc_base = raw_alloc; /* Store original malloc pointer for free() */
     ctx->stack_base = stack_base;
@@ -180,7 +211,9 @@ fcontext_stack_t *fcontext_create_mmap(size_t stack_size, size_t guard_size, fco
 #endif
 
     void *sp = (char *)stack_base + stack_size;
-    ctx->context = make_fcontext(sp, stack_size, entry_fn);
+    /* Create context pointing to trampoline, then init it */
+    ctx->context = make_fcontext(sp, stack_size, (fcontext_fn_t)context_trampoline);
+    ctx->context = fcontext_init_trampoline(ctx->context, entry_fn);
     ctx->alloc_type = FCONTEXT_ALLOC_MMAP;
     ctx->alloc_base = region;
     ctx->stack_base = stack_base;

tests/test_fcontext_alignment.c

@@ -10,15 +10,15 @@
  * Distributed under the Boost Software License, Version 1.0.
  */
 
-#include <stdio.h>
-#include <stdlib.h>
+#include "fcontext.h"
 #include <assert.h>
 #include <stdint.h>
-#include "fcontext.h"
+#include <stdio.h>
+#include <stdlib.h>
 
 static int alignment_ok = 0;
 
-void align_fiber(fcontext_transfer_t t) {
+fcontext_transfer_t align_fiber(fcontext_transfer_t t) {
     uintptr_t sp = 0;
 
     /* Capture Stack Pointer using inline assembly */
@@ -60,29 +60,29 @@ void align_fiber(fcontext_transfer_t t) {
      * Call pushes 8-byte return address.
      * On entry: RSP % 16 == 8.
      */
-    if ((sp & 0xF) == 0x8) {
+    if((sp & 0xF) == 0x8) {
         printf("  [fiber] x86_64 ABI verified (RSP %% 16 == 8)\n");
         alignment_ok = 1;
-    } else if ((sp & 0xF) == 0x0) {
+    } else if((sp & 0xF) == 0x0) {
         /* Some runtimes/compilers might force 16-byte alignment on entry for SIMD optimization */
         printf("  [fiber] 16-byte alignment verified (RSP %% 16 == 0)\n");
         alignment_ok = 1;
     }
 #elif defined(__aarch64__)
     /* AAPCS64: SP must be 16-byte aligned at all times */
-    if ((sp & 0xF) == 0x0) {
+    if((sp & 0xF) == 0x0) {
         printf("  [fiber] ARM64 AAPCS verified (SP %% 16 == 0)\n");
         alignment_ok = 1;
     }
 #else
     /* Default check: 8-byte or 16-byte alignment is generally acceptable for others */
-    if ((sp & 0x7) == 0x0) {
+    if((sp & 0x7) == 0x0) {
         printf("  [fiber] Basic alignment verified (>= 8 bytes)\n");
         alignment_ok = 1;
     }
 #endif
 
-    jump_fcontext(t.prev_context, NULL);
+    return t;
 }
 
 int main(void) {
@@ -95,7 +95,7 @@ int main(void) {
 
     fcontext_destroy(state);
 
-    if (alignment_ok) {
+    if(alignment_ok) {
         printf("\n✓ PASS: Stack alignment requirements met\n");
         return 0;
     } else {

tests/test_fcontext_basic.c

@@ -11,18 +11,18 @@
  *  http://www.boost.org/LICENSE_1_0.txt)
  */
 
-#include <stdio.h>
-#include <stdlib.h>
+#include "fcontext.h"
 #include <assert.h>
 #include <stdint.h>
-#include "fcontext.h"
+#include <stdio.h>
+#include <stdlib.h>
 
 static int test_phase = 0;
 
 /**
  * Fiber function - entry point for new context
  */
-void fiber_func(fcontext_transfer_t t) {
+fcontext_transfer_t fiber_func(fcontext_transfer_t t) {
     printf("  [fiber] entered context\n");
     fflush(stdout);
     assert(test_phase == 0);
@@ -39,6 +39,9 @@ void fiber_func(fcontext_transfer_t t) {
 
     printf("  [fiber] finishing\n");
     fflush(stdout);
+
+    /* Return to trampoline to switch back */
+    return t;
 }
 
 int main(void) {
@@ -47,7 +50,7 @@ int main(void) {
 
     /* Create context with guarded stack using mmap */
     fcontext_stack_t *state = fcontext_create(24 * 1024, fiber_func);
-    if (state == NULL) {
+    if(state == NULL) {
         fprintf(stderr, "Failed to create context\n");
         return 1;
     }

tests/test_fcontext_guard.c

@@ -37,7 +37,7 @@ static void *g_stack_base = NULL;
 static size_t g_stack_size = 0;
 
 /* Recursive function to consume stack until overflow */
-void overflow_fiber(fcontext_transfer_t t) {
+fcontext_transfer_t overflow_fiber(fcontext_transfer_t t) {
     /* Allocate 1KB on stack */
     /* Force alignment to avoid potential SIGILL from SIMD instructions on unaligned stack */
 #if defined(_MSC_VER)
@@ -58,6 +58,7 @@ void overflow_fiber(fcontext_transfer_t t) {
 
     /* Prevent tail call optimization */
     (void)buffer[0];
+    return t;
 }
 
 #ifndef _WIN32

tests/test_fcontext_low_level.c

@@ -9,58 +9,58 @@
  * Distributed under the Boost Software License, Version 1.0.
  */
 
+#include "fcontext.h"
+#include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
-#include <assert.h>
-#include "fcontext.h"
 
 static int fiber_ran = 0;
 
-void low_level_fiber(fcontext_transfer_t t) {
+fcontext_transfer_t low_level_fiber(fcontext_transfer_t t) {
     printf("  [fiber] entered low-level fiber\n");
-    
+
     /* Verify data passed */
-    int* value = (int*)t.data;
+    int *value = (int *)t.data;
     assert(*value == 123);
-    
+
     fiber_ran = 1;
-    
+
     printf("  [fiber] jumping back\n");
-    jump_fcontext(t.prev_context, NULL);
+    return jump_fcontext(t.prev_context, NULL);
 }
 
 int main(void) {
     printf("=== fcontext Low-Level API Test (make_fcontext) ===\n");
-    
+
     /* 1. Manually allocate stack */
     size_t stack_size = 16 * 1024;
-    void* stack_buffer = malloc(stack_size);
-    if (!stack_buffer) {
+    void *stack_buffer = malloc(stack_size);
+    if(!stack_buffer) {
         fprintf(stderr, "Failed to allocate stack\n");
         return 1;
     }
-    
+
     /* 2. Calculate aligned stack pointer (top of stack) */
-    void* sp = (char*)stack_buffer + stack_size;
+    void *sp = (char *)stack_buffer + stack_size;
     sp = fcontext_align_stack_pointer(sp);
-    
+
     printf("[main] Stack allocated at %p, SP aligned to %p\n", stack_buffer, sp);
-    
+
     /* 3. Create context directly */
     fcontext_t ctx = make_fcontext(sp, stack_size, low_level_fiber);
     assert(ctx != NULL);
-    
+
     /* 4. Jump to context */
     int data = 123;
     printf("[main] Jumping to fiber...\n");
     jump_fcontext(ctx, &data);
-    
+
     printf("[main] Returned from fiber\n");
     assert(fiber_ran == 1);
-    
+
     /* 5. Cleanup */
     free(stack_buffer);
-    
+
     printf("\n✓ PASS: Low-level make_fcontext test completed.\n");
     return 0;
 }

tests/test_fcontext_ontop.c

@@ -6,10 +6,10 @@
  * Distributed under the Boost Software License, Version 1.0.
  */
 
+#include "fcontext.h"
+#include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
-#include <assert.h>
-#include "fcontext.h"
 
 typedef struct {
     int value;
@@ -18,63 +18,63 @@ typedef struct {
 /* Function to be executed "on top" of the fiber */
 fcontext_transfer_t ontop_func(fcontext_transfer_t t) {
     printf("  [ontop] Executing ontop function\n");
-    test_data_t* data = (test_data_t*)t.data;
-    
+    test_data_t *data = (test_data_t *)t.data;
+
     /* Verify we received the expected data from main */
     assert(data->value == 3);
-    
+
     /* Modify data to prove we ran */
     data->value += 10;
-    
+
     /* Return transfer to resume the fiber */
     return t;
 }
 
-void fiber_func(fcontext_transfer_t t) {
+fcontext_transfer_t fiber_func(fcontext_transfer_t t) {
     printf("  [fiber] Started\n");
-    test_data_t* data = (test_data_t*)t.data;
-    
+    test_data_t *data = (test_data_t *)t.data;
+
     /* 1. Verify initial value passed from main */
     assert(data->value == 1);
-    
+
     /* Modify and yield back */
     data->value = 2;
     printf("  [fiber] Yielding back to main\n");
-    
+
     /* Use fcontext_swap for standard yield */
     t = fcontext_swap(t.prev_context, data);
-    
+
     /* 2. We are back. ontop_fcontext should have executed ontop_func before we resumed. */
     printf("  [fiber] Resumed\n");
-    data = (test_data_t*)t.data;
-    
+    data = (test_data_t *)t.data;
+
     /* Verify ontop_func modification (3 + 10 = 13) */
     assert(data->value == 13);
-    
+
     printf("  [fiber] Finished\n");
-    fcontext_swap(t.prev_context, data);
+    return fcontext_swap(t.prev_context, data);
 }
 
 int main(void) {
     printf("=== fcontext ontop Test ===\n");
-    
-    fcontext_stack_t* stack = fcontext_create(16 * 1024, fiber_func);
-    if (!stack) {
+
+    fcontext_stack_t *stack = fcontext_create(16 * 1024, fiber_func);
+    if(!stack) {
         fprintf(stderr, "Failed to create context\n");
         return 1;
     }
-    
-    test_data_t data = { .value = 1 };
-    
+
+    test_data_t data = {.value = 1};
+
     printf("[main] Jumping to fiber\n");
     fcontext_transfer_t t = fcontext_swap(stack->context, &data);
-    
+
     /* Fiber yielded. Value should be 2. */
     assert(data.value == 2);
-    
+
     printf("[main] Fiber yielded. Calling ontop_fcontext\n");
     data.value = 3;
-    
+
     /* Use ontop_fcontext to resume fiber and run ontop_func first */
     /* Note: We manually handle ASAN hooks here since ontop_fcontext is raw API */
 #ifdef __SANITIZE_ADDRESS__
@@ -85,11 +85,11 @@ int main(void) {
 #ifdef __SANITIZE_ADDRESS__
     __sanitizer_finish_switch_fiber(fake_stack_save, NULL, NULL);
 #endif
-    
+
     printf("[main] Fiber returned\n");
-    
+
     fcontext_destroy(stack);
-    
+
     printf("\n✓ PASS: ontop_fcontext test completed.\n");
     return 0;
 }