sva-engine 0.7.2

// Concern: orders nodes dependencies-first and picks which a reading materializes | Non-concern: what a held node contains (render.rs), reading one (refs.rs) | IO: (Instances, asks) -> Schedule

use std::collections::{BTreeMap, BTreeSet, HashMap};

use sva_ast::{Arg, Expr};
use sva_formula::{Held, NodeId, Var};

use crate::cast::Cast;
use crate::error::EngineError;
use crate::instantiate::{Cx, Instances, Node};
use crate::query::Ask;
use crate::typing::{Typing, Value, When};

/// One-hop instance names. `self(...)` is a same-node read, never an edge.
fn direct_refs(inst: &Instances, e: &Expr, cx: Cx, out: &mut Vec<String>) {
    if inst
        .follow(e, cx, |e2, cx2| direct_refs(inst, e2, cx2, out))
        .is_some()
    {
        return;
    }
    match inst.node(e, cx) {
        Node::Lit(_) | Node::Name(_) => {}
        Node::Bin(_, l, r) => {
            direct_refs(inst, l, cx, out);
            direct_refs(inst, r, cx, out);
        }
        Node::Call { args, .. } => {
            for a in args {
                let (Arg::Pos(x) | Arg::Named(_, x)) = a;
                direct_refs(inst, x, cx, out);
            }
        }
        Node::Read { path, arg, .. } => {
            out.push(path.to_string());
            direct_refs(inst, arg, cx, out);
        }
        Node::Own { arg, .. } => direct_refs(inst, arg, cx, out),
        Node::Signal { of, arg, .. } => {
            direct_refs(inst, of.expr, inst.signal(of, cx), out);
            direct_refs(inst, arg, cx, out);
        }
    }
}

/// One dependencies-first walk: the groups, what each node reads, and which are loops.
pub struct Order {
    pub groups: Vec<Vec<String>>,
    deps: BTreeMap<String, Vec<String>>,
}

impl Order {
    pub fn deps(&self, path: &str) -> &[String] {
        self.deps.get(path).map_or(&[], Vec::as_slice)
    }

    pub(crate) fn within(&self, kept: &BTreeSet<String>) -> Order {
        Order {
            groups: self
                .groups
                .iter()
                .filter(|group| group.iter().all(|path| kept.contains(path)))
                .cloned()
                .collect(),
            deps: self.deps.clone(),
        }
    }

    /// A group of one whose node does not ref itself is an ordinary node; anything else is a loop.
    pub fn is_loop(&self, group: &[String]) -> bool {
        match group {
            [only] => self.deps(only).iter().any(|d| d == only),
            _ => true,
        }
    }
}

pub fn direct_deps(inst: &Instances, path: &str) -> Result<Vec<String>, EngineError> {
    let (e, cx) = inst
        .at(path)
        .ok_or_else(|| EngineError::UnknownNode(path.to_string()))?;
    let mut out = Vec::new();
    direct_refs(inst, e, cx, &mut out);
    out.sort();
    out.dedup();
    for target in &out {
        if !inst.holds(target) {
            return Err(EngineError::UnknownNode(target.clone()));
        }
    }
    Ok(out)
}

struct Frame {
    node: String,
    refs: Vec<String>,
    idx: usize,
}

/// A node two roots reach is grouped once, whichever reached it first.
pub fn schedule_from(inst: &Instances, roots: &[String]) -> Result<Order, EngineError> {
    let mut walk = Walk {
        inst,
        deps: BTreeMap::new(),
        index: HashMap::new(),
        low: HashMap::new(),
        open: Vec::new(),
        next: 0,
        groups: Vec::new(),
    };
    for root in roots {
        walk.from(root)?;
    }
    Ok(Order {
        groups: walk.groups,
        deps: walk.deps,
    })
}

struct Walk<'a> {
    inst: &'a Instances<'a>,
    deps: BTreeMap<String, Vec<String>>,
    index: HashMap<String, usize>,
    low: HashMap<String, usize>,
    open: Vec<String>,
    next: usize,
    groups: Vec<Vec<String>>,
}

