vj-sim 0.3.0

//! The implicit solve order of function arguments (task C13; IEEE 1800-2023
//! 18.5.11).
//!
//! "Random variables used as function arguments get solved first": a call's
//! return value is a state variable, so every random variable the call's
//! arguments mention must have its value before the constraint holding the call
//! can be solved. The LRM puts it as an implicit variable ordering: the
//! argument variables have priority over the other random variables of the
//! constraint, constraints over only higher-priority variables are solved
//! first, and a circular ordering is an error.
//!
//! [`stages`] computes that order as a level per variable: stage 0 has no
//! argument obligation, a variable constrained together with a call over `a`
//! sits one stage after `a` (the longest chain decides). The registry solves
//! stage by stage, the values of earlier stages being constants (state) for
//! the later ones.
//!
//! Only the *outermost* call of an expression yields edges: the variables of its
//! whole argument tree (nested calls included) precede the conjunct's other
//! variables. A variable that is an argument of any call of the conjunct is not
//! "other" for any of them, so `F(a) == G(b)` puts `a` and `b` in one stage (the
//! equation between two return values is then a check on a drawn pair, not a
//! circular order). The `dist` selector variables (numbered at or after
//! `n_vars`) take no part: they are not random variables of the class.

use super::dist::SELECTOR_BASE;
use super::ir::{ArrayNode, Expr, VarId};
use std::collections::BTreeSet;

/// The implicit order is circular (LRM 18.5.11).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Cycle;

/// The variables of `e` outside any call, and of the argument trees of its
/// outermost calls.
fn split(e: &Expr, others: &mut BTreeSet<VarId>, args: &mut BTreeSet<VarId>) {
    match e {
        Expr::Array(ArrayNode::Call { args: a, .. }) => {
            for x in a {
                x.for_each_var(&mut |v, _| {
                    args.insert(v);
                });
            }
        }
        Expr::Var(v, _) => {
            others.insert(*v);
        }
        other => {
            let _ = other.try_map_children::<()>(&mut |c| {
                split(c, others, args);
                Ok(c.clone())
            });
        }
    }
}

/// The stage of each of the `n_vars` random variables (module docs). A
/// conjunct set with no call over a random variable puts everything in stage 0.
pub fn stages(n_vars: usize, conjuncts: &[Expr]) -> Result<Vec<u32>, Cycle> {
    let mut edges: BTreeSet<(usize, usize)> = BTreeSet::new();
    for c in conjuncts {
        let (mut others, mut args) = (BTreeSet::new(), BTreeSet::new());
        split(c, &mut others, &mut args);
        let real =
            |v: &VarId| (*v < SELECTOR_BASE && (*v as usize) < n_vars).then_some(*v as usize);
        let args: Vec<usize> = args.iter().filter_map(real).collect();
        for o in others.iter().filter_map(real) {
            if args.contains(&o) {
                continue;
            }
            for &a in &args {
                edges.insert((a, o));
            }
        }
    }
    // Longest path by Kahn's algorithm.
    let mut indeg = vec![0usize; n_vars];
    let mut succ: Vec<Vec<usize>> = vec![Vec::new(); n_vars];
    for &(a, b) in &edges {
        indeg[b] += 1;
        succ[a].push(b);
    }
    let mut stage = vec![0u32; n_vars];
    let mut ready: Vec<usize> = (0..n_vars).filter(|&v| indeg[v] == 0).collect();
    let mut done = 0;
    while let Some(v) = ready.pop() {
        done += 1;
        for &w in &succ[v] {
            stage[w] = stage[w].max(stage[v] + 1);
            indeg[w] -= 1;
            if indeg[w] == 0 {
                ready.push(w);
            }
        }
    }
    if done != n_vars {
        return Err(Cycle);
    }
    Ok(stage)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::constraints::ir::{BinOp, Ty};

    fn s32() -> Ty {
        Ty::signed(32)
    }
    fn var(i: u32) -> Expr {
        Expr::Var(i, s32())
    }
    fn call(args: Vec<Expr>) -> Expr {
        Expr::Array(ArrayNode::Call {
            func: 0,
            args,
            ty: s32(),
        })
    }
    fn eq(a: Expr, b: Expr) -> Expr {
        Expr::Binary(BinOp::Eq, Box::new(a), Box::new(b), Ty::BIT)
    }

    #[test]
    fn no_call_means_one_stage() {
        assert_eq!(stages(3, &[eq(var(0), var(1))]), Ok(vec![0, 0, 0]));
    }

    #[test]
    fn an_argument_is_solved_before_the_other_variables_of_its_constraint() {
        // x(0) <= F(y(1)): y first.
        let c = eq(var(0), call(vec![var(1)]));
        assert_eq!(stages(2, &[c]), Ok(vec![1, 0]));
    }

    #[test]
    fn chains_take_the_longest_path() {
        // x == F(y), y == F(z): z, y, x.
        let cs = [
            eq(var(0), call(vec![var(1)])),
            eq(var(1), call(vec![var(2)])),
        ];
        assert_eq!(stages(3, &cs), Ok(vec![2, 1, 0]));
    }

    #[test]
    fn two_arguments_of_one_constraint_share_a_stage() {
        // F(a) == G(b): no order between a and b (a check on the pair).
        let c = eq(call(vec![var(0)]), call(vec![var(1)]));
        assert_eq!(stages(2, &[c]), Ok(vec![0, 0]));
        // x == F(x): a variable that is also an argument is not "other".
        assert_eq!(stages(1, &[eq(var(0), call(vec![var(0)]))]), Ok(vec![0]));
    }

    #[test]
    fn nested_calls_put_every_inner_variable_before_the_rest() {
        // x == F(G(a, b)): a and b first.
        let c = eq(var(0), call(vec![call(vec![var(1), var(2)])]));
        assert_eq!(stages(3, &[c]), Ok(vec![1, 0, 0]));
    }

    #[test]
    fn a_circular_order_is_an_error() {
        // x == F(y) and y == F(x): LRM 18.5.11 says circular is an error.
        let cs = [
            eq(var(0), call(vec![var(1)])),
            eq(var(1), call(vec![var(0)])),
        ];
        assert_eq!(stages(2, &cs), Err(Cycle));
    }

    #[test]
    fn selector_variables_take_no_part() {
        let c = eq(var(SELECTOR_BASE), call(vec![var(0)]));
        assert_eq!(stages(1, &[c]), Ok(vec![0]));
    }
}