Skip to content

Instantly share code, notes, and snippets.

@JiayinCao
Last active March 27, 2024 11:03
Show Gist options
  • Star 24 You must be signed in to star a gist
  • Fork 5 You must be signed in to fork a gist
  • Save JiayinCao/07475d3423952b702d1efc5268b0df4e to your computer and use it in GitHub Desktop.
Save JiayinCao/07475d3423952b702d1efc5268b0df4e to your computer and use it in GitHub Desktop.
Tiny Fiber ( A razor thin cross platform fiber library )
/*
MIT License
Copyright (c) 2023 Jiayin Cao
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/
// Tiny fiber is a lazor thin library that support fiber interfaces on different 64 bits platforms.
// User can include this one header file in their project to use fiber in their own projects.
// It is currently supported on the following platform
// - X64 Windows
// - X64 Linux
// - X64 Mac
// - Arm64 Linux
// - Arm64 Mac
//
// Below is a simple example demonstrating the basic usage of the library
//
// #include <iostream>
// #include "tiny_fiber.h"
//
// tiny_fiber::FiberHandle thread_fiber;
// tiny_fiber::FiberHandle fiber;
//
// void fibermain(void* arg) {
// tiny_fiber::FiberHandle fiber = *reinterpret_cast<tiny_fiber::FiberHandle*>(arg);
// std::cout<<"hello fiber"<<std::endl;
// tiny_fiber::SwitchFiber(fiber, thread_fiber);
// }
//
// int main(int argc, char** argv) {
// const int stack_size = 1024 * 16;
// thread_fiber = tiny_fiber::CreateFiberFromThread();
// fiber = tiny_fiber::CreateFiber(stack_size, fibermain, &fiber);
// tiny_fiber::SwitchFiber(thread_fiber, fiber);
// return 0;
// }
//
#pragma once
#if defined(__aarch64__)
#define TINY_FIBER_ARM64
#elif defined(__x86_64__) || defined(_M_X64)
#define TINY_FIBER_X64
#else
#error "Target architecture is not supported."
#endif
#if defined(__APPLE__)
#define TINY_FIBER_APPLE
#elif defined(WIN32) || defined(_WIN32) || defined(__WIN32__) || defined(__NT__)
#define TINY_FIBER_WIN
#elif defined(__linux__)
#define TINY_FIBER_LINUX
#else
#error "Target Platform is not supported."
#endif
#if defined(TINY_FIBER_ARM64) && !defined(TINY_FIBER_WIN)
#define FIBER_ASM_ARM64_IMPLEMENTATION
#elif defined(TINY_FIBER_X64) && !defined(TINY_FIBER_WIN)
#define FIBER_ASM_X64_IMPLEMENTATION
#endif
#if defined(TINY_FIBER_MALLOC) && defined(TINY_FIBER_FREE)
// custom stack allocation is fine
#elif !defined(TINY_FIBER_MALLOC) && !defined(TINY_FIBER_FREE)
// default stack allocation is fine either
#else
#error "Do not define either of TINY_FIBER_MALLOC and TINY_FIBER_FREE"
#endif
#include <memory>
#if defined(TINY_FIBER_WIN)
#include <Windows.h>
#endif
#ifndef TINY_FIBER_MALLOC
#define TINY_FIBER_MALLOC malloc
#define TINY_FIBER_FREE free
#endif
namespace tiny_fiber {
// an alias for CPU registers, all need to be 64 bits long
typedef unsigned long long Register;
// check the register size, making sure it is 8 bytes.
static_assert(sizeof(Register) == 8, "Incorrect register size");
// helper macro to check the offset of a register in fiber context
#define CHECK_OFFSET(REGISTER, OFFSET) \
static_assert(offsetof(FiberContexInternal, REGISTER) == OFFSET, \
"Incorrect register offset")
// Both x64 and arm64 require stack memory pointer to be 16 bytes aligned
#define FIBER_STACK_ALIGNMENT 16
// MacOS requires a _ to be defined at the beginning of a function name implemented in asm
#if defined(TINY_FIBER_APPLE)
#define _switch_fiber_internal switch_fiber_internal
#endif
#ifdef FIBER_ASM_ARM64_IMPLEMENTATION
#define FIBER_ASM_IMPLEMENTATION
// alias for registers that need to be saved during a fiber switch
#define FIBER_REG_X0 0x00
#define FIBER_REG_X16 0x08
#define FIBER_REG_X17 0x10
#define FIBER_REG_X18 0x18
#define FIBER_REG_X19 0x20
#define FIBER_REG_X20 0x28
#define FIBER_REG_X21 0x30
#define FIBER_REG_X22 0x38
#define FIBER_REG_X23 0x40
#define FIBER_REG_X24 0x48
#define FIBER_REG_X25 0x50
#define FIBER_REG_X26 0x58
#define FIBER_REG_X27 0x60
#define FIBER_REG_X28 0x68
#define FIBER_REG_X29 0x70
#define FIBER_REG_X30 0x78
#define FIBER_REG_V8 0x80
#define FIBER_REG_V9 0x88
#define FIBER_REG_V10 0x90
#define FIBER_REG_V11 0x98
#define FIBER_REG_V12 0xa0
#define FIBER_REG_V13 0xa8
#define FIBER_REG_V14 0xb0
#define FIBER_REG_V15 0xb8
#define FIBER_REG_SP 0xc0
/**
* On Arm 64 architecture, each register should be 64 bits. And following registers need to be
* saved during a context switch
* - [X16 - X30]
* https://developer.arm.com/documentation/102374/0101/Procedure-Call-Standard
* - [D16 - D23]
* https://developer.arm.com/documentation/den0024/a/The-ABI-for-ARM-64-bit-Architecture/Register-use-in-the-AArch64-Procedure-Call-Standard/Parameters-in-NEON-and-floating-point-registers
*/
struct FiberContexInternal {
// funciton paramters
Register x0;
// general purpose registers
Register x16;
Register x17;
Register x18;
Register x19;
Register x20;
Register x21;
Register x22;
Register x23;
Register x24;
Register x25;
Register x26;
Register x27;
Register x28;
Register x29; // frame pointer (FP)
Register x30; // link register (R30)
// float-point and neon registers
Register v8;
Register v9;
Register v10;
Register v11;
Register v12;
Register v13;
Register v14;
Register v15;
// stack pointer
Register sp;
};
// Making sure our offsets are all correct
CHECK_OFFSET(x0, FIBER_REG_X0);
CHECK_OFFSET(x16, FIBER_REG_X16);
CHECK_OFFSET(x17, FIBER_REG_X17);
CHECK_OFFSET(x18, FIBER_REG_X18);
CHECK_OFFSET(x19, FIBER_REG_X19);
CHECK_OFFSET(x20, FIBER_REG_X20);
CHECK_OFFSET(x21, FIBER_REG_X21);
CHECK_OFFSET(x22, FIBER_REG_X22);
CHECK_OFFSET(x23, FIBER_REG_X23);
CHECK_OFFSET(x24, FIBER_REG_X24);
CHECK_OFFSET(x25, FIBER_REG_X25);
CHECK_OFFSET(x26, FIBER_REG_X26);
CHECK_OFFSET(x27, FIBER_REG_X27);
CHECK_OFFSET(x28, FIBER_REG_X28);
CHECK_OFFSET(x29, FIBER_REG_X29);
CHECK_OFFSET(x30, FIBER_REG_X30);
CHECK_OFFSET(v8, FIBER_REG_V8);
CHECK_OFFSET(v9, FIBER_REG_V9);
CHECK_OFFSET(v10, FIBER_REG_V10);
CHECK_OFFSET(v11, FIBER_REG_V11);
CHECK_OFFSET(v12, FIBER_REG_V12);
CHECK_OFFSET(v13, FIBER_REG_V13);
CHECK_OFFSET(v14, FIBER_REG_V14);
CHECK_OFFSET(v15, FIBER_REG_V15);
CHECK_OFFSET(sp, FIBER_REG_SP);
// create a new fiber
extern "C" inline bool _create_fiber_internal(void* stack, uint32_t stack_size, void (*target)(void*), void* arg, FiberContexInternal* context) {
// it is the users responsibility to make sure the stack is 16 bytes aligned, which is required by the Arm64 architecture
if((((uintptr_t)stack) & (FIBER_STACK_ALIGNMENT - 1)) != 0)
return false;
uintptr_t* stack_top = (uintptr_t*)((uint8_t*)(stack) + stack_size);
context->x30 = (uintptr_t)target;
context->x0 = (uintptr_t)arg;
context->sp = (uintptr_t)stack_top;
return true;
}
// switching to a different fiber
asm(
R"(
.text
.align 4
_switch_fiber_internal:
str x16, [x0, 0x08] /* FIBER_REG_X16 */
str x17, [x0, 0x10] /* FIBER_REG_X17 */
str x18, [x0, 0x18] /* FIBER_REG_X18 */
str x19, [x0, 0x20] /* FIBER_REG_X19 */
str x20, [x0, 0x28] /* FIBER_REG_X20 */
str x21, [x0, 0x30] /* FIBER_REG_X21 */
str x22, [x0, 0x38] /* FIBER_REG_X22 */
str x23, [x0, 0x40] /* FIBER_REG_X23 */
str x24, [x0, 0x48] /* FIBER_REG_X24 */
str x25, [x0, 0x50] /* FIBER_REG_X25 */
str x26, [x0, 0x58] /* FIBER_REG_X26 */
str x27, [x0, 0x60] /* FIBER_REG_X27 */
str x28, [x0, 0x68] /* FIBER_REG_X28 */
str x29, [x0, 0x70] /* FIBER_REG_X29 */
str d8, [x0, 0x80] /* FIBER_REG_V8 */
str d9, [x0, 0x88] /* FIBER_REG_V9 */
str d10, [x0, 0x90] /* FIBER_REG_V10 */
str d11, [x0, 0x98] /* FIBER_REG_V11 */
str d12, [x0, 0xa0] /* FIBER_REG_V12 */
str d13, [x0, 0xa8] /* FIBER_REG_V13 */
str d14, [x0, 0xb0] /* FIBER_REG_V14 */
str d15, [x0, 0xb8] /* FIBER_REG_V15 */
// store sp and lr
mov x2, sp
str x2, [x0, 0xc0] /* FIBER_REG_SP */
str x30, [x0, 0x78] /* FIBER_REG_X30 */
// load context 'to'
mov x7, x1
ldr x16, [x7, 0x08] /* FIBER_REG_X16 */
ldr x17, [x7, 0x10] /* FIBER_REG_X17 */
ldr x18, [x7, 0x18] /* FIBER_REG_X18 */
ldr x19, [x7, 0x20] /* FIBER_REG_X19 */
ldr x20, [x7, 0x28] /* FIBER_REG_X20 */
ldr x21, [x7, 0x30] /* FIBER_REG_X21 */
ldr x22, [x7, 0x38] /* FIBER_REG_X22 */
ldr x23, [x7, 0x40] /* FIBER_REG_X23 */
ldr x24, [x7, 0x48] /* FIBER_REG_X24 */
ldr x25, [x7, 0x50] /* FIBER_REG_X25 */
ldr x26, [x7, 0x58] /* FIBER_REG_X26 */
ldr x27, [x7, 0x60] /* FIBER_REG_X27 */
ldr x28, [x7, 0x68] /* FIBER_REG_X28 */
ldr x29, [x7, 0x70] /* FIBER_REG_X29 */
ldr d8, [x7, 0x80] /* FIBER_REG_V8 */
ldr d9, [x7, 0x88] /* FIBER_REG_V9 */
ldr d10, [x7, 0x90] /* FIBER_REG_V10 */
ldr d11, [x7, 0x98] /* FIBER_REG_V11 */
ldr d12, [x7, 0xa0] /* FIBER_REG_V12 */
ldr d13, [x7, 0xa8] /* FIBER_REG_V13 */
ldr d14, [x7, 0xb0] /* FIBER_REG_V14 */
ldr d15, [x7, 0xb8] /* FIBER_REG_V15 */
// Load first parameter, this is only used for the first time a fiber gains control
ldr x0, [x7, 0x00] /* FIBER_REG_X0 */
ldr x30, [x7, 0x78] /* FIBER_REG_X30 */
ldr x2, [x7, 0xc0] /* FIBER_REG_SP */
mov sp, x2
ret
)");
#elif defined(FIBER_ASM_X64_IMPLEMENTATION)
#define FIBER_ASM_IMPLEMENTATION
#define FIBER_REG_RBX 0x00
#define FIBER_REG_RBP 0x08
#define FIBER_REG_R12 0x10
#define FIBER_REG_R13 0x18
#define FIBER_REG_R14 0x20
#define FIBER_REG_R15 0x28
#define FIBER_REG_RDI 0x30
#define FIBER_REG_RSP 0x38
#define FIBER_REG_RIP 0x40
//! Fiber context that saves all the callee saved registers
/**
* The specific set of register is architecture and OS dependent. SORT uses this implementation for
* Intel Mac and X64 Ubuntu, which use 'System V AMD64 ABI'.
* https://en.wikipedia.org/wiki/X86_calling_conventions#cite_note-AMD-28
*
* System V Application Binary Interface AMD64 Architecture Processor Supplement
* Page 23, AMD64 ABI Draft 1.0
* https://github.com/hjl-tools/x86-psABI/wiki/x86-64-psABI-1.0.pdf
*/
struct FiberContexInternal {
// callee-saved registers
Register rbx;
Register rbp;
Register r12;
Register r13;
Register r14;
Register r15;
// parameter registers
Register rdi;
// stack and instruction registers
Register rsp;
Register rip;
};
CHECK_OFFSET(rbx, FIBER_REG_RBX);
CHECK_OFFSET(rbp, FIBER_REG_RBP);
CHECK_OFFSET(r12, FIBER_REG_R12);
CHECK_OFFSET(r13, FIBER_REG_R13);
CHECK_OFFSET(r14, FIBER_REG_R14);
CHECK_OFFSET(r15, FIBER_REG_R15);
CHECK_OFFSET(rdi, FIBER_REG_RDI);
CHECK_OFFSET(rsp, FIBER_REG_RSP);
CHECK_OFFSET(rip, FIBER_REG_RIP);
// create a new fiber
inline bool _create_fiber_internal(void* stack, uint32_t stack_size, void (*target)(void*), void* arg, FiberContexInternal* context) {
// it is the users responsibility to make sure the stack is 16 bytes aligned, which is required by the Arm64 architecture
if((((uintptr_t)stack) & (FIBER_STACK_ALIGNMENT - 1)) != 0)
return false;
uintptr_t* stack_top = (uintptr_t*)((uint8_t*)(stack) + stack_size);
context->rip = (uintptr_t)target;
context->rdi = (uintptr_t)arg;
context->rsp = (uintptr_t)&stack_top[-3];
stack_top[-2] = 0;
return true;
}
// switching to a different fiber
asm(
R"(
.text
.align 4
_switch_fiber_internal:
// Save context 'from'
// Store callee-preserved registers
movq %rbx, 0x00(%rdi) /* FIBER_REG_RBX */
movq %rbp, 0x08(%rdi) /* FIBER_REG_RBP */
movq %r12, 0x10(%rdi) /* FIBER_REG_R12 */
movq %r13, 0x18(%rdi) /* FIBER_REG_R13 */
movq %r14, 0x20(%rdi) /* FIBER_REG_R14 */
movq %r15, 0x28(%rdi) /* FIBER_REG_R15 */
/* call stores the return address on the stack before jumping */
movq (%rsp), %rcx
movq %rcx, 0x40(%rdi) /* FIBER_REG_RIP */
/* skip the pushed return address */
leaq 8(%rsp), %rcx
movq %rcx, 0x38(%rdi) /* FIBER_REG_RSP */
// Load context 'to'
movq %rsi, %r8
// Load callee-preserved registers
movq 0x00(%r8), %rbx /* FIBER_REG_RBX */
movq 0x08(%r8), %rbp /* FIBER_REG_RBP */
movq 0x10(%r8), %r12 /* FIBER_REG_R12 */
movq 0x18(%r8), %r13 /* FIBER_REG_R13 */
movq 0x20(%r8), %r14 /* FIBER_REG_R14 */
movq 0x28(%r8), %r15 /* FIBER_REG_R15 */
// Load first parameter, this is only used for the first time a fiber gains control
movq 0x30(%r8), %rdi /* FIBER_REG_RDI */
// Load stack pointer
movq 0x38(%r8), %rsp /* FIBER_REG_RSP */
// Load instruction pointer, and jump
movq 0x40(%r8), %rcx /* FIBER_REG_RIP */
jmp *%rcx
)");
#elif defined(TINY_FIBER_WIN)
struct FiberContexInternal {
// handle to hold fiber
LPVOID raw_fiber_handle = nullptr;
};
#endif
#if defined(FIBER_ASM_IMPLEMENTATION)
#if __has_attribute (optnone)
#define FORCENOTINLINE __attribute__((optnone))
#elif __has_attribute (noinline)
#define FORCENOTINLINE __attribute__((noinline))
#else
#define FORCENOTINLINE
#endif
extern "C"
{
extern void FORCENOTINLINE _switch_fiber_internal(FiberContexInternal* from, const FiberContexInternal* to);
}
#endif
//! Abstruction for fiber struct.
struct Fiber {
/**< fiber context, this is platform dependent. */
FiberContexInternal context;
#if defined(FIBER_ASM_IMPLEMENTATION)
/**< Pointer to stack. */
void* stack_ptr = nullptr;
/**< Stack size. */
unsigned int stack_size = 0;
#elif defined(TINY_FIBER_WIN)
// whether the fiber is originated from a thread or not
bool is_fiber_from_thread = false;
#endif
};
// Define fiber handle
typedef Fiber* FiberHandle;
//! Create a new fiber.
//!
//! Allocate stack memory for the fiber. If there is no valid function pointer provided, it will fail.
inline FiberHandle CreateFiber(uint32_t stack_size, void (*fiber_func)(void*), void* arg) {
if(stack_size == 0 || !fiber_func || !arg)
return nullptr;
Fiber* ptr = (Fiber*)TINY_FIBER_MALLOC(sizeof(Fiber));
ptr->context = {};
#if defined(FIBER_ASM_IMPLEMENTATION)
ptr->stack_ptr = TINY_FIBER_MALLOC(stack_size + FIBER_STACK_ALIGNMENT - 1);
ptr->stack_size = stack_size;
if(!ptr->stack_ptr)
return nullptr;
// Make sure the stack meets the alignment requirement
uintptr_t aligned_stack_ptr = (uintptr_t)ptr->stack_ptr;
aligned_stack_ptr += FIBER_STACK_ALIGNMENT - 1;
aligned_stack_ptr &= ~(FIBER_STACK_ALIGNMENT - 1);
if(!_create_fiber_internal((void*)aligned_stack_ptr, stack_size, fiber_func, arg, &ptr->context))
return nullptr;
#elif defined(TINY_FIBER_WIN)
ptr->context.raw_fiber_handle = ::CreateFiber(stack_size, fiber_func, arg);
ptr->is_fiber_from_thread = false;
#endif
return ptr;
}
//! Convert the current thread to fiber.
//!
//! Note, on Ubuntu and Mac, this doesn't really convert the current thread to a new fiber, it really just creates a brand new fiber that has nothing to do
//! with the current thread. However, as long as the thread first switch to a fiber from this created fiber, this would allow the thread fiber to capture
//! the thread context.
inline FiberHandle CreateFiberFromThread() {
Fiber* ptr = (Fiber*)TINY_FIBER_MALLOC(sizeof(Fiber));
#if defined(FIBER_ASM_IMPLEMENTATION)
ptr->context = {};
ptr->stack_ptr = nullptr;
ptr->stack_size = 0;
#elif defined(TINY_FIBER_WIN)
ptr->context.raw_fiber_handle = ConvertThreadToFiber(nullptr);
ptr->is_fiber_from_thread = true;
#endif
return ptr;
}
//! Switch from one fiber to another different fiber
//!
//! The upper level logic should have checked for same fiber switch. The low level implementation defined here won't do such a check. In the case of
//! switching to a same fiber, it should not crash the system. But there could be some loss of performance. Another more important thing that derserve our
//! attention is that the high level code will have to make sure the from_fiber is the current executing fiber, otherwise there will be undefined behavior.
inline void SwitchFiber(Fiber* from_fiber, const Fiber* to_fiber) {
#if defined(FIBER_ASM_IMPLEMENTATION)
_switch_fiber_internal(&from_fiber->context, &to_fiber->context);
#elif defined(TINY_FIBER_WIN)
::SwitchToFiber(to_fiber->context.raw_fiber_handle);
#endif
}
//! Delete a fiber
//!
//! If the fiber is converted from a thread, delete fiber will convert the fiber back to a regular thread then.
inline void DeleteFiber(FiberHandle fiber_handle){
if(!fiber_handle)
return;
#if defined(TINY_FIBER_WIN)
if (fiber_handle->is_fiber_from_thread)
ConvertFiberToThread();
else
::DeleteFiber(fiber_handle->context.raw_fiber_handle);
fiber_handle->context.raw_fiber_handle = nullptr;
fiber_handle->is_fiber_from_thread = false;
#elif defined(FIBER_ASM_IMPLEMENTATION)
if(fiber_handle->stack_ptr){
TINY_FIBER_FREE(fiber_handle->stack_ptr);
fiber_handle->stack_ptr = nullptr;
}
#endif
TINY_FIBER_FREE(fiber_handle);
}
} // namespace tiny_fiber
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment