narju/

sched.rs

use std::any::Any;
use std::collections::VecDeque;
use std::fmt;

use crate::floor::host::{Atom, HostType, Registry};
use crate::floor::machine::{Fault, Machine, Mail, Step, Store};
use crate::floor::{rc_val, Val};

pub type Id = usize;

/// Repeated from [`crate::world`] because an address names one and is below it.
pub type ConnId = u64;

fn sym(s: &str) -> Val {
    Val::Sym(s.into())
}

fn pair(a: Val, b: Val) -> Val {
    Val::Pair(rc_val(a), rc_val(b))
}

/// Where a task is: in this heap, or on the far side of one of this node's
/// links. A host peer is not here - `('host name)` is an ordinary cons, a name
/// the program writes, while a task address is a capability the runtime mints.
///
/// Opaque because crossing a link `Task(3)` must arrive as `Remote(conn, 3)`
/// ([`crate::link::across`]), and a tagged cons is a shape a program can build
/// as data, so the rewrite could not tell an address from something like one.
#[derive(Debug, Clone, PartialEq)]
pub enum Addr {
    Task(Id),
    Remote(ConnId, Id),
}

impl Addr {
    pub const NAME: &'static str = "addr";

    pub fn val(self) -> Val {
        Val::Atom(Atom::new(self))
    }

    pub fn of(v: &Val) -> Option<&Addr> {
        match v {
            Val::Atom(a) => a.downcast::<Addr>(),
            _ => None,
        }
    }

    /// The conn is written even though the receiver discards it: what a link
    /// does with an address is the link's business, not the codec's.
    fn decode(bytes: &[u8]) -> Option<Atom> {
        fn be(b: &[u8]) -> Option<u64> {
            b.try_into().ok().map(u64::from_be_bytes)
        }
        fn id(b: &[u8]) -> Option<Id> {
            be(b).and_then(|n| Id::try_from(n).ok())
        }
        match bytes.split_first()? {
            (0, rest) => Some(Atom::new(Addr::Task(id(rest)?))),
            (1, rest) if rest.len() == 16 => {
                Some(Atom::new(Addr::Remote(be(&rest[..8])?, id(&rest[8..])?)))
            }
            _ => None,
        }
    }
}

impl fmt::Display for Addr {
    /// Not a form the reader accepts, deliberately: an address cannot be
    /// written down, so its printed form does not pretend to be source.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Addr::Task(id) => write!(f, "#<task {id}>"),
            Addr::Remote(conn, id) => write!(f, "#<task {id} on {conn}>"),
        }
    }
}

