[WIP]: Add experimental fibers

Cargo.toml

@@ -72,6 +72,10 @@ targets = [
 name = "coroutine"
 harness = false
 
+[[bench]]
+name = "fiber"
+harness = false
+
 # We want debug symbols on release binaries by default. We can always strip them
 # out later if necessary.
 [profile.release]

benches/fiber.rs

@@ -0,0 +1,75 @@
+use std::convert::identity;
+
+use corosensei::stack::DefaultStack;
+use corosensei::{fiber, fiber_with_stack, Fiber};
+use criterion::measurement::Measurement;
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+
+fn fiber_switch<M: Measurement + 'static>(name: &str, c: &mut Criterion<M>) {
+    struct IdFiber(Fiber<(IdFiber, usize)>);
+
+    let mut identity = IdFiber(fiber::<std::convert::Infallible>().switch(|fiber| {
+        let (mut yielder, mut input): (IdFiber, usize) = fiber.switch(identity);
+        loop {
+            (yielder, input) = yielder.0.switch(move |fib| (IdFiber(fib), input))
+        }
+    }));
+
+    c.bench_function(name, |b| {
+        b.iter(|| {
+            identity.0.switch_in_place(
+                |yielder| (IdFiber(yielder), black_box(0usize)),
+                |(f, _)| (f.0, ()),
+            );
+        })
+    });
+}
+
+fn fiber_call<M: Measurement + 'static>(name: &str, c: &mut Criterion<M>) {
+    // Don't count time spent allocating a stack.
+    let mut stack = Some(DefaultStack::default());
+
+    c.bench_function(name, move |b| {
+        b.iter(|| {
+            let identity =
+                fiber_with_stack(move |input, stack| (stack, input), stack.take().unwrap());
+            stack = Some(identity.switch(|fib| (fib, black_box(0usize))).0)
+        })
+    });
+}
+
+fn fiber_switch_time(c: &mut Criterion) {
+    fiber_switch("fiber_switch_time", c);
+}
+fn fiber_call_time(c: &mut Criterion) {
+    fiber_call("fiber_call_time", c);
+}
+
+criterion_group!(
+    name = time;
+    config = Criterion::default();
+    targets = fiber_switch_time, fiber_call_time
+);
+
+cfg_if::cfg_if! {
+    if #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] {
+        use criterion_cycles_per_byte::CyclesPerByte;
+
+        fn fiber_switch_cycles(c: &mut Criterion<CyclesPerByte>) {
+            fiber_switch("fiber_switch_cycles", c);
+        }
+        fn fiber_call_cycles(c: &mut Criterion<CyclesPerByte>) {
+            fiber_call("fiber_call_cycles", c);
+        }
+
+        criterion_group!(
+            name = cycles;
+            config = Criterion::default().with_measurement(CyclesPerByte);
+            targets = fiber_switch_cycles, fiber_call_cycles
+        );
+
+        criterion_main!(cycles, time);
+    } else {
+        criterion_main!(time);
+    }
+}

examples/fiber.rs

@@ -0,0 +1,60 @@
+// This example is basically a reworked `example/basic.rs`
+
+use std::convert::identity;
+
+use corosensei::{fiber, Fiber};
+
+struct Arg {
+    exec: Fiber<Option<Yield>>,
+    i: i32,
+}
+
+struct Yield {
+    exec: Fiber<Arg>,
+    i: i32,
+}
+
+impl Arg {
+    fn new(exec: Fiber<Option<Yield>>, input: i32) -> Self {
+        Self { exec, i: input }
+    }
+}
+
+fn main() {
+    println!("[main] creating fiber");
+
+    let exec = fiber();
+
+    let mut exec = exec.switch(|exec| {
+        let Arg { mut exec, i } = exec.switch(identity);
+        println!("[fiber] fiber started with input {}", i);
+        for mut i in 0..5 {
+            println!("[fiber] yielding {}", i);
+            Arg { exec, i } = exec.switch(move |exec| Some(Yield { exec, i }));
+            println!("[fiber] got {} from parent", i)
+        }
+        println!("[fiber] exiting fiber");
+        (exec, None)
+    });
+
+    let mut counter = 100;
+    loop {
+        println!("[main] resuming fiber with argument {}", counter);
+        match exec.switch(move |exec| Arg::new(exec, counter)) {
+            Some(Yield { exec: new, i }) => {
+                exec = new;
+                println!("[main] got {:?} from fiber", i)
+            }
+            None => break,
+        }
+
+        counter += 1;
+    }
+
+    println!("[main] exiting");
+}
+
+#[test]
+fn fiber() {
+    main()
+}

