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::stack::{Func, Stack, StackBox};
use crate::yield_::yield_now;
pub const DEFAULT_STACK_SIZE: usize = 0x1000;
pub struct GeneratorObj<'a, A, T, const LOCAL: bool> {
gen: StackBox<GeneratorImpl<'a, A, T>>,
}
pub type Generator<'a, A, T> = GeneratorObj<'a, A, T, false>;
unsafe impl<A: Send, T: Send> Send for Generator<'static, A, T> {}
impl<'a, A, T> Generator<'a, A, T> {
pub fn scoped_init<F: FnOnce(Scope<'a, A, T>) -> T + Send + 'a>(&mut self, f: F)
where
T: Send + 'a,
A: Send + 'a,
{
self.gen.scoped_init(f);
}
pub fn init_code<F: FnOnce() -> T + Send + 'a>(&mut self, f: F)
where
T: Send + 'a,
{
self.gen.init_code(f);
}
}
pub type LocalGenerator<'a, A, T> = GeneratorObj<'a, A, T, true>;
impl<'a, A, T> LocalGenerator<'a, A, T> {
pub fn scoped_init<F: FnOnce(Scope<'a, A, T>) -> T + 'a>(&mut self, f: F)
where
T: 'a,
A: 'a,
{
self.gen.scoped_init(f);
}
}
impl<'a, A, T, const LOCAL: bool> GeneratorObj<'a, A, T, LOCAL> {
#[inline]
pub unsafe fn from_raw(raw: *mut usize) -> Self {
GeneratorObj {
gen: StackBox::from_raw(raw as *mut GeneratorImpl<'a, A, T>),
}
}
#[inline]
pub fn into_raw(self) -> *mut usize {
let ret = self.gen.as_ptr() as *mut usize;
std::mem::forget(self);
ret
}
#[inline]
pub fn prefetch(&self) {
self.gen.prefetch();
}
#[inline]
pub fn set_para(&mut self, para: A) {
self.gen.set_para(para);
}
#[inline]
pub fn set_local_data(&mut self, data: *mut u8) {
self.gen.set_local_data(data);
}
#[inline]
pub fn get_local_data(&self) -> *mut u8 {
self.gen.get_local_data()
}
#[inline]
pub fn get_panic_data(&mut self) -> Option<Box<dyn Any + Send>> {
self.gen.get_panic_data()
}
#[inline]
pub fn resume(&mut self) -> Option<T> {
self.gen.resume()
}
#[inline]
pub fn raw_send(&mut self, para: Option<A>) -> Option<T> {
self.gen.raw_send(para)
}
pub fn send(&mut self, para: A) -> T {
self.gen.send(para)
}
pub fn cancel(&mut self) {
self.gen.cancel()
}
#[inline]
pub fn is_done(&self) -> bool {
self.gen.is_done()
}
pub fn stack_usage(&self) -> (usize, usize) {
self.gen.stack_usage()
}
}
impl<'a, T, const LOCAL: bool> Iterator for GeneratorObj<'a, (), T, LOCAL> {
type Item = T;
fn next(&mut self) -> Option<T> {
self.resume()
}
}
impl<'a, A, T, const LOCAL: bool> fmt::Debug for GeneratorObj<'a, A, T, LOCAL> {
#[cfg(nightly)]
#[allow(unused_unsafe)]
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
use std::intrinsics::type_name;
write!(
f,
"Generator<{}, Output={}, Local={}> {{ ... }}",
unsafe { type_name::<A>() },
unsafe { type_name::<T>() },
LOCAL
)
}
#[cfg(not(nightly))]
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "Generator {{ ... }}")
}
}
pub struct Gn<A = ()> {
dummy: PhantomData<A>,
}
impl<A> Gn<A> {
pub fn new_scoped<'a, T, F>(f: F) -> Generator<'a, A, T>
where
F: FnOnce(Scope<A, T>) -> T + Send + 'a,
T: Send + 'a,
A: Send + 'a,
{
Self::new_scoped_opt(DEFAULT_STACK_SIZE, f)
}
pub fn new_scoped_local<'a, T, F>(f: F) -> LocalGenerator<'a, A, T>
where
F: FnOnce(Scope<A, T>) -> T + 'a,
T: 'a,
A: 'a,
{
Self::new_scoped_opt_local(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 + Send + 'a,
T: Send + 'a,
A: Send + 'a,
{
let mut gen = GeneratorImpl::<A, T>::new(Stack::new(size));
gen.scoped_init(f);
Generator { gen }
}
pub fn new_scoped_opt_local<'a, T, F>(size: usize, f: F) -> LocalGenerator<'a, A, T>
where
F: FnOnce(Scope<A, T>) -> T + 'a,
T: 'a,
A: 'a,
{
let mut gen = GeneratorImpl::<A, T>::new(Stack::new(size));
gen.scoped_init(f);
LocalGenerator { gen }
}
}
impl<A: Any> Gn<A> {
#[cfg_attr(feature = "cargo-clippy", allow(clippy::new_ret_no_self))]
#[deprecated(since = "0.6.18", note = "please use `scope` version instead")]
pub fn new<'a, T: Any, F>(f: F) -> Generator<'a, A, T>
where
F: FnOnce() -> T + Send + '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 + Send + 'a,
{
let mut gen = GeneratorImpl::<A, T>::new(Stack::new(size));
gen.init_context();
gen.init_code(f);
Generator { gen }
}
}
#[repr(C)]
struct GeneratorImpl<'a, A, T> {
context: Context,
stack: Stack,
para: Option<A>,
ret: Option<T>,
f: Option<Func>,
phantom: PhantomData<&'a T>,
}
impl<'a, A: Any, T: Any> GeneratorImpl<'a, A, T> {
fn init_context(&mut self) {
unsafe {
std::ptr::write(
self.context.para.as_mut_ptr(),
&mut self.para as &mut dyn Any,
);
std::ptr::write(self.context.ret.as_mut_ptr(), &mut self.ret as &mut dyn Any);
}
}
}
impl<'a, A, T> GeneratorImpl<'a, A, T> {
fn new(mut stack: Stack) -> StackBox<Self> {
unsafe {
let mut stack_box = stack
.alloc_uninit_box::<GeneratorImpl<'a, A, T>>()
.assume_init();
stack_box.init(GeneratorImpl {
para: None,
stack,
ret: None,
f: None,
context: Context::new(),
phantom: PhantomData,
});
stack_box
}
}
#[inline]
pub fn prefetch(&self) {
self.context.regs.prefetch();
}
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_code(move || f(scope));
}
fn init_code<F: FnOnce() -> T + 'a>(&mut self, f: F)
where
T: 'a,
{
if self.f.is_none() && self.context._ref == 0 {
self.cancel();
} else {
let _ = self.f.take();
}
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;
let func = StackBox::new_fn_once(&mut self.stack, 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);
}
});
self.f = Some(func);
self.context.regs.init_with(
gen_init,
0,
&mut self.f as *mut _ as *mut usize,
&self.stack,
);
}
#[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() {
panic::resume_unwind(err);
}
}
#[inline]
fn is_started(&self) -> bool {
self.f.is_none()
}
#[inline]
fn set_para(&mut self, para: A) {
self.para = Some(para);
}
#[inline]
fn set_local_data(&mut self, data: *mut u8) {
self.context.local_data = data;
}
#[inline]
fn get_local_data(&self) -> *mut u8 {
self.context.local_data
}
#[inline]
fn get_panic_data(&mut self) -> Option<Box<dyn Any + Send>> {
self.context.err.take()
}
#[inline]
fn resume(&mut self) -> Option<T> {
if self.is_done() {
return None;
}
self.context._ref += 1;
self.resume_gen();
self.ret.take()
}
#[inline]
fn raw_send(&mut self, para: Option<A>) -> Option<T> {
if self.is_done() {
return None;
}
self.para = para;
self.context._ref += 1;
self.resume_gen();
self.ret.take()
}
fn send(&mut self, para: A) -> T {
let ret = self.raw_send(Some(para));
ret.expect("send got None return")
}
#[inline]
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);
}
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]
fn is_done(&self) -> bool {
self.is_started() && (self.context._ref & 0x3) != 0
}
fn stack_usage(&self) -> (usize, usize) {
(self.stack.size(), self.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() {
trace!("generator is not done while drop");
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);
}
}
}
fn gen_init(_: usize, f: *mut usize) -> ! {
let clo = move || {
let f: &mut Option<Func> = unsafe { &mut *(f as *mut _) };
let func = f.take().unwrap();
func.call_once();
};
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");
}