impl HostType for Addr {
    fn type_name(&self) -> &'static str {
        Addr::NAME
    }

    fn eq_atom(&self, other: &dyn HostType) -> bool {
        other.as_any().downcast_ref::<Addr>() == Some(self)
    }

    fn encode(&self, out: &mut Vec<u8>) {
        match self {
            Addr::Task(id) => {
                out.push(0);
                out.extend_from_slice(&(*id as u64).to_be_bytes());
            }
            Addr::Remote(conn, id) => {
                out.push(1);
                out.extend_from_slice(&conn.to_be_bytes());
                out.extend_from_slice(&(*id as u64).to_be_bytes());
            }
        }
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

pub fn extend(r: Registry) -> Registry {
    r.ty(Addr::NAME, Addr::decode)
}

pub fn addr(id: Id) -> Val {
    Addr::Task(id).val()
}

enum Route {
    Local(Id),
    /// Not for anything here. The heap does not interpret it; that is the
    /// executor's job, and why a host capability needs no floor verb of its own.
    Away,
    Bad,
}

fn route(v: &Val) -> Route {
    match Addr::of(v) {
        Some(Addr::Task(id)) => Route::Local(*id),
        Some(Addr::Remote(_, _)) => Route::Away,
        // A host peer is `('host name)`, and any other tagged cons is some
        // executor's business too - the heap does not know what exists above it.
        None => match v {
            Val::Pair(tag, _) if matches!(&**tag, Val::Sym(_)) => Route::Away,
            _ => Route::Bad,
        },
    }
}

/// `('task-down addr . result)`, where result is `(ok . v)`, `(throw . v)`,
/// `(throw . noproc)` for a task already gone when the monitor was set, or
/// `(throw . noconnection)` for one behind a link that died.
pub fn down(to: Val, result: Val) -> Val {
    pair(sym("task-down"), pair(to, result))
}

fn task_down(id: Id, result: Val) -> Val {
    down(addr(id), result)
}

/// Something a running task asked for that only its heap can do.
enum Request {
    Send {
        to: Id,
        msg: Val,
    },
    Away {
        to: Val,
        msg: Val,
    },
    Spawn {
        id: Id,
        thunk: Val,
    },
    Watch {
        watcher: Id,
        target: Id,
    },
    /// A watch on what is not in this heap, passed out for the executor to keep.
    /// Undecoded, because the heap does not know what kinds of address there are.
    WatchAway {
        watcher: Id,
        target: Val,
    },
    Limit {
        target: Id,
        turns: usize,
    },
}

/// Nothing here reaches the heap, because while a task runs the heap is its
/// caller rather than something it can borrow. Sends and spawns accumulate as
/// requests the heap drains when the turn ends.
pub struct Post {
    me: Id,
    inbox: VecDeque<Val>,
    out: Vec<Request>,
    reached: bool,
    /// The id the heap will hand the next task it creates. Lent for the turn so
    /// `spawn` can answer at once; sound only because one task runs at a time.
    next: Id,
}

impl Mail for Post {
    fn spawn(&mut self, thunk: Val) -> Result<Val, Val> {
        let id = self.next;
        self.next += 1;
        self.out.push(Request::Spawn { id, thunk });
        Ok(addr(id))
    }

    fn send(&mut self, to: &Val, msg: Val) -> Result<Val, Val> {
        match route(to) {
            // A send to oneself lands now. Deferred, it would let a task
            // park on a mailbox holding a message it has itself posted.
            Route::Local(id) if id == self.me => self.inbox.push_back(msg),
            Route::Local(id) => self.out.push(Request::Send { to: id, msg }),
            Route::Away => self.out.push(Request::Away {
                to: to.clone(),
                msg,
            }),
            Route::Bad => return Err(pair(sym("bad-address"), to.clone())),
        }
        Ok(Val::Nil)
    }

    /// Reaching here is what "responsive" means, whether or not there was
    /// anything to read: a task in a tight loop over a full mailbox is working
    /// rather than diverging, so blocking would be the wrong test.
    fn receive(&mut self) -> Option<Val> {
        self.reached = true;
        self.inbox.pop_front()
    }

    /// A watch goes wherever a message would, but what comes back is weaker off
    /// this heap. A dead link does not mean the far task died, only that nothing
    /// more can be learned, so it reports `noconnection` rather than an exit
    /// value: the caller must stop waiting either way, and must not be told a
    /// result nobody observed. A host peer that ends reports `noproc` and no
    /// value, an adapter having nothing it could be said to have evaluated.
    fn monitor(&mut self, target: &Val) -> Result<Val, Val> {
        match route(target) {
            Route::Local(target) => {
                self.out.push(Request::Watch {
                    watcher: self.me,
                    target,
                });
                Ok(Val::Nil)
            }
            Route::Away => {
                self.out.push(Request::WatchAway {
                    watcher: self.me,
                    target: target.clone(),
                });
                Ok(Val::Nil)
            }
            Route::Bad => Err(pair(sym("bad-address"), target.clone())),
        }
    }

    /// Counted in turns rather than steps because the step budget is a fairness
    /// knob and this is not: a task on an expensive semantics gets less done per
    /// turn, but turns-to-answer is what its caller can reason about.
    fn limit(&mut self, target: &Val, turns: i64) -> Result<Val, Val> {
        match route(target) {
            Route::Local(target) => {
                self.out.push(Request::Limit {
                    target,
                    turns: turns as usize,
                });
                Ok(Val::Nil)
            }
            Route::Away => Err(pair(sym("not-local"), target.clone())),
            Route::Bad => Err(pair(sym("bad-address"), target.clone())),
        }
    }
}

struct Task {
    machine: Machine,
    store: Store<Post>,
    /// Blocked on an empty mailbox, and so not in the run queue.
    parked: bool,
    watchers: Vec<Id>,
    /// Consecutive turns in which `receive` was not reached.
    turns: usize,
    /// `None` means a task may compute for as long as it likes.
    limit: Option<usize>,
}

impl Task {
    /// The thunk is applied to the task's own address, which is how a task comes
    /// to know it - so naming oneself costs no verb and no scheduler call.
    fn new(id: Id, thunk: Val) -> Task {
        Task {
            machine: Machine::applying(thunk, addr(id)),
            store: Store::new(Post {
                me: id,
                inbox: VecDeque::new(),
                out: Vec::new(),
                reached: false,
                next: 0,
            }),
            parked: false,
            watchers: Vec::new(),
            turns: 0,
            limit: None,
        }
    }
}

pub struct Heap {
    /// A slot goes to `None` when its task ends and is never reused, so an
    /// address outlives its task and a send to a dead one is a drop, not a
    /// misdelivery.
    tasks: Vec<Option<Task>>,
    runq: VecDeque<Id>,
    /// Swapped into whichever task is running, so a closure that captured a cell
    /// keeps meaning that cell after it is spawned.
    cells: Vec<Val>,
    /// Floor steps a task may take before it is rotated out.
    budget: usize,
    exits: Vec<(Id, Val)>,
    /// Messages for addresses that are not local, in send order.
    away: Vec<(Val, Val)>,
    /// Watches on what is not here, as `(watcher, address)`.
    watches: Vec<(Id, Val)>,
    live: usize,
}

impl Heap {
    pub fn new(budget: usize) -> Heap {
        Heap {
            tasks: Vec::new(),
            runq: VecDeque::new(),
            cells: Vec::new(),
            budget,
            exits: Vec::new(),
            away: Vec::new(),
            watches: Vec::new(),
            live: 0,
        }
    }

    /// Seed a task from outside. The thunk is applied to nil.
    pub fn spawn(&mut self, thunk: Val) -> Id {
        let id = self.tasks.len();
        self.tasks.push(Some(Task::new(id, thunk)));
        self.runq.push_back(id);
        self.live += 1;
        id
    }

    /// A heap with live tasks and an empty queue is waiting on a message from
    /// outside; a heap with none is finished.
    pub fn idle(&self) -> bool {
        self.live == 0
    }

    /// Waking it if it was parked. An address for anything else is dropped,
    /// exactly as a send to a task that has ended is.
    pub fn post(&mut self, to: &Val, msg: Val) {
        if let Route::Local(id) = route(to) {
            self.deliver(id, msg);
        }
    }

    pub fn alive(&self, id: Id) -> bool {
        matches!(self.tasks.get(id), Some(Some(_)))
    }

    /// Tasks that have ended since this was last called, in the order they ended.
    pub fn take_exits(&mut self) -> Vec<(Id, Val)> {
        std::mem::take(&mut self.exits)
    }

    /// Routing these is the executor's job: another heap, another machine and a
    /// host adapter are all just the far end of a channel, which is what lets a
    /// syscall be an ordinary `send` rather than a floor verb.
    pub fn take_away(&mut self) -> Vec<(Val, Val)> {
        std::mem::take(&mut self.away)
    }

    /// As `(watcher, address)`. Whether such a thing is alive is not a question
    /// this heap can put, so the answer is owed by whoever holds link or peer.
    pub fn take_watches(&mut self) -> Vec<(Id, Val)> {
        std::mem::take(&mut self.watches)
    }

    /// Until the queue empties - every task finished or parked on an empty mailbox.
    pub fn run(&mut self) -> Result<(), Fault> {
        while let Some(id) = self.runq.pop_front() {
            self.turn(id)?;
        }
        Ok(())
    }

    fn turn(&mut self, id: Id) -> Result<(), Fault> {
        let Some(mut task) = self.tasks[id].take() else {
            return Ok(());
        };

        std::mem::swap(&mut self.cells, &mut task.store.cells);
        task.store.mail.next = self.tasks.len();
        let outcome = task.machine.turn(&mut task.store, self.budget);
        std::mem::swap(&mut self.cells, &mut task.store.cells);
        let out = std::mem::take(&mut task.store.mail.out);

        let blocked = matches!(outcome, Ok((Step::Blocked, _)));
        let ended = match outcome {
            // A raise that reached the end of the task is what a failure is: it
            // propagated past everything that might have caught it.
            Ok((Step::Done(Val::Raise(v)), _)) => Some(pair(sym("throw"), Val::clone(&v))),
            Ok((Step::Done(v), _)) => Some(pair(sym("ok"), v)),
            // Not catchable, and not the task's fault: a machine invariant was
            // violated, so the heap stops rather than reports.
            Err(bug) => return Err(bug),
            Ok((Step::Running, _)) => None,
            Ok((Step::Blocked, _)) => None,
        };

        if std::mem::take(&mut task.store.mail.reached) {
            task.turns = 0;
        } else {
            task.turns += 1;
        }
        // Killed rather than asked to stop, because the case this exists for
        // cannot hear the asking: reporting a divergent tower needs evaluation.
        let ended = ended.or_else(|| {
            task.limit
                .is_some_and(|n| task.turns > n)
                .then(|| pair(sym("throw"), sym("unresponsive")))
        });

        match ended {
            Some(result) => {
                let watchers = std::mem::take(&mut task.watchers);
                // What the turn sent, before what its ending reports. The other
                // order lets a task's death overtake its own reply: the caller
                // sees `callee-down` for a call that was in fact answered.
                self.drain(out);
                for w in watchers {
                    self.deliver(w, task_down(id, result.clone()));
                }
                self.exits.push((id, result));
                self.live -= 1;
            }
            None => {
                self.tasks[id] = Some(task);
                self.drain(out);
                let task = self.tasks[id].as_mut().expect("just restored");
                if blocked && task.store.mail.inbox.is_empty() {
                    task.parked = true;
                } else {
                    self.runq.push_back(id);
                }
            }
        }
        Ok(())
    }

    fn drain(&mut self, out: Vec<Request>) {
        for req in out {
            match req {
                Request::Send { to, msg } => self.deliver(to, msg),
                Request::Away { to, msg } => self.away.push((to, msg)),
                Request::WatchAway { watcher, target } => self.watches.push((watcher, target)),
                Request::Spawn { id, thunk } => {
                    debug_assert_eq!(id, self.tasks.len(), "ids are handed out in order");
                    self.tasks.push(Some(Task::new(id, thunk)));
                    self.runq.push_back(id);
                    self.live += 1;
                }
                Request::Watch { watcher, target } => {
                    match self.tasks.get_mut(target).and_then(Option::as_mut) {
                        Some(task) => task.watchers.push(watcher),
                        // Answering now rather than never is what makes a
                        // monitor race-free to set.
                        None => {
                            let down = task_down(target, pair(sym("throw"), sym("noproc")));
                            self.deliver(watcher, down);
                        }
                    }
                }
                // Silent for a task that has ended, as a `send` to one is.
                Request::Limit { target, turns } => {
                    if let Some(Some(task)) = self.tasks.get_mut(target) {
                        task.limit = Some(turns);
                    }
                }
            }
        }
    }

    fn deliver(&mut self, to: Id, msg: Val) {
        let Some(Some(task)) = self.tasks.get_mut(to) else {
            return;
        };
        task.store.mail.inbox.push_back(msg);
        let wake = task.parked;
        task.parked = false;
        if wake {
            self.runq.push_back(to);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::floor::machine::NoMail;
    use crate::floor::Env;

    fn val(src: &str) -> Val {
        let Val::Code(e) = crate::surface::compile(src).expect("compiles") else {
            panic!("not code")
        };
        Machine::new(Env::default(), e)
            .run(&mut Store::<NoMail>::default())
            .expect("no fault")
    }

    fn ok(v: Val) -> Val {
        pair(sym("ok"), v)
    }

    fn run(h: &mut Heap) -> Vec<(Id, Val)> {
        h.run().expect("no bug");
        h.take_exits()
    }

    #[test]
    fn a_task_runs_to_completion() {
        let mut h = Heap::new(1000);
        let id = h.spawn(val("(lambda t (_) (+ 1 2))"));
        assert_eq!(run(&mut h), vec![(id, ok(Val::Num(3)))]);
        assert!(!h.alive(id));
    }

    #[test]
    fn a_throw_is_a_result_not_a_heap_failure() {
        let mut h = Heap::new(1000);
        let id = h.spawn(val("(lambda t (_) (throw 'nope))"));
        assert_eq!(run(&mut h), vec![(id, pair(sym("throw"), sym("nope")))]);
    }

    #[test]
    fn an_empty_mailbox_is_the_only_reason_to_park() {
        let mut h = Heap::new(1000);
        let id = h.spawn(val("(lambda t (_) (receive 0))"));
        // Quiescent with the task still alive: it is parked, not spinning.
        assert_eq!(run(&mut h), vec![]);
        assert!(h.alive(id));

        h.post(&addr(id), Val::Num(7));
        assert_eq!(run(&mut h), vec![(id, ok(Val::Num(7)))]);
    }

    #[test]
    fn a_task_is_told_its_own_address() {
        let mut h = Heap::new(1000);
        let id = h.spawn(val("(lambda t (me) me)"));
        assert_eq!(run(&mut h), vec![(id, ok(addr(id)))]);
    }

    #[test]
    fn a_send_to_oneself_is_visible_to_the_next_receive() {
        let mut h = Heap::new(1000);
        let id = h.spawn(val("(lambda t (me) (begin (send me 5) (receive)))"));
        assert_eq!(run(&mut h), vec![(id, ok(Val::Num(5)))]);
    }

    #[test]
    fn a_spawned_task_gets_the_next_address_and_runs() {
        let mut h = Heap::new(1000);
        let parent = h.spawn(val("(lambda t (me)
               (let ((child (spawn (lambda c (_) (send me 'hi)))))
                 (cons child (receive 0))))"));
        let exits = run(&mut h);
        assert_eq!(
            exits,
            vec![(1, ok(Val::Nil)), (parent, ok(pair(addr(1), sym("hi")))),]
        );
    }

    #[test]
    fn a_monitor_reports_a_death_it_did_not_see() {
        let mut h = Heap::new(1000);
        let watcher = h.spawn(val("(lambda t (_)
               (let ((child (spawn (lambda c (_) 9))))
                 (begin (monitor child) (receive 0))))"));
        let exits = run(&mut h);
        let down = task_down(1, ok(Val::Num(9)));
        assert_eq!(exits, vec![(1, ok(Val::Num(9))), (watcher, ok(down))]);
    }

    #[test]
    fn a_dying_task_is_heard_before_it_is_mourned() {
        let mut h = Heap::new(1000);
        let watcher = h.spawn(val("(lambda t (me)
               (let ((child (spawn (lambda c (_) (begin (send me 'answer) 9)))))
                 (begin (monitor child)
                        (let ((first (receive 0))) (cons first (receive 0))))))"));
        let heard = pair(sym("answer"), task_down(1, ok(Val::Num(9))));
        assert_eq!(
            run(&mut h),
            vec![(1, ok(Val::Num(9))), (watcher, ok(heard))]
        );
    }

    #[test]
    fn monitoring_a_dead_task_answers_at_once() {
        let mut h = Heap::new(1000);
        let id = h.spawn(val(
            "(lambda t (_) (begin (monitor (receive)) (receive 0)))",
        ));
        h.post(&addr(0), addr(99));
        let down = task_down(99, pair(sym("throw"), sym("noproc")));
        assert_eq!(run(&mut h), vec![(id, ok(down))]);
    }

    #[test]
    fn a_budget_rotates_tasks_rather_than_running_one_to_the_end() {
        // Two tasks that each take more than a turn's worth of steps. The
        // first to finish is not simply the first started unless one ran
        // to completion, so interleaving shows in the step count alone.
        let src = "(lambda t (_)
                     (let ((down (lambda d (n) (if (eq? n 0) 'done (d (- n 1))))))
                       (down 200)))";
        let mut h = Heap::new(8);
        let a = h.spawn(val(src));
        let b = h.spawn(val(src));
        let exits = run(&mut h);
        assert_eq!(exits.len(), 2);
        assert_eq!(exits[0].0, a);
        assert_eq!(exits[1].0, b);
        assert_eq!(exits[0].1, ok(sym("done")));
    }

    #[test]
    fn a_channel_carries_data_and_refuses_anything_else() {
        let mut h = Heap::new(1000);
        let cell = h.spawn(val("(lambda t (_) (send '(task 0) (cell-new 1)))"));
        assert_eq!(
            run(&mut h),
            vec![(
                cell,
                pair(sym("throw"), pair(sym("not-data"), Val::Cell(0)))
            )]
        );

        // Nested, because the check is structural rather than a tag test
        // on the payload's head.
        let mut h = Heap::new(1000);
        let clo = h.spawn(val(
            "(lambda t (_) (car (catch (send '(task 0) (cons 1 (lambda f (x) x))))))",
        ));
        assert_eq!(run(&mut h), vec![(clo, ok(sym("not-data")))]);
    }

    #[test]
    fn an_address_that_is_not_one_is_a_throw() {
        let mut h = Heap::new(1000);
        let id = h.spawn(val("(lambda t (_) (send 0 'hi))"));
        assert_eq!(
            run(&mut h),
            vec![(
                id,
                pair(sym("throw"), pair(sym("bad-address"), Val::Num(0)))
            )]
        );
    }

    #[test]
    fn a_task_across_a_link_leaves_the_heap() {
        let mut h = Heap::new(1000);
        let to = Addr::Remote(2, 5).val();
        h.spawn(val("(lambda t (_) (receive))"));
        h.post(&addr(0), to.clone());
        h.spawn(val("(lambda t (_) (send (receive) 'hi))"));
        h.post(&addr(1), to.clone());
        h.run().expect("no bug");
        assert_eq!(h.take_away(), vec![(to, sym("hi"))]);
    }

    #[test]
    fn an_address_the_heap_does_not_know_goes_out_rather_than_failing() {
        let mut h = Heap::new(1000);
        h.spawn(val("(lambda t (_) (send '(host nowhere) 'hi))"));
        h.run().expect("no bug");
        assert_eq!(h.take_away(), vec![(val("'(host nowhere)"), sym("hi"))]);
    }

    #[test]
    fn tasks_in_one_heap_share_cells() {
        let mut h = Heap::new(1000);
        let id = h.spawn(val("(lambda t (_)
               (let ((c (cell-new 0)))
                 (let ((child (spawn (lambda k (_) (cell-set! c 42)))))
                   (begin (monitor child) (begin (receive 0) (cell-read c))))))"));
        assert_eq!(
            run(&mut h),
            vec![(1, ok(Val::Num(42))), (id, ok(Val::Num(42)))]
        );
    }
}