src/arch/aarch64.rs

@@ -71,12 +71,13 @@
 //! ```
 
 use core::arch::{asm, global_asm};
+use core::{ffi, ptr};
 
 use super::{allocate_obj_on_stack, push};
 use crate::stack::{Stack, StackPointer};
 use crate::unwind::{
     asm_may_unwind_root, asm_may_unwind_yield, cfi_reset_args_size_root, cfi_reset_args_size_yield,
-    InitialFunc, StackCallFunc, TrapHandler,
+    InitialFunc, StackCallFunc, SwitchFiberFunc, TrapHandler,
 };
 use crate::util::EncodedValue;
 
@@ -185,6 +186,55 @@ extern "C" {
     fn stack_call_trampoline(arg: *mut u8, sp: StackPointer, f: StackCallFunc);
 }
 
+#[inline]
+pub unsafe fn fiber_init_stack(stack_base: StackPointer) -> StackPointer {
+    unsafe extern "C" fn entry(
+        sp: StackPointer,
+        arg: EncodedValue,
+        // There could be no output buffer on a fresh fiber.
+        // Assuming first `f` function never returns on this fiber.
+        _obj: *mut ffi::c_void,
+        f: SwitchFiberFunc,
+    ) {
+        f(sp, arg, ptr::null_mut())
+    }
+
+    let mut sp = stack_base.get();
+
+    // Put the entry function pointer
+    push(&mut sp, Some(entry as StackWord));
+
+    StackPointer::new_unchecked(sp)
+}
+
+#[inline]
+pub unsafe fn fiber_switch(
+    stack_ptr: StackPointer,
+    mut arg: EncodedValue,
+    ret: *mut ffi::c_void,
+    mut f: SwitchFiberFunc,
+) {
+    let mut sp = stack_ptr.get();
+    asm!(
+        "stp x19, x29, [sp, #-16]!",
+        "adr lr, 2f",
+        "str lr, [sp, #-8]!",
+        "ldr lr, [x0], #8",
+        "mov x9, sp",
+        "mov sp, x0",
+        "mov x0, x9",
+        "ret",
+        "2:",
+        "ldp x19, x29, [sp], #16",
+        inlateout("x0") sp, inlateout("x1") arg, inlateout("x3") f,
+        lateout("x20") _, lateout("x21") _, lateout("x22") _, lateout("x23") _,
+        lateout("x24") _, lateout("x25") _, lateout("x26") _, lateout("x27") _,
+        lateout("x28") _,
+        clobber_abi("C"),
+    );
+    f(StackPointer::new_unchecked(sp), arg, ret)
+}
+
 #[inline]
 pub unsafe fn init_stack<T>(stack: &impl Stack, func: InitialFunc<T>, obj: T) -> StackPointer {
     let mut sp = stack.base().get();

src/arch/riscv.rs

@@ -71,12 +71,13 @@
 //! ```
 
 use core::arch::{asm, global_asm};
+use core::{ffi, ptr};
 
 use super::{allocate_obj_on_stack, push};
 use crate::stack::{Stack, StackPointer};
 use crate::unwind::{
     asm_may_unwind_root, asm_may_unwind_yield, cfi_reset_args_size_root, cfi_reset_args_size_yield,
-    InitialFunc, StackCallFunc, TrapHandler,
+    InitialFunc, StackCallFunc, SwitchFiberFunc, TrapHandler,
 };
 use crate::util::EncodedValue;
 
@@ -255,6 +256,63 @@ extern "C" {
     fn stack_call_trampoline(arg: *mut u8, sp: StackPointer, f: StackCallFunc);
 }
 
+#[inline]
+pub unsafe fn fiber_init_stack(stack_base: StackPointer) -> StackPointer {
+    unsafe extern "C" fn entry(
+        sp: StackPointer,
+        arg: EncodedValue,
+        // There could be no output buffer on a fresh fiber.
+        // Assuming first `f` function never returns on this fiber.
+        _obj: *mut ffi::c_void,
+        f: SwitchFiberFunc,
+    ) {
+        f(sp, arg, ptr::null_mut())
+    }
+
+    let mut sp = stack_base.get();
+
+    // Put the entry function pointer
+    push(&mut sp, Some(entry as StackWord));
+
+    StackPointer::new_unchecked(sp)
+}
+
+#[inline]
+pub unsafe fn fiber_switch(
+    stack_ptr: StackPointer,
+    mut arg: EncodedValue,
+    ret: *mut ffi::c_void,
+    mut f: SwitchFiberFunc,
+) {
+    let mut sp = stack_ptr.get();
+    asm!(
+        addi!("sp", "sp", -3),
+        s!("s0", 2, "sp"),
+        s!("s1", 1, "sp"),
+        "lla ra, 2f",
+        s!("ra", 0, "sp"),
+        l!("ra", 0, "a0"),
+        "mv t1, sp",
+        addi!("sp", "a0", 1),
+        "mv a0, t1",
+        "ret",
+        "2:",
+        l!("s1", 0, "sp"),
+        l!("s0", 1, "sp"),
+        addi!("sp", "sp", 2),
+
+        inlateout("a0") sp, inlateout("a1") arg, inlateout("a3") f,
+        lateout("s2") _, lateout("s3") _, lateout("s4") _, lateout("s5") _,
+        lateout("s6") _, lateout("s7") _, lateout("s8") _, lateout("s9") _,
+        lateout("s10") _, lateout("s11") _,
+        lateout("fs0") _, lateout("fs1") _, lateout("fs2") _, lateout("fs3") _,
+        lateout("fs4") _, lateout("fs5") _, lateout("fs6") _, lateout("fs7") _,
+        lateout("fs8") _, lateout("fs9") _, lateout("fs10") _, lateout("fs11") _,
+        clobber_abi("C"),
+    );
+    f(StackPointer::new_unchecked(sp), arg, ret)
+}
+
 #[inline]
 pub unsafe fn init_stack<T>(stack: &impl Stack, func: InitialFunc<T>, obj: T) -> StackPointer {
     let mut sp = stack.base().get();

src/arch/x86_64.rs

@@ -99,12 +99,13 @@
 //! ```
 
 use core::arch::{asm, global_asm};
+use core::{ffi, ptr};
 
 use super::{allocate_obj_on_stack, push};
 use crate::stack::{Stack, StackPointer};
 use crate::unwind::{
     asm_may_unwind_root, asm_may_unwind_yield, cfi_reset_args_size_root, cfi_reset_args_size_yield,
-    InitialFunc, StackCallFunc, TrapHandler,
+    InitialFunc, StackCallFunc, SwitchFiberFunc, TrapHandler,
 };
 use crate::util::EncodedValue;
 
@@ -283,6 +284,61 @@ extern "C" {
     fn stack_call_trampoline(arg: *mut u8, sp: StackPointer, f: StackCallFunc);
 }
 
+#[inline]
+pub unsafe fn fiber_init_stack(stack_base: StackPointer) -> StackPointer {
+    unsafe extern "sysv64" fn entry(
+        sp: StackPointer,
+        arg: EncodedValue,
+        // There could be no output buffer on a fresh fiber.
+        // Assuming first `f` function never returns on this fiber.
+        _obj: *mut ffi::c_void,
+        f: SwitchFiberFunc,
+    ) {
+        f(sp, arg, ptr::null_mut())
+    }
+
+    let mut sp = stack_base.get();
+
+    // Zero initialize return pointer
+    push(&mut sp, Some(0));
+    // Put the entry function pointer
+    push(&mut sp, Some(entry as StackWord));
+
+    StackPointer::new_unchecked(sp)
+}
+
+#[inline]
+pub unsafe fn fiber_switch(
+    stack_ptr: StackPointer,
+    mut arg: EncodedValue,
+    ret: *mut ffi::c_void,
+    mut f: SwitchFiberFunc,
+) {
+    let mut sp = stack_ptr.get();
+    // TODO: edit the comment below
+    // The function `fiber_switch` upon its call saves `ret: rdx` output pointer argument and
+    // callee-saved registers (defined by to the AMD64 System-V ABI) to the stack.  So it switches
+    // stacks, restores callee-saved registers and an output pointer argument, then finally it jumps
+    // to the input function pointer in rcx.
+    asm!(
+        "lea rax, [rip + 2f]",
+        "push rbx",
+        "push rbp",
+        "push rax",
+        "mov rax, rsp",
+        "mov rsp, rdi",
+        "mov rdi, rax",
+        "ret",
+        "2:",
+        "pop rbp",
+        "pop rbx",
+        inout("rdi") sp, inout("rsi") arg, inout("rcx") f,
+        lateout("r12") _, lateout("r13") _, lateout("r14") _, lateout("r15") _,
+        clobber_abi("sysv64"),
+    );
+    f(StackPointer::new_unchecked(sp), arg, ret)
+}
+
 /// Sets up the initial state on a stack so that the given function is
 /// executed on the first switch to this stack.
 ///