#![feature(asm)]
#![feature(naked_functions)]
use std::ptr;
pub const DEFAULT_STACK_SIZE: usize = 1024 * 1024 * 2;
static mut RUNTIME: usize = 0;
pub use greenie_proc::greenify;
pub struct Runtime {
threads: Vec<Thread>,
current: usize,
}
#[derive(PartialEq, Eq, Debug)]
pub enum State {
Available,
Running,
Ready,
Suspended,
Sleep(std::time::Duration, std::time::Instant),
}
struct Thread {
id: usize,
stack: Vec<u8>,
ctx: ThreadContext,
state: State,
generator: Option<Rc<Generator>>,
}
#[derive(Debug, Default)]
#[repr(C)]
struct ThreadContext {
rsp: u64,
r15: u64,
r14: u64,
r13: u64,
r12: u64,
rbx: u64,
rbp: u64,
}
impl Thread {
fn new(id: usize, stack_size: usize) -> Self {
Thread {
id,
stack: vec![0_u8; stack_size],
ctx: ThreadContext::default(),
state: State::Available,
generator: None,
}
}
}
impl Runtime {
pub fn new(max_threads: usize, stack_size: Option<usize>) -> Self {
let base_thread = Thread {
id: 0,
stack: vec![0_u8; DEFAULT_STACK_SIZE],
ctx: ThreadContext::default(),
state: State::Running,
generator: None,
};
let mut threads = vec![base_thread];
let mut available_threads: Vec<Thread> = (1..max_threads)
.map(|i| {
Thread::new(
i,
match stack_size {
Some(size) => size,
_ => DEFAULT_STACK_SIZE,
},
)
})
.collect();
threads.append(&mut available_threads);
Runtime {
threads,
current: 0,
}
}
pub fn init(&self) {
unsafe {
let r_ptr: *const Runtime = self;
RUNTIME = r_ptr as usize;
}
}
pub fn run(&mut self) -> ! {
while self.t_yield() {}
std::process::exit(0);
}
fn t_return(&mut self) {
if self.current != 0 {
self.threads[self.current].state = State::Available;
self.threads[self.current].generator = None;
self.t_yield();
}
}
fn t_yield_generator<T: 'static>(&mut self, val: T) -> Result<(), &'static str> {
if self.threads[self.current].generator.is_none() {
return Err("Not a generator");
}
let to = if let Some(generator) = self.threads[self.current].generator.as_mut() {
if generator.is_join {
return Err("Not a generator");
}
let heap_value = Box::new(val);
generator.state.set(GeneratorState::Yielded(
heap_value as Box<dyn std::any::Any>,
));
generator.to
} else {
unreachable!()
};
self.threads[self.current].state = State::Suspended;
self.threads[to].state = State::Running;
let old_pos = self.current;
self.current = to;
unsafe {
switch(&mut self.threads[old_pos].ctx, &self.threads[to].ctx);
}
Ok(())
}
fn t_return_generator<T: 'static>(&mut self, val: T) -> Result<(), &'static str> {
let to = if let Some(generator) = self.threads[self.current].generator.as_mut() {
let heap_value = Box::new(val);
generator.state.set(GeneratorState::Complete(
heap_value as Box<dyn std::any::Any>,
));
generator.to
} else {
return Err("Can't return value from green thread");
};
self.threads[self.current].state = State::Available;
self.threads[self.current].generator = None;
self.threads[to].state = State::Running;
let old_pos = self.current;
self.current = to;
unsafe {
switch(&mut self.threads[old_pos].ctx, &self.threads[to].ctx);
}
Ok(())
}
pub fn yield_to(&mut self, id: usize) {
if self.threads[id].state == State::Suspended || self.threads[id].state == State::Ready {
self.threads[id].state = State::Running;
let old_pos = self.current;
self.current = id;
unsafe {
switch(&mut self.threads[old_pos].ctx, &self.threads[id].ctx);
}
}
}
fn t_yield(&mut self) -> bool {
let mut pos = self.current;
while self.threads[pos].state != State::Ready {
pos += 1;
if pos == self.threads.len() {
pos = 0;
}
if pos == self.current {
return false;
}
if let State::Sleep(dur, now) = &self.threads[pos].state {
if now.elapsed() >= *dur {
break;
}
}
}
if self.threads[self.current].state != State::Available {
self.threads[self.current].state = State::Ready;
}
self.threads[pos].state = State::Running;
let old_pos = self.current;
self.current = pos;
unsafe {
switch(&mut self.threads[old_pos].ctx, &self.threads[pos].ctx);
}
true
}
pub fn suspend(&mut self, id: usize) {
if id == self.current {
self.threads[id].state = State::Suspended;
self.t_yield();
} else {
self.threads[id].state = State::Suspended;
}
}
pub fn resume(&mut self, id: usize) {
if id == self.current {
} else {
self.threads[id].state = State::Ready;
}
}
pub fn spawn(&mut self, f: fn()) -> usize {
let (i, available) = self
.threads
.iter_mut()
.enumerate()
.find(|(_, t)| t.state == State::Available)
.expect("no available thread.");
let size = available.stack.len();
let s_ptr = available.stack.as_mut_ptr();
unsafe {
ptr::write(s_ptr.offset((size - 8) as isize) as *mut u64, guard as u64);
ptr::write(s_ptr.offset((size - 16) as isize) as *mut u64, f as u64);
available.ctx.rsp = s_ptr.offset((size - 16) as isize) as u64;
}
available.state = State::Ready;
i
}
}
#[cfg_attr(any(target_os = "windows", target_os = "linux"), naked)]
fn guard() {
unsafe {
let rt_ptr = RUNTIME as *mut Runtime;
let rt = &mut *rt_ptr;
rt.t_return();
};
}
pub fn yield_thread() {
unsafe {
let rt_ptr = RUNTIME as *mut Runtime;
(*rt_ptr).t_yield();
};
}
pub fn thread_sleep(dur: std::time::Duration) {
let rt = get_rt();
rt.threads[rt.current].state = State::Sleep(dur, std::time::Instant::now());
yield_thread();
}
pub fn spawn_greenie(f: fn()) {
unsafe {
let rt_ptr = RUNTIME as *mut Runtime;
let rt = &mut *rt_ptr;
let _ = rt.spawn(f);
yield_thread();
}
}
#[naked]
#[inline(never)]
unsafe fn switch(old: *mut ThreadContext, new: *const ThreadContext) {
asm!("
mov %rsp, 0x00($0)
mov %r15, 0x08($0)
mov %r14, 0x10($0)
mov %r13, 0x18($0)
mov %r12, 0x20($0)
mov %rbx, 0x28($0)
mov %rbp, 0x30($0)
mov 0x00($1), %rsp
mov 0x08($1), %r15
mov 0x10($1), %r14
mov 0x18($1), %r13
mov 0x20($1), %r12
mov 0x28($1), %rbx
mov 0x30($1), %rbp
ret
"
: "=*m"(old)
: "r"(new)
:
: "alignstack"
);
}
fn get_rt() -> &'static mut Runtime {
unsafe {
let rt_ptr = RUNTIME as *mut Runtime;
&mut *rt_ptr
}
}
pub enum GeneratorState {
Ready,
Yielded(Box<dyn std::any::Any>),
Complete(Box<dyn std::any::Any>),
}
impl Default for GeneratorState {
fn default() -> Self {
Self::Ready
}
}
pub struct Generator {
state: std::cell::Cell<GeneratorState>,
complete: std::cell::Cell<bool>,
thread_id: usize,
to: usize,
is_join: bool,
}
use std::rc::Rc;
pub fn generator_yield<T: 'static>(val: T) -> Result<(), &'static str> {
let rt = get_rt();
rt.t_yield_generator(val)
}
pub fn generator_return<T: 'static>(val: T) -> Result<(), &'static str> {
let rt = get_rt();
rt.t_return_generator(val)
}
impl Generator {
pub fn spawn(closure: fn()) -> Rc<Self> {
let rt = get_rt();
let to = rt.current;
let thread_id = rt.spawn(closure);
rt.threads[thread_id].state = State::Suspended;
let generator = Rc::new(Generator {
state: std::cell::Cell::new(GeneratorState::Ready),
thread_id,
to,
complete: std::cell::Cell::new(false),
is_join: false,
});
rt.threads[thread_id].generator = Some(generator.clone());
generator
}
fn spawn_joinable(id: usize) -> Rc<Self> {
let rt = get_rt();
let to = rt.current;
let thread_id = id;
rt.threads[thread_id].state = State::Suspended;
let generator = Rc::new(Generator {
state: std::cell::Cell::new(GeneratorState::Ready),
thread_id,
to,
complete: std::cell::Cell::new(false),
is_join: true,
});
rt.threads[thread_id].generator = Some(generator.clone());
generator
}
pub fn resume(&self) -> Result<GeneratorState, &'static str> {
if self.complete.get() {
return Err("Generator already complete");
}
let rt = get_rt();
rt.threads[rt.current].state = State::Suspended;
rt.threads[self.thread_id].state = State::Running;
let old = rt.current;
rt.current = self.thread_id;
assert!(rt.threads[self.thread_id].generator.is_some());
unsafe {
switch(&mut rt.threads[old].ctx, &rt.threads[self.thread_id].ctx);
}
let state = self.state.take();
if let GeneratorState::Complete(_) = &state {
self.complete.set(true);
}
Ok(state)
}
}
#[macro_export]
macro_rules! iterate_generator {
(for ($x: ident in $generator: expr) $b: block) => {
loop {
match $generator.resume() {
Ok(GeneratorState::Yielded($x)) => $b,
Ok(GeneratorState::Complete($x)) => {
break $x
},
Err(e) => {
eprintln!("{}",e);
std::process::exit(1);
}
_ => unreachable!()
}
}
};
(for ($x: ident in $generator: expr) $b: block and $b2: block) => {
loop {
match $generator.resume() {
Ok(GeneratorState::Yielded($x)) => $b,
Ok(GeneratorState::Complete($x)) => {
$b2
break;
},
Ok(GeneratorState::Ready) => panic!("impossible"),
Err(e) => {
eprintln!("{}",e);
std::process::exit(1);
}
_ => unreachable!()
}
}
};
(enumerate for ($c: ident, $x: ident in $generator: expr) $b: block) => {
let mut $c = 0;
loop {
match $generator.resume() {
Ok(GeneratorState::Yielded($x)) => $b,
Ok(GeneratorState::Complete($x)) => {
break $x
},
Err(e) => {
eprintln!("{}",e);
std::process::exit(1);
}
_ => unreachable!()
}
$c += 1;
}
};
(enumerate for ($c: ident, $x: ident in $generator: expr) $b: block and $b2: block) => {
let mut $c = 0;
loop {
match $generator.resume() {
Ok(GeneratorState::Yielded($x)) => $b,
Ok(GeneratorState::Complete($x)) => {
$b2
break;
},
Ok(GeneratorState::Ready) => panic!("impossible"),
Err(e) => {
eprintln!("{}",e);
std::process::exit(1);
}
_ => unreachable!()
}
$c += 1;
}
};
}