impl Walk<'_> {
    fn from(&mut self, root: &str) -> Result<(), EngineError> {
        if !self.inst.holds(root) {
            return Err(EngineError::UnknownNode(root.to_string()));
        }
        if self.index.contains_key(root) {
            return Ok(());
        }
        self.index.insert(root.to_string(), self.next);
        self.low.insert(root.to_string(), self.next);
        self.next += 1;
        self.open.push(root.to_string());
        let seed = direct_deps(self.inst, root)?;
        self.deps.insert(root.to_string(), seed.clone());
        let mut stack = vec![Frame {
            node: root.to_string(),
            refs: seed,
            idx: 0,
        }];

        while let Some(frame) = stack.last_mut() {
            if frame.idx < frame.refs.len() {
                let target = frame.refs[frame.idx].clone();
                let node = frame.node.clone();
                frame.idx += 1;
                match self.index.get(&target).copied() {
                    None => {
                        self.index.insert(target.clone(), self.next);
                        self.low.insert(target.clone(), self.next);
                        self.next += 1;
                        self.open.push(target.clone());
                        let refs = direct_deps(self.inst, &target)?;
                        self.deps.insert(target.clone(), refs.clone());
                        stack.push(Frame {
                            node: target,
                            refs,
                            idx: 0,
                        });
                    }
                    Some(at) if self.open.contains(&target) => {
                        let mine = self.low[&node];
                        self.low.insert(node, mine.min(at));
                    }
                    Some(_) => {}
                }
                continue;
            }

            let node = frame.node.clone();
            let mine = self.low[&node];
            stack.pop();
            if let Some(parent) = stack.last() {
                let above = self.low[&parent.node];
                self.low.insert(parent.node.clone(), above.min(mine));
            }
            if mine == self.index[&node] {
                let at = self
                    .open
                    .iter()
                    .rposition(|n| *n == node)
                    .expect("a root of its group is still open");
                let mut group = self.open.split_off(at);
                group.sort();
                self.groups.push(group);
            }
        }
        Ok(())
    }
}

/// What a render has to hold, and what it can leave as a closed form.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Schedule {
    /// What a reading holds for itself, the root among them.
    pub wanted: Vec<NodeId>,
    /// The nodes a reading asked a closed form of, none already held.
    pub compose: Vec<NodeId>,
}

/// A closed form is held as samples only under a buffer reading, a ledger, or the root.
pub fn plan(typing: &Typing, root: NodeId, asks: &[Ask]) -> Schedule {
    let mut wanted: BTreeSet<NodeId> = BTreeSet::new();
    let mut compose: Vec<NodeId> = Vec::new();
    let audio = asks.is_empty();
    for ask in asks {
        let Some(id) = typing.id(&ask.node) else {
            continue;
        };
        // FORMAT 14.2: `bindings` and `arguments` are structural and `flops` counts; none composes.
        if matches!(
            ask.representation,
            crate::query::Representation::Bindings
                | crate::query::Representation::Arguments
                | crate::query::Representation::Flops
        ) {
            continue;
        }
        match ask.representation.consumes(typing.ty(id).is_closed_form()) {
            sva_samples::Consumes::ClosedForm if !compose.contains(&id) => compose.push(id),
            sva_samples::Consumes::ClosedForm => {}
            _ => {
                wanted.insert(id);
            }
        }
        if let crate::query::Representation::Ledger { depth } = ask.representation {
            attributed(typing, id, depth, &mut wanted);
        }
    }
    if audio || !wanted.is_empty() {
        wanted.insert(root);
    }
    compose.retain(|id| !wanted.contains(id));
    Schedule {
        wanted: wanted.into_iter().collect(),
        compose,
    }
}

/// A ledger names every ref under its target, so each is a buffer of its own.
fn attributed(typing: &Typing, id: NodeId, depth: usize, wanted: &mut BTreeSet<NodeId>) {
    // Level by level, as the reading walks: a node is held at its shortest chain's depth.
    let mut seen = BTreeSet::from([id]);
    let mut level = vec![id];
    for _ in 0..depth {
        let mut next = Vec::new();
        while let Some(held) = level.pop() {
            for operand in read_operands(typing, held) {
                if !seen.insert(operand) {
                    continue;
                }
                wanted.insert(operand);
                match typing.name(operand) == typing.name(held) {
                    true => level.push(operand),
                    false => next.push(operand),
                }
            }
        }
        level = next;
    }
}

/// A node reading its own output is one program, whatever its width: a component lowered on
/// its own would read the loop at that component's width instead of the node's.
pub(crate) fn holds_self(typing: &Typing, id: NodeId, seen: &mut BTreeSet<NodeId>) -> bool {
    if !seen.insert(id) {
        return false;
    }
    match typing.value(id) {
        Value::SelfAt { .. } => true,
        Value::Cast(Cast::Sample, _) | Value::Read { .. } => false,
        Value::Cast(_, source) => holds_self(typing, *source, seen),
        Value::Op { args, .. } => args.iter().any(|a| holds_self(typing, *a, seen)),
        Value::Filter {
            x, cutoff, q, gain, ..
        } => [x, cutoff, q, gain]
            .into_iter()
            .any(|operand| holds_self(typing, *operand, seen)),
        Value::Solver { varying, .. } => varying.iter().any(|(_, a)| holds_self(typing, *a, seen)),
        Value::ClosedForm(_) | Value::Noise(_) | Value::Stored(_) => false,
    }
}

