cogo/

scoped.rs

// modified from crossbeam

use std::cell::RefCell;
use std::fmt;
use std::mem;
use std::panic;
use std::rc::Rc;
use std::sync::Arc;
use std::thread;

use crate::coroutine_impl::{spawn, Coroutine};
use crate::join::JoinHandle;
use crossbeam::atomic::AtomicCell;

/// Like `coroutine::spawn`, but without the closure bounds.
pub unsafe fn spawn_unsafe<'a, F>(f: F) -> JoinHandle<()>
    where
        F: FnOnce() + Send + 'a,
{
    let closure: Box<dyn FnOnce() + 'a> = Box::new(f);
    let closure: Box<dyn FnOnce() + Send> = mem::transmute(closure);
    spawn(move || closure())
}

pub struct Scope<'a> {
    dtors: RefCell<Option<DtorChain<'a>>>,
}

struct DtorChain<'a> {
    dtor: Box<dyn FnOnce() + 'a>,
    next: Option<Box<DtorChain<'a>>>,
}

enum JoinState {
    Running(JoinHandle<()>),
    Joined,
}

impl JoinState {
    fn join(&mut self) {
        let mut state = JoinState::Joined;
        mem::swap(self, &mut state);
        if let JoinState::Running(handle) = state {
            let res = handle.join();

            // TODO: when panic happened, the logic need to refine
            if !thread::panicking() {
                res.unwrap_or_else(|e| panic::resume_unwind(e));
            }
        }
    }
}

/// A handle to a scoped coroutine
pub struct ScopedJoinHandle<T> {
    inner: Rc<RefCell<JoinState>>,
    packet: Arc<AtomicCell<Option<T>>>,
    co: Coroutine,
}

/// Create a new `scope`, for deferred destructors.
///
/// Scopes, in particular, support scoped coroutine spawning.
///
pub fn scope<'a, F, R>(f: F) -> R
    where
        F: FnOnce(&Scope<'a>) -> R,
{
    let mut scope = Scope {
        dtors: RefCell::new(None),
    };
    let ret = f(&scope);
    scope.drop_all();
    ret
}

impl<'a> fmt::Debug for Scope<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Scope {{ ... }}")
    }
}

impl<T> fmt::Debug for ScopedJoinHandle<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ScopedJoinHandle {{ ... }}")
    }
}

impl<'a> Scope<'a> {
    // This method is carefully written in a transactional style, so
    // that it can be called directly and, if any dtor panics, can be
    // resumed in the unwinding this causes. By initially running the
    // method outside of any destructor, we avoid any leakage problems
    // due to @rust-lang/rust#14875.
    fn drop_all(&mut self) {
        loop {
            // use a separate scope to ensure that the RefCell borrow
            // is relinquished before running `dtor`
            let dtor = {
                let mut dtors = self.dtors.borrow_mut();
                if let Some(mut node) = dtors.take() {
                    *dtors = node.next.take().map(|b| *b);
                    node.dtor
                } else {
                    return;
                }
            };
            dtor();
        }
    }

    /// Schedule code to be executed when exiting the scope.
    ///
    /// This is akin to having a destructor on the stack, except that it is
    /// *guaranteed* to be run.
    pub fn defer<F>(&self, f: F)
        where
            F: FnOnce() + 'a,
    {
        let mut dtors = self.dtors.borrow_mut();
        *dtors = Some(DtorChain {
            dtor: Box::new(f),
            next: dtors.take().map(Box::new),
        });
    }

    /// Create a scoped coroutine.
    ///
    /// `spawn` is similar to the `spawn` function in this library. The
    /// difference is that this coroutine is scoped, meaning that it's guaranteed to terminate
    /// before the current stack frame goes away, allowing you to reference the parent stack frame
    /// directly. This is ensured by having the parent join on the child coroutine before the
    /// scope exits.
    fn spawn_impl<F, T>(&self, f: F) -> ScopedJoinHandle<T>
        where
            F: FnOnce() -> T + Send + 'a,
            T: Send + 'a,
    {
        let their_packet = Arc::new(AtomicCell::new(None));
        let my_packet = their_packet.clone();

        let join_handle = unsafe {
            spawn_unsafe(move || {
                their_packet.swap(Some(f()));
            })
        };

        let co = join_handle.coroutine().clone();
        let deferred_handle = Rc::new(RefCell::new(JoinState::Running(join_handle)));
        let my_handle = deferred_handle.clone();

        self.defer(move || {
            let mut state = deferred_handle.borrow_mut();
            state.join();
        });

        ScopedJoinHandle {
            inner: my_handle,
            packet: my_packet,
            co,
        }
    }

    /// Create a scoped coroutine.
    ///
    /// `spawn` is similar to the `spawn` function in this library. The
    /// difference is that this coroutine is scoped, meaning that it's guaranteed to terminate
    /// before the current stack frame goes away, allowing you to reference the parent stack frame
    /// directly. This is ensured by having the parent join on the child coroutine before the
    /// scope exits.
    pub unsafe fn spawn<F, T>(&self, f: F) -> ScopedJoinHandle<T>
        where
            F: FnOnce() -> T + Send + 'a,
            T: Send + 'a,
    {
        self.spawn_impl(f)
    }
}

impl<T> ScopedJoinHandle<T> {
    /// Join the scoped coroutine, returning the result it produced.
    pub fn join(self) -> T {
        self.inner.borrow_mut().join();
        self.packet.take().unwrap()
    }

    /// Get the underlying coroutine handle.
    pub fn coroutine(&self) -> &Coroutine {
        &self.co
    }
}

impl<'a> Drop for Scope<'a> {
    fn drop(&mut self) {
        self.drop_all()
    }
}