use std::any::Any;
use std::fmt;
use std::marker::PhantomData;
use std::panic;
use std::thread;
use crate::reg_context::RegContext;
use crate::rt::{Context, ContextStack, Error};
use crate::scope::Scope;
use crate::yield_::yield_now;
pub const DEFAULT_STACK_SIZE: usize = 0x1000;
pub struct Gn<A = ()> {
dummy: PhantomData<A>,
}
pub type Generator<'a, A, T> = Box<GeneratorImpl<'a, A, T>>;
unsafe impl<A, T> Send for GeneratorImpl<'static, A, T> {}
impl<A> Gn<A> {
pub fn new_scoped<'a, T, F>(f: F) -> Generator<'a, A, T>
where
F: FnOnce(Scope<A, T>) -> T + 'a,
T: 'a,
A: 'a,
{
Self::new_scoped_opt(DEFAULT_STACK_SIZE, f)
}
pub fn new_scoped_opt<'a, T, F>(size: usize, f: F) -> Generator<'a, A, T>
where
F: FnOnce(Scope<A, T>) -> T + 'a,
T: 'a,
A: 'a,
{
let mut g = GeneratorImpl::<A, T>::new(size);
g.scoped_init(f);
g
}
}
impl<A: Any> Gn<A> {
#[cfg_attr(feature = "cargo-clippy", allow(clippy::new_ret_no_self))]
pub fn new<'a, T: Any, F>(f: F) -> Generator<'a, A, T>
where
F: FnOnce() -> T + 'a,
{
Self::new_opt(DEFAULT_STACK_SIZE, f)
}
pub fn new_opt<'a, T: Any, F>(size: usize, f: F) -> Generator<'a, A, T>
where
F: FnOnce() -> T + 'a,
{
let mut g = GeneratorImpl::<A, T>::new(size);
g.init_context();
g.init(f);
g
}
}
pub struct GeneratorImpl<'a, A, T> {
context: Context,
para: Option<A>,
ret: Option<T>,
f: Option<Box<dyn FnOnce() + 'a>>,
}
impl<'a, A: Any, T: Any> GeneratorImpl<'a, A, T> {
pub fn init_context(&mut self) {
self.context.para = &mut self.para as &mut dyn Any;
self.context.ret = &mut self.ret as &mut dyn Any;
}
}
impl<'a, A, T> GeneratorImpl<'a, A, T> {
pub fn new(size: usize) -> Box<Self> {
Box::new(GeneratorImpl {
para: None,
ret: None,
f: None,
context: Context::new(size),
})
}
#[inline]
pub fn prefetch(&self) {
self.context.regs.prefetch();
}
pub fn scoped_init<F: FnOnce(Scope<'a, A, T>) -> T + 'a>(&mut self, f: F)
where
T: 'a,
A: 'a,
{
use std::mem::transmute;
let scope = unsafe { transmute(Scope::new(&mut self.para, &mut self.ret)) };
self.init(move || f(scope));
}
pub fn init<F: FnOnce() -> T + 'a>(&mut self, f: F)
where
T: 'a,
{
if self.f.is_none() && self.context._ref == 0 {
unsafe {
self.cancel();
}
}
self.context.parent = &mut self.context;
self.context._ref = 0;
let ret = &mut self.ret as *mut _;
let context = &mut self.context as *mut Context;
self.f = Some(Box::new(move || {
let r = f();
let ret = unsafe { &mut *ret };
let _ref = unsafe { (*context)._ref };
if _ref == 0xf {
::std::mem::forget(r);
*ret = None;
} else {
*ret = Some(r);
}
}));
let stk = &self.context.stack;
self.context
.regs
.init_with(gen_init, 0, &mut self.f as *mut _ as *mut usize, stk);
}
#[inline]
fn resume_gen(&mut self) {
let env = ContextStack::current();
let cur = &mut env.top().regs;
debug_assert!(!self.context.parent.is_null());
let top = unsafe { &mut *self.context.parent };
env.push_context(&mut self.context);
RegContext::swap(cur, &top.regs);
if !self.context.local_data.is_null() {
return;
}
if let Some(err) = self.context.err.take() {
#[cold]
panic::resume_unwind(err);
}
}
#[inline]
fn is_started(&self) -> bool {
self.f.is_none()
}
#[inline]
pub fn set_para(&mut self, para: A) {
self.para = Some(para);
}
#[inline]
pub fn set_local_data(&mut self, data: *mut u8) {
self.context.local_data = data;
}
#[inline]
pub fn get_local_data(&self) -> *mut u8 {
self.context.local_data
}
#[inline]
pub fn get_panic_data(&mut self) -> Option<Box<dyn Any + Send>> {
self.context.err.take()
}
#[inline]
pub fn resume(&mut self) -> Option<T> {
if self.is_done() {
#[cold]
return None;
}
self.context._ref += 1;
self.resume_gen();
self.ret.take()
}
#[inline]
pub fn raw_send(&mut self, para: Option<A>) -> Option<T> {
if self.is_done() {
#[cold]
return None;
}
self.para = para;
self.context._ref += 1;
self.resume_gen();
self.ret.take()
}
pub fn send(&mut self, para: A) -> T {
let ret = self.raw_send(Some(para));
ret.expect("send got None return")
}
#[inline]
unsafe fn raw_cancel(&mut self) {
self.context._ref = 2;
let old = ::std::panic::take_hook();
::std::panic::set_hook(Box::new(|_| {}));
self.resume_gen();
::std::panic::set_hook(old);
}
pub unsafe fn cancel(&mut self) {
if self.is_done() {
return;
}
if !self.is_started() {
self.f.take();
self.context._ref = 1;
} else {
self.raw_cancel();
}
}
#[inline]
pub fn is_done(&self) -> bool {
self.is_started() && (self.context._ref & 0x3) != 0
}
pub fn stack_usage(&self) -> (usize, usize) {
(
self.context.stack.size(),
self.context.stack.get_used_size(),
)
}
}
impl<'a, A, T> Drop for GeneratorImpl<'a, A, T> {
fn drop(&mut self) {
if thread::panicking() {
return;
}
if !self.is_started() {
return;
}
if !self.is_done() {
warn!("generator is not done while drop");
unsafe { self.raw_cancel() }
}
assert!(self.is_done());
let (total_stack, used_stack) = self.stack_usage();
if used_stack < total_stack {
} else {
error!("stack overflow detected!");
panic!(Error::StackErr);
}
}
}
impl<'a, T> Iterator for GeneratorImpl<'a, (), T> {
type Item = T;
fn next(&mut self) -> Option<T> {
self.resume()
}
}
impl<'a, A, T> fmt::Debug for GeneratorImpl<'a, A, T> {
#[cfg(nightly)]
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
use std::intrinsics::type_name;
write!(
f,
"Generator<{}, Output={}> {{ ... }}",
unsafe { type_name::<A>() },
unsafe { type_name::<T>() }
)
}
#[cfg(not(nightly))]
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "Generator {{ ... }}")
}
}
fn gen_init(_: usize, f: *mut usize) -> ! {
let clo = move || {
let f: &mut Option<Box<dyn FnOnce()>> = unsafe { &mut *(f as *mut _) };
let func = f.take().unwrap();
func();
};
fn check_err(cause: Box<dyn Any + Send + 'static>) {
if let Some(&Error::Cancel) = cause.downcast_ref::<Error>() {
return;
}
error!("set panicked inside generator");
ContextStack::current().top().err = Some(cause);
}
if let Err(cause) = panic::catch_unwind(clo) {
check_err(cause);
}
yield_now();
unreachable!("Should never comeback");
}