Implementing cooperative multitasking in Rust

// Warning, untested and doesn't work
macro_rules! asm_function {
    (pub unsafe extern "C" fn $name:ident($($arg:ident : $arg_type:ty),*) -> $ret_type:ty { $asm:expr } ) => (
        extern "C" {
            pub fn $name($($arg: $arg_type),*) -> $ret_type;
        }

        global_asm!(
            concat!(
                "   .global ", $name, "\n",
                $name, ":\n",
                $asm
            )
        );
    );
    (pub unsafe extern "C" fn $name:ident($($arg:ident : $arg_type:ty),*) { $asm: expr } ) => (
        asm_function!(pub unsafe extern "C" $name ($($arg : $arg_type),*)) -> () $asm
    )
}

// We make a newtype here for stack pointers to help prevent them being
// misused (particularly useful is that they are given move semantics, so
// users can't make copies and try to switch to the same stack pointer twice).
#[repr(C)]
pub struct StackPtr(usize);

extern "C" {
    // Our start function takes 3 arguments: A pointer to the end of a buffer that 
    // we will use as a stack. A pointer to a function that we will call, which in turn
    //  - Takes two arguments (a pointer to the stack we called it on, and a pointer
    //    to some argument).
    //  - Must never return.
    // Finally we take a pointer to the arguments to pass to the start function.

    // Unlike exit_to this function can return, at least as the calling code can observe, 
    // when a stack uses `exit_to` or `switch` to switch back to us. It returns no value.

    // It's already unsafe to call this function, so it doesn't make it any more unsafe 
    // to call a user supplied function which is also unsafe. Hence we allow for that.
    pub fn start(buf: *mut u8, 
                 to_call: unsafe extern "C" fn(StackPtr, *mut()) -> !, 
                 args: *mut ());

    // The simplest function. We take a stack pointer (that we will switch to), and we 
    // never  return to this stack again (we can't since we forget where the top of it is).
    pub fn exit_to(stack: StackPtr) -> !;

    // Takes a stack pointer to switch to, and a pointer to which we are going to store 
    // our current stack pointer. Like start this function can return indirectly.
    pub fn switch(target: StackPtr, save: *mut StackPtr);
}

// This is a good point to note that all our global runtime state will be 
// unique to each OS thread, avoiding concurrency issues.

// thread_local! is a macro that makes thread local global variables (each thread 
// gets it own copy), which are lazily initialized on first use.
thread_local! {
    // We're going to need a list of stacks waiting to execute.
    static WAITING_THREADS: RefCell<Vec<StackPtr>> = RefCell::new(vec!());

    // This holds a stack which just figures out what to execute next and switches
    // to it. Which ended up being a very convenient way to implement yielding execution.
    static YIELD_THREAD: UnsafeCell<StackPtr> = UnsafeCell::new(make_yield_thread());
}

/// Takes a closure and runs it on a new coroutine
pub fn start<F: FnOnce() -> () + 'static + panic::UnwindSafe >(f: F) {
    let stack = Stack::new(32 * 1024 - 8);

    primitives::start(stack, move |stack, parent| {
        // Add our parent to the list of threads waiting to execute.
        // The `.with` syntax here is just how variables defined in thread_local! have
        // to be accessed (the argument to the closure is a reference to the contents of
        // the variable).
        WAITING_THREADS.with(|threads| {
            threads.borrow_mut().push(parent)
        });

        // We catch panic's here so the coroutine panicking doesn't kill the entire OS thread.
        let res = panic::catch_unwind(f);

        if let Err(_) = res {
            eprintln!("Coroutine panicked, killing");
        }

        unsafe {
            // get the stack ptr for the yield thread from the thread local 
            let thread = YIELD_THREAD.with(|it| ptr::read(it.get()));
            // and switch to it
            primitives::exit_to(thread);
        }
    });
}

fn make_yield_thread() -> StackPtr {
    // This is a unstable way of making thread locals that don't need to be lazily
    // initialized and accessed via a closure/.with function that was convenient here.
    // We use this to store the result of a `switch` call, so the new stack can return it.
    #[thread_local]
    static mut RET: StackPtr = NULL_STACK_PTR;

    // Stack a new coroutine
    let stack = Stack::new(16 * 1024 - 8);
    primitives::start(stack, move |_stack, parent| {
        // Immediately switches back to it's parent, we start doing work once a thread yields.
        unsafe{ primitives::switch(parent, &mut RET) };

        WAITING_THREADS.with( |threads|
            loop {
                // Whenever we switch to this coroutine switch to the first waiting thread.
                let maybe_thread = threads.borrow_mut().pop();
                match maybe_thread {
                    Some(thread) => {
                        let ptr = YIELD_THREAD.with(|it| it.get());
                        unsafe{ primitives::unsafe_switch(thread, ptr) };
                    },
                    None => {
                        // For now (with no non blocking IO) we know a waiting thread exists.
                        unreachable!();
                    }
                }
            }
        )
    });
    unsafe{ return ptr::read(&RET) };
}