/// Which nodes under `id` a render has to hold before it can hold `id` itself: the buffers its
/// program reads, an operation, filter or read under it being part of that program.
pub(crate) fn materialized_operands(typing: &Typing, id: NodeId) -> Vec<NodeId> {
    let sampled = |set: Vec<NodeId>| -> Vec<NodeId> {
        let mut out = Vec::new();
        for op in set {
            if typing.ty(op).is_closed_form() {
                continue;
            }
            let inlined = matches!(
                typing.value(op),
                Value::Op { .. } | Value::Filter { .. } | Value::Read { .. }
            );
            match inlined || holds_self(typing, op, &mut BTreeSet::new()) {
                true => out.extend(materialized_operands(typing, op)),
                false => out.push(op),
            }
        }
        out
    };
    match typing.value(id) {
        Value::ClosedForm(_) | Value::Noise(_) | Value::Stored(_) => Vec::new(),
        Value::SelfAt { at, .. } => sampled(at.moving()),
        Value::Solver { varying, .. } => sampled(varying.iter().map(|(_, a)| *a).collect()),
        Value::Cast(Cast::Sample, source) => vec![*source],
        Value::Read { source, at, .. } => {
            let mut out = match (at, anywhere(typing, *source)) {
                (When::Moving(_) | When::Step(_), true) => Vec::new(),
                _ => vec![*source],
            };
            out.extend(sampled(at.moving()));
            out
        }
        Value::Cast(_, source) => sampled(vec![*source]),
        Value::Op { args, .. } => sampled(args.clone()),
        Value::Filter {
            x, cutoff, q, gain, ..
        } => sampled(vec![*x, *cutoff, *q, *gain]),
    }
}

/// Whether a read of `id` has a value at any instant, stored or not.
pub(crate) fn anywhere(typing: &Typing, id: NodeId) -> bool {
    match typing.value(id) {
        Value::Noise(_) => true,
        Value::Cast(Cast::Sample, of) => typing.ty(*of).held == Held::Form(Var::T),
        Value::ClosedForm(form) => form.var == Var::T,
        _ => false,
    }
}

/// Every ref one node reads; `materialized_operands` answers a narrower one.
pub(crate) fn read_operands(typing: &Typing, id: NodeId) -> Vec<NodeId> {
    match typing.value(id) {
        Value::ClosedForm(form) => crate::refs::nodes_in(&form.body),
        _ if typing.ty(id).is_closed_form() => {
            let mut out = Vec::new();
            reads_under(typing, id, &mut BTreeSet::new(), &mut out);
            out.dedup();
            out
        }
        _ => materialized_operands(typing, id),
    }
}

/// A closed-form-typed node that holds no written one still reads refs through its operands.
fn reads_under(typing: &Typing, id: NodeId, seen: &mut BTreeSet<NodeId>, out: &mut Vec<NodeId>) {
    if !seen.insert(id) {
        return;
    }
    match typing.value(id) {
        Value::ClosedForm(form) => out.extend(crate::refs::nodes_in(&form.body)),
        Value::Read { source, .. } => out.push(*source),
        Value::Cast(_, source) => read_through(typing, *source, seen, out),
        Value::Op { args, .. } => {
            for arg in args {
                read_through(typing, *arg, seen, out);
            }
        }
        Value::Filter {
            x, cutoff, q, gain, ..
        } => {
            for operand in [x, cutoff, q, gain] {
                read_through(typing, *operand, seen, out);
            }
        }
        Value::Solver { varying, .. } => {
            for (_, arg) in varying {
                read_through(typing, *arg, seen, out);
            }
        }
        Value::SelfAt { .. } | Value::Noise(_) | Value::Stored(_) => {}
    }
}

/// A written form under an operand is the term read, whatever transform sits between.
fn read_through(typing: &Typing, id: NodeId, seen: &mut BTreeSet<NodeId>, out: &mut Vec<NodeId>) {
    let Value::ClosedForm(_) = typing.value(id) else {
        return reads_under(typing, id, seen, out);
    };
    if seen.insert(id) {
        out.push(id);
    }
}