use crate::{Frame, VMResult, Value, VM};
use std::collections::{HashSet, VecDeque};
#[derive(Debug, Clone)]
pub enum SchedReq {
Go { name_idx: u16, args: Vec<Value> },
Make { cap: usize },
Send { ch: i64, val: Value },
Recv { ch: i64 },
RecvOk { ch: i64 },
Close { ch: i64 },
Select {
cases: Vec<SelectCase>,
has_default: bool,
},
}
#[derive(Debug, Clone)]
pub struct SelectCase {
pub recv: bool,
pub ch: i64,
pub val: Value,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SchedError {
Deadlock,
Panic(String),
}
impl std::fmt::Display for SchedError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SchedError::Deadlock => {
write!(f, "fatal error: all goroutines are asleep - deadlock!")
}
SchedError::Panic(m) => write!(f, "{m}"),
}
}
}
struct Channel {
cap: usize,
buf: VecDeque<Value>,
closed: bool,
send_q: VecDeque<(usize, Value)>,
recv_q: VecDeque<(usize, bool)>,
}
impl Channel {
fn new(cap: usize) -> Self {
Channel {
cap,
buf: VecDeque::new(),
closed: false,
send_q: VecDeque::new(),
recv_q: VecDeque::new(),
}
}
}
pub struct Scheduler<F: FnMut() -> VM> {
make_vm: F,
vms: Vec<VM>,
ready: VecDeque<usize>,
blocked: HashSet<usize>,
done: HashSet<usize>,
chans: Vec<Channel>,
select_waiters: Vec<(usize, Vec<SelectCase>)>,
recv_zero: Value,
}
impl<F: FnMut() -> VM> Scheduler<F> {
pub fn new(make_vm: F) -> Self {
Scheduler {
make_vm,
vms: Vec::new(),
ready: VecDeque::new(),
blocked: HashSet::new(),
done: HashSet::new(),
chans: Vec::new(),
select_waiters: Vec::new(),
recv_zero: Value::Int(0),
}
}
pub fn with_recv_zero(mut self, zero: Value) -> Self {
self.recv_zero = zero;
self
}
pub fn run(mut self, main_vm: VM) -> Result<(), SchedError> {
self.vms.push(main_vm);
self.ready.push_back(0);
self.drive()
}
pub fn run_capturing(mut self, main_vm: VM, global: &str) -> Result<Option<Value>, SchedError> {
self.vms.push(main_vm);
self.ready.push_back(0);
self.drive()?;
let main = &self.vms[0];
let val = main
.chunk
.names
.iter()
.position(|n| n == global)
.and_then(|i| main.globals.get(i).cloned());
Ok(val)
}
fn drive(&mut self) -> Result<(), SchedError> {
while let Some(gid) = self.ready.pop_front() {
self.vms[gid].clear_halt();
let popped = match self.vms[gid].run() {
VMResult::Error(e) => return Err(SchedError::Panic(e)),
VMResult::Ok(v) => Some(v),
VMResult::Halted => None,
};
match self.vms[gid].take_sched() {
Some(req) => {
if let Some(v) = popped {
self.vms[gid].stack.push(v);
}
self.service(gid, req)?;
}
None => {
self.done.insert(gid);
if gid == 0 {
return Ok(());
}
}
}
if self.ready.is_empty() && !self.blocked.is_empty() {
return Err(SchedError::Deadlock);
}
}
Ok(())
}
fn service(&mut self, gid: usize, req: SchedReq) -> Result<(), SchedError> {
match req {
SchedReq::Make { cap } => {
let id = self.chans.len();
self.chans.push(Channel::new(cap));
self.vms[gid].stack.push(Value::Int(id as i64));
self.ready.push_back(gid);
}
SchedReq::Go { name_idx, args } => {
let ngid = self.vms.len();
let mut vm = (self.make_vm)();
self.position_goroutine(&mut vm, name_idx, &args)?;
self.vms.push(vm);
self.ready.push_back(ngid);
self.ready.push_back(gid);
}
SchedReq::Send { ch, val } => {
if self.try_send(ch, &val)? {
self.ready.push_back(gid);
self.recheck_selects()?;
} else {
self.chan_mut(ch)?.send_q.push_back((gid, val));
self.blocked.insert(gid);
}
}
SchedReq::Recv { ch } => {
if let Some((v, _ok)) = self.try_recv(ch)? {
self.vms[gid].stack.push(v);
self.ready.push_back(gid);
self.recheck_selects()?;
} else {
self.chan_mut(ch)?.recv_q.push_back((gid, false));
self.blocked.insert(gid);
}
}
SchedReq::RecvOk { ch } => {
if let Some((v, ok)) = self.try_recv(ch)? {
self.vms[gid].stack.push(v);
self.vms[gid].stack.push(Value::Int(ok as i64));
self.ready.push_back(gid);
self.recheck_selects()?;
} else {
self.chan_mut(ch)?.recv_q.push_back((gid, true));
self.blocked.insert(gid);
}
}
SchedReq::Select { cases, has_default } => {
if let Some((idx, rv)) = self.try_select(&cases)? {
self.deliver_select(gid, idx, rv);
self.recheck_selects()?;
} else if has_default {
self.deliver_select(gid, cases.len(), Value::Undef);
} else {
self.select_waiters.push((gid, cases));
self.blocked.insert(gid);
}
}
SchedReq::Close { ch } => {
let c = self.chan_mut(ch)?;
c.closed = true;
let woken: Vec<(usize, bool)> = c.recv_q.drain(..).collect();
for (r, wants_ok) in woken {
self.vms[r].stack.push(self.recv_zero.clone());
if wants_ok {
self.vms[r].stack.push(Value::Int(0));
}
self.wake(r);
}
self.ready.push_back(gid);
self.recheck_selects()?;
}
}
Ok(())
}
fn deliver_select(&mut self, gid: usize, idx: usize, rv: Value) {
self.vms[gid].stack.push(rv);
self.vms[gid].stack.push(Value::Int(idx as i64));
self.wake(gid);
}
fn try_recv(&mut self, ch: i64) -> Result<Option<(Value, bool)>, SchedError> {
enum O {
V(Value),
Wake(usize, Value),
Zero,
No,
}
let o = {
let c = self.chan_mut(ch)?;
if let Some(v) = c.buf.pop_front() {
if let Some((s, sv)) = c.send_q.pop_front() {
c.buf.push_back(sv);
O::Wake(s, v)
} else {
O::V(v)
}
} else if let Some((s, sv)) = c.send_q.pop_front() {
O::Wake(s, sv)
} else if c.closed {
O::Zero
} else {
O::No
}
};
Ok(match o {
O::V(v) => Some((v, true)),
O::Wake(s, v) => {
self.wake(s);
Some((v, true))
}
O::Zero => Some((self.recv_zero.clone(), false)),
O::No => None,
})
}
fn try_send(&mut self, ch: i64, val: &Value) -> Result<bool, SchedError> {
enum O {
Deliver(usize, bool),
Buffered,
No,
}
let o = {
let c = self.chan_mut(ch)?;
if c.closed {
return Err(SchedError::Panic(
"panic: send on closed channel".to_string(),
));
}
if let Some((r, wants_ok)) = c.recv_q.pop_front() {
O::Deliver(r, wants_ok)
} else if c.buf.len() < c.cap {
c.buf.push_back(val.clone());
O::Buffered
} else {
O::No
}
};
Ok(match o {
O::Deliver(r, wants_ok) => {
self.vms[r].stack.push(val.clone());
if wants_ok {
self.vms[r].stack.push(Value::Int(1));
}
self.wake(r);
true
}
O::Buffered => true,
O::No => false,
})
}
fn try_select(&mut self, cases: &[SelectCase]) -> Result<Option<(usize, Value)>, SchedError> {
for (i, c) in cases.iter().enumerate() {
if c.recv {
if let Some((v, _ok)) = self.try_recv(c.ch)? {
return Ok(Some((i, v)));
}
} else if self.try_send(c.ch, &c.val)? {
return Ok(Some((i, Value::Undef)));
}
}
Ok(None)
}
fn recheck_selects(&mut self) -> Result<(), SchedError> {
loop {
let mut progressed = false;
let waiters = std::mem::take(&mut self.select_waiters);
let mut still = Vec::new();
for (gid, cases) in waiters {
match self.try_select(&cases)? {
Some((idx, rv)) => {
self.deliver_select(gid, idx, rv);
progressed = true;
}
None => still.push((gid, cases)),
}
}
self.select_waiters = still;
if !progressed {
break;
}
}
Ok(())
}
fn wake(&mut self, gid: usize) {
self.blocked.remove(&gid);
self.ready.push_back(gid);
}
fn chan_mut(&mut self, ch: i64) -> Result<&mut Channel, SchedError> {
self.chans
.get_mut(ch as usize)
.ok_or_else(|| SchedError::Panic(format!("panic: invalid channel {ch}")))
}
fn position_goroutine(
&self,
vm: &mut VM,
name_idx: u16,
args: &[Value],
) -> Result<(), SchedError> {
let entry = vm.chunk.find_sub(name_idx).ok_or_else(|| {
SchedError::Panic(format!("panic: goroutine target sub {name_idx} not found"))
})?;
let base = vm.stack.len();
for a in args {
vm.stack.push(a.clone());
}
vm.frames.push(Frame {
return_ip: vm.chunk.ops.len(), stack_base: base,
slots: Vec::new(),
entry_ip: Some(entry),
});
vm.ip = entry;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Chunk, ChunkBuilder, Op};
fn run(chunk: Chunk, global: &str) -> Result<Option<Value>, SchedError> {
let fc = chunk.clone();
Scheduler::new(move || VM::new(fc.clone())).run_capturing(VM::new(chunk), global)
}
#[test]
fn producer_consumer_over_unbuffered_channel() {
let mut b = ChunkBuilder::new();
let prod = b.add_name("producer");
let ch = b.add_name("ch");
let sum = b.add_name("sum");
b.emit(Op::LoadInt(0), 1);
b.emit(Op::ChanMake, 1);
b.emit(Op::SetVar(ch), 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::Go(prod, 1), 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecv, 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecv, 1);
b.emit(Op::Add, 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecv, 1);
b.emit(Op::Add, 1);
b.emit(Op::SetVar(sum), 1);
let skip = b.emit(Op::Jump(0), 1);
let entry = b.current_pos();
b.add_sub_entry(prod, entry);
b.emit(Op::SetSlot(0), 2); for v in 1..=3 {
b.emit(Op::GetSlot(0), 2);
b.emit(Op::LoadInt(v), 2);
b.emit(Op::ChanSend, 2);
}
b.emit(Op::LoadUndef, 2);
b.emit(Op::ReturnValue, 2);
b.patch_jump(skip, b.current_pos());
assert_eq!(run(b.build(), "sum"), Ok(Some(Value::Int(6))));
}
#[test]
fn buffered_channel_holds_values_without_a_receiver() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let sum = b.add_name("sum");
b.emit(Op::LoadInt(2), 1);
b.emit(Op::ChanMake, 1);
b.emit(Op::SetVar(ch), 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::LoadInt(10), 1);
b.emit(Op::ChanSend, 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::LoadInt(20), 1);
b.emit(Op::ChanSend, 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecv, 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecv, 1);
b.emit(Op::Add, 1);
b.emit(Op::SetVar(sum), 1);
assert_eq!(run(b.build(), "sum"), Ok(Some(Value::Int(30))));
}
#[test]
fn select_picks_the_ready_case() {
let mut b = ChunkBuilder::new();
let ch1 = b.add_name("ch1");
let ch2 = b.add_name("ch2");
let sum = b.add_name("sum");
b.emit(Op::LoadInt(0), 1);
b.emit(Op::ChanMake, 1);
b.emit(Op::SetVar(ch1), 1);
b.emit(Op::LoadInt(1), 1);
b.emit(Op::ChanMake, 1);
b.emit(Op::SetVar(ch2), 1);
b.emit(Op::GetVar(ch2), 1);
b.emit(Op::LoadInt(7), 1);
b.emit(Op::ChanSend, 1);
b.emit(Op::GetVar(ch1), 1);
b.emit(Op::LoadInt(1), 1);
b.emit(Op::LoadInt(0), 1);
b.emit(Op::GetVar(ch2), 1);
b.emit(Op::LoadInt(1), 1);
b.emit(Op::LoadInt(0), 1);
b.emit(Op::Select(2, 0), 1);
b.emit(Op::Pop, 1);
b.emit(Op::SetVar(sum), 1);
assert_eq!(run(b.build(), "sum"), Ok(Some(Value::Int(7))));
}
#[test]
fn select_default_when_no_case_ready() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let idx = b.add_name("idx");
b.emit(Op::LoadInt(0), 1);
b.emit(Op::ChanMake, 1);
b.emit(Op::SetVar(ch), 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::LoadInt(1), 1);
b.emit(Op::LoadInt(0), 1);
b.emit(Op::Select(1, 1), 1); b.emit(Op::SetVar(idx), 1); b.emit(Op::Pop, 1); assert_eq!(run(b.build(), "idx"), Ok(Some(Value::Int(1))));
}
#[test]
fn repro_closed_drained_recv_is_indistinguishable_from_a_received_zero() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let out = b.add_name("out");
b.emit(Op::LoadInt(2), 1);
b.emit(Op::ChanMake, 1);
b.emit(Op::SetVar(ch), 1);
for v in [10, 20] {
b.emit(Op::GetVar(ch), 1);
b.emit(Op::LoadInt(v), 1);
b.emit(Op::ChanSend, 1);
}
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanClose, 1);
for _ in 0..2 {
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecv, 1);
b.emit(Op::Pop, 1);
}
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecv, 1);
b.emit(Op::SetVar(out), 1);
let drained = run(b.build(), "out");
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let out = b.add_name("out");
b.emit(Op::LoadInt(1), 1);
b.emit(Op::ChanMake, 1);
b.emit(Op::SetVar(ch), 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::LoadInt(0), 1);
b.emit(Op::ChanSend, 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecv, 1);
b.emit(Op::SetVar(out), 1);
let real_zero = run(b.build(), "out");
assert_eq!(drained, Ok(Some(Value::Int(0))));
assert_eq!(real_zero, Ok(Some(Value::Int(0))));
assert_eq!(
drained, real_zero,
"Op::ChanRecv cannot report closed-and-drained; use Op::ChanRecvOk"
);
}
fn make_chan(b: &mut ChunkBuilder, ch: u16, cap: i64) {
b.emit(Op::LoadInt(cap), 1);
b.emit(Op::ChanMake, 1);
b.emit(Op::SetVar(ch), 1);
}
fn recv_ok(b: &mut ChunkBuilder, ch: u16, v: u16, ok: u16) {
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecvOk, 1);
b.emit(Op::SetVar(ok), 1);
b.emit(Op::SetVar(v), 1);
}
#[test]
fn recv_ok_reports_true_for_a_buffered_value_and_false_once_drained() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let (v0, k0) = (b.add_name("v0"), b.add_name("k0"));
let (v1, k1) = (b.add_name("v1"), b.add_name("k1"));
let (v2, k2) = (b.add_name("v2"), b.add_name("k2"));
make_chan(&mut b, ch, 2);
for v in [10, 20] {
b.emit(Op::GetVar(ch), 1);
b.emit(Op::LoadInt(v), 1);
b.emit(Op::ChanSend, 1);
}
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanClose, 1);
recv_ok(&mut b, ch, v0, k0);
recv_ok(&mut b, ch, v1, k1);
recv_ok(&mut b, ch, v2, k2);
let chunk = b.build();
for (name, want) in [
("v0", 10),
("k0", 1),
("v1", 20),
("k1", 1),
("v2", 0),
("k2", 0),
] {
assert_eq!(
run(chunk.clone(), name),
Ok(Some(Value::Int(want))),
"{name}"
);
}
}
#[test]
fn recv_ok_on_a_channel_closed_with_no_values_is_false() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let (v, k) = (b.add_name("v"), b.add_name("k"));
make_chan(&mut b, ch, 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanClose, 1);
recv_ok(&mut b, ch, v, k);
let chunk = b.build();
assert_eq!(run(chunk.clone(), "v"), Ok(Some(Value::Int(0))));
assert_eq!(run(chunk, "k"), Ok(Some(Value::Int(0))));
}
#[test]
fn recv_ok_of_a_real_zero_from_an_open_channel_is_true() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let (v, k) = (b.add_name("v"), b.add_name("k"));
make_chan(&mut b, ch, 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::LoadInt(0), 1);
b.emit(Op::ChanSend, 1);
recv_ok(&mut b, ch, v, k);
let chunk = b.build();
assert_eq!(run(chunk.clone(), "v"), Ok(Some(Value::Int(0))));
assert_eq!(run(chunk, "k"), Ok(Some(Value::Int(1))));
}
#[test]
fn recv_ok_on_an_open_empty_channel_blocks() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let (v, k) = (b.add_name("v"), b.add_name("k"));
make_chan(&mut b, ch, 1);
recv_ok(&mut b, ch, v, k);
assert_eq!(run(b.build(), "v"), Err(SchedError::Deadlock));
}
#[test]
fn a_blocked_recv_ok_is_woken_true_by_a_sender_and_false_by_a_close() {
let mut b = ChunkBuilder::new();
let sender = b.add_name("sender");
let ch = b.add_name("ch");
let (v0, k0) = (b.add_name("v0"), b.add_name("k0"));
let (v1, k1) = (b.add_name("v1"), b.add_name("k1"));
make_chan(&mut b, ch, 0);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::Go(sender, 1), 1);
recv_ok(&mut b, ch, v0, k0);
recv_ok(&mut b, ch, v1, k1);
let skip = b.emit(Op::Jump(0), 1);
let entry = b.current_pos();
b.add_sub_entry(sender, entry);
b.emit(Op::SetSlot(0), 2);
b.emit(Op::GetSlot(0), 2);
b.emit(Op::LoadInt(5), 2);
b.emit(Op::ChanSend, 2);
b.emit(Op::GetSlot(0), 2);
b.emit(Op::ChanClose, 2);
b.emit(Op::LoadUndef, 2);
b.emit(Op::ReturnValue, 2);
b.patch_jump(skip, b.current_pos());
let chunk = b.build();
for (name, want) in [("v0", 5), ("k0", 1), ("v1", 0), ("k1", 0)] {
assert_eq!(
run(chunk.clone(), name),
Ok(Some(Value::Int(want))),
"{name}"
);
}
}
#[test]
fn range_over_a_closed_channel_yields_every_value_then_stops() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let (v, k) = (b.add_name("v"), b.add_name("k"));
let sum = b.add_name("sum");
let n = b.add_name("n");
make_chan(&mut b, ch, 3);
for x in 1..=3 {
b.emit(Op::GetVar(ch), 1);
b.emit(Op::LoadInt(x), 1);
b.emit(Op::ChanSend, 1);
}
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanClose, 1);
b.emit(Op::LoadInt(0), 1);
b.emit(Op::SetVar(sum), 1);
b.emit(Op::LoadInt(0), 1);
b.emit(Op::SetVar(n), 1);
let top = b.current_pos();
recv_ok(&mut b, ch, v, k);
b.emit(Op::GetVar(k), 1);
let exit = b.emit(Op::JumpIfFalse(0), 1);
b.emit(Op::GetVar(sum), 1);
b.emit(Op::GetVar(v), 1);
b.emit(Op::Add, 1);
b.emit(Op::SetVar(sum), 1);
b.emit(Op::GetVar(n), 1);
b.emit(Op::LoadInt(1), 1);
b.emit(Op::Add, 1);
b.emit(Op::SetVar(n), 1);
let back = b.emit(Op::Jump(0), 1);
b.patch_jump(back, top);
b.patch_jump(exit, b.current_pos());
let chunk = b.build();
assert_eq!(run(chunk.clone(), "n"), Ok(Some(Value::Int(3))));
assert_eq!(run(chunk, "sum"), Ok(Some(Value::Int(6))));
}
#[test]
fn recv_ok_leaves_plain_recv_alone_on_the_same_channel() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let (v, k) = (b.add_name("v"), b.add_name("k"));
let plain = b.add_name("plain");
make_chan(&mut b, ch, 2);
for x in [7, 9] {
b.emit(Op::GetVar(ch), 1);
b.emit(Op::LoadInt(x), 1);
b.emit(Op::ChanSend, 1);
}
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecv, 1);
b.emit(Op::SetVar(plain), 1);
recv_ok(&mut b, ch, v, k);
let chunk = b.build();
assert_eq!(run(chunk.clone(), "plain"), Ok(Some(Value::Int(7))));
assert_eq!(run(chunk.clone(), "v"), Ok(Some(Value::Int(9))));
assert_eq!(run(chunk, "k"), Ok(Some(Value::Int(1))));
}
#[test]
fn recv_ok_uses_the_frontends_recv_zero_when_drained() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
let (v, k) = (b.add_name("v"), b.add_name("k"));
make_chan(&mut b, ch, 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanClose, 1);
recv_ok(&mut b, ch, v, k);
let chunk = b.build();
let fc = chunk.clone();
let got = Scheduler::new(move || VM::new(fc.clone()))
.with_recv_zero(Value::str(""))
.run_capturing(VM::new(chunk.clone()), "v");
assert_eq!(got, Ok(Some(Value::str(""))));
let fc = chunk.clone();
let got = Scheduler::new(move || VM::new(fc.clone()))
.with_recv_zero(Value::str(""))
.run_capturing(VM::new(chunk), "k");
assert_eq!(got, Ok(Some(Value::Int(0))));
}
#[test]
fn deadlock_is_detected() {
let mut b = ChunkBuilder::new();
let ch = b.add_name("ch");
b.emit(Op::LoadInt(0), 1);
b.emit(Op::ChanMake, 1);
b.emit(Op::SetVar(ch), 1);
b.emit(Op::GetVar(ch), 1);
b.emit(Op::ChanRecv, 1);
b.emit(Op::Pop, 1);
assert_eq!(run(b.build(), "x"), Err(SchedError::Deadlock));
}
}