thread_local! {
    static INACTIVE_STACKS: RefCell<Vec<Stack>> = RefCell::new(vec!());
}

// Get a existing, unused stack if one exists.
let mut maybe_stack = None;
INACTIVE_STACKS.with(|stacks| maybe_stack = stacks.borrow_mut().pop());
// Otherwise allocate a new one.
let stack = maybe_stack.unwrap_or_else(|| Stack::new(32 * 1024 - 8));

INACTIVE_STACKS.with(|stacks| {
    stacks.borrow_mut().push(stack);
});

thread_local! {
    // Global MIO objects
    static MIO_POLL: RefCell<mio::Poll> = RefCell::new(mio::Poll::new().unwrap());
    static MIO_EVENTS: RefCell<mio::Events> = RefCell::new(mio::Events::with_capacity(1024));

    // A count of coroutines currently blocked waiting on mio.
    static MIO_ACTIVE: Cell<usize> = Cell::new(0);

    // Information to keep track of what stack we should start when mio tells us something is
    // ready to execute.
    static MIO_TOKEN_MAP: RefCell<Vec<StackPtr>> = RefCell::new(vec!());
    static MIO_FREE_TOKEN: RefCell<Vec<usize>> = RefCell::new(vec!());
}

pub fn await<E: mio::Evented>(handle: &E, interest: mio::Ready) {
    // Handle is something like a TcpStream, interest is something like 
    // Readable (unblock when the stream is ready to be read from)


    MIO_POLL.with(|poll| {
        MIO_TOKEN_MAP.with(|map| {
            // Get a token which mio will return when we are ready to execute again.
            // (and allocate space in the map from tokens to stack pointers if needed).

            let maybe_token = MIO_FREE_TOKEN.with(|free|
                free.borrow_mut().pop()
            );

            let token = maybe_token.unwrap_or_else(|| {
                let mut map = map.borrow_mut();
                map.push(NULL_STACK_PTR);
                map.len() - 1
            });

            {
                // Register the token with mio, handle and interest are passed in by
                // the user (described above). The logic to reregister (needed if the
                // token has been used before) is a bit ugly here but it works.

                let poll = poll.borrow();
                let res = poll.register(handle, mio::Token(token), interest,
                    mio::PollOpt::oneshot());
                if !res.is_ok() {
                    poll.reregister(handle, mio::Token(token), interest,
                        mio::PollOpt::oneshot()).expect("Failed to register");
                }
            }

            // Increment counter.
            MIO_ACTIVE.with(|active| active.set(active.get() + 1));

            // Switch threads, storing the stack pointer to switch back
            // in the map associated with the token we registered with mio.
            unsafe {
                let thread = YIELD_THREAD.with(|it| ptr::read(it.get()));
                let ptr = {
                    // Need scope to force drop of map (?)
                    let mut mbm = map.borrow_mut();
                    let ptr = &mut mbm[token] as *mut StackPtr;
                    ptr
                };
                primitives::unsafe_switch(thread, ptr);
            }
        })
    });
}

MIO_EVENTS.with(|events| {
    // Poll mio for events (blocking until something is ready).
    MIO_POLL.with(|poll|
        poll.borrow().poll(&mut *events.borrow_mut(), None).expect("mio poll failed")
    );

    // For each event, switch to it's thread in turn. Notice that we will stay in this 
    // portion of the code until we have handled every single event. (This is why
    // implementing this as a separate stack was so convenient).
    for event in events.borrow().iter() {
        let thread = get_stack(event);
        MIO_ACTIVE.with(|active| active.set(active.get() - 1));

        let ptr = YIELD_THREAD.with(|it| it.get());
        unsafe{ primitives::unsafe_switch(thread, ptr) };
    }
});

pub fn run() {
    let mut finished = false;
    loop {
        // Check if there are any waiting of blocked threads
        WAITING_THREADS.with(|threads|
            MIO_ACTIVE.with(|active|
                finished = active.get() == 0 && threads.borrow().is_empty()
            )
        );

        // If not, return
        if finished { break; }

        // Otherwise go run one of them (waiting for one to unblock if necessary).
        unsafe {
            let thread = YIELD_THREAD.with(|it| ptr::read(it.get()));
            primitives::unsafe_switch(thread, new_waiting());
        }
    }
}