vj-sim 0.1.0

//! Soft constraints and their priorities (Task C15; IEEE 1800-2023 18.5.13,
//! 18.5.13.1, 18.5.13.2).
//!
//! # Representation
//!
//! A soft constraint is an ordinary boolean conjunct wrapped as
//! `Var(SOFT_BASE + k) || body`: the hidden variable `k` is the soft's *discard
//! flag*. Resolving a call replaces every flag by a constant
//! ([`resolve_softs`]): a discarded soft becomes `true` ("treated as if
//! replaced by the value 1", 18.5.13) and a kept one is its body, so the
//! solver never sees a soft constraint, only the hard problem one choice of
//! flags describes. The wrapper may sit anywhere in a conjunct, so `p -> soft
//! q` and `if (c) soft q; else soft r;` need nothing more.
//!
//! # Priorities (18.5.13.1)
//!
//! [`Priority`] orders the soft constraints of a program: the object (an
//! object contained through a `rand` handle ranks below its container), then
//! the class level (a derived class above its bases, an inline constraint
//! above all), then the declaration order. [`object_order`] is the rule for
//! contained objects: the container above every contained object, a later
//! handle above an earlier one, and an object held through several handles at
//! the rank of the last of them. The class-level and declaration-order rules
//! are what the builder's `Builder::set_priority_scope` and its running
//! sequence number implement.
//!
//! # Solving (18.5.13's conceptual model)
//!
//! [`solve_relaxing`] is the staged relaxation: try every soft constraint; if
//! that has no solution, try the hard constraints alone (their failure is an
//! ordinary failure) and then take the soft constraints in priority order,
//! keeping each one only if it can be satisfied together with the hard
//! constraints and the higher-priority softs already kept. This is the "c1
//! and c2, else c1, else c2, else neither" model applied recursively, and
//! it needs nothing of the solver but a feasibility answer.
//!
//! # `disable soft` (18.5.13.2)
//!
//! A [`Disable`] discards every *lower*-priority soft constraint that directly
//! names the same random variable ([`SoftKey`]), unconditionally;
//! [`apply_disables`] does it once, statically, because priorities and the
//! variables a constraint names are static.

use super::dist::{Dist, SELECTOR_BASE};
use super::ir::{ArrayId, ArrayNode, BinOp, Expr, Ty, VarId};
use num_bigint::BigUint;
use num_traits::One;

/// The first discard-flag id: soft `k`'s flag is `Var(SOFT_BASE + k)`. Above
/// [`SELECTOR_BASE`] so no pass that already skips selector variables can
/// mistake it for a field.
pub const SOFT_BASE: VarId = 1 << 31;

/// The rank of one soft constraint (or `disable soft`); larger is stronger.
/// Ordered by object, then class level, then declaration order.
#[derive(
    Clone,
    Copy,
    Debug,
    Default,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    serde::Serialize,
    serde::Deserialize,
)]
pub struct Priority {
    /// The rank of the object the constraint belongs to: contained objects
    /// rank below their container ([`object_order`]).
    pub object: u32,
    /// The class level within the object: a derived class's constraints rank
    /// above its base classes'; an inline constraint above every class level.
    pub level: u32,
    /// The declaration sequence: later is higher.
    pub seq: u32,
}

/// A random variable a `disable soft` can name, or a soft constraint can
/// directly mention (18.5.13.2: "including `array.size`").
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub enum SoftKey {
    /// A scalar random variable (or one element of a fixed array).
    Var(VarId),
    /// Some element of a dynamic array or queue, or the array as a whole when
    /// it is a `disable soft` target.
    Array(ArrayId),
    /// The size of a dynamic array or queue.
    ArraySize(ArrayId),
}

impl SoftKey {
    /// Whether a `disable soft` on `self` reaches a soft constraint that
    /// mentions `mention` (18.5.13.2: only where the variable *directly*
    /// appears). Naming a whole array reaches its elements and its size;
    /// naming `A.size` only the size.
    pub fn disables(self, mention: SoftKey) -> bool {
        match (self, mention) {
            (SoftKey::Var(a), SoftKey::Var(b)) => a == b,
            (SoftKey::Array(a), SoftKey::Array(b) | SoftKey::ArraySize(b)) => a == b,
            (SoftKey::ArraySize(a), SoftKey::ArraySize(b)) => a == b,
            _ => false,
        }
    }
}

/// One soft constraint of a program.
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Soft {
    /// Its priority.
    pub priority: Priority,
    /// The [`Dist`]s (indices into the program's dist table) it contains.
    pub dists: Vec<u32>,
    /// The random variables it directly mentions.
    pub keys: Vec<SoftKey>,
    /// Discarded for every call by a higher-priority `disable soft`
    /// ([`apply_disables`]).
    pub disabled: bool,
    /// The `disable soft`s ([`Disable`]s, by their index in the program's
    /// table) that reach it: [`Self::disabled`] is whether any is in force. A
    /// `disable soft` in a constraint block `constraint_mode(0)` switched off
    /// is not (LRM 18.9), so a call recomputes the flag from this list.
    pub disabled_by: Vec<u32>,
}

/// One `disable soft x;` constraint.
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Disable {
    /// Its priority: only softs below it are discarded.
    pub priority: Priority,
    /// The variable it names.
    pub key: SoftKey,
}

/// Marks every soft constraint a `disable soft` of higher priority reaches.
pub fn apply_disables(softs: &mut [Soft], disables: &[Disable]) {
    for s in softs {
        s.disabled_by = disables
            .iter()
            .enumerate()
            .filter(|(_, d)| d.priority > s.priority && s.keys.iter().any(|&k| d.key.disables(k)))
            .map(|(i, _)| i as u32)
            .collect();
        s.disabled = !s.disabled_by.is_empty();
    }
}

/// The random variables `e` (and the `dists` it holds) directly mention.
pub fn keys_of(e: &Expr, dists: &[&Dist]) -> Vec<SoftKey> {
    let mut keys: Vec<SoftKey> = Vec::new();
    let mut add = |k: SoftKey| {
        if !keys.contains(&k) {
            keys.push(k);
        }
    };
    e.for_each_var(&mut |v, _| {
        if v < SELECTOR_BASE {
            add(SoftKey::Var(v));
        }
    });
    e.any_node(&mut |n| {
        if let Expr::Array(a) = n {
            match a {
                ArrayNode::Size(id) => add(SoftKey::ArraySize(*id)),
                ArrayNode::Elem { array, .. }
                | ArrayNode::Foreach { array, .. }
                | ArrayNode::Reduce { array, .. } => add(SoftKey::Array(*array)),
                ArrayNode::Distinct { arrays, .. } => {
                    arrays.iter().for_each(|a| add(SoftKey::Array(*a)));
                }
                // A pick's elements are plain variables and states, found above.
                ArrayNode::Loop(..) | ArrayNode::Call { .. } | ArrayNode::Pick { .. } => {}
            }
        }
        false
    });
    for d in dists {
        for &v in &d.vars {
            add(SoftKey::Var(v));
        }
    }
    keys
}

/// The soft constraint with discard flag `id` wrapped around `body`.
pub fn wrap_soft(id: u32, body: Expr) -> Expr {
    Expr::Binary(
        BinOp::LogOr,
        Box::new(Expr::Var(SOFT_BASE + id, Ty::BIT)),
        Box::new(body),
        Ty::BIT,
    )
}

/// `e` with every soft wrapper resolved: `discard[k]` replaces soft `k` by
/// `true`, a kept soft by its body. `Err(k)` names a flag outside `discard`.
pub fn resolve_softs(e: &Expr, discard: &[bool]) -> Result<Expr, u32> {
    let mapped = e.try_map_children(&mut |c| resolve_softs(c, discard))?;
    if let Expr::Binary(BinOp::LogOr, a, b, _) = &mapped
        && let Expr::Var(v, _) = **a
        && v >= SOFT_BASE
    {
        let k = v - SOFT_BASE;
        return match discard.get(k as usize) {
            Some(true) => Ok(Expr::Const(Ty::BIT, BigUint::one())),
            Some(false) => Ok((**b).clone()),
            None => Err(k),
        };
    }
    Ok(mapped)
}

/// Whether `e` holds a soft wrapper.
pub fn has_soft(e: &Expr) -> bool {
    e.any_node(&mut |n| matches!(n, Expr::Var(v, _) if *v >= SOFT_BASE))
}

/// Solves a problem with `softs`: `attempt(discard, seed)` solves it with the
/// soft constraints flagged in `discard` replaced by `true` and returns its
/// solution (it must not have any effect but advancing `seed`);
/// `relaxable(err)` says whether an error means "no solution for this choice
/// of softs" and not a failure of the call.
///
/// Every soft first; if there is no solution, the hard constraints alone (a
/// failure there is the call's), then each soft by priority, kept when it fits
/// with everything kept so far (module docs). Each attempt starts from the
/// call's seed, so the result is exactly the draw a program with only the kept
/// constraints would make, and with no contradiction it is the draw of the
/// all-hard program.
pub fn solve_relaxing<T, E>(
    softs: &[Soft],
    seed: &mut u64,
    mut attempt: impl FnMut(&[bool], &mut u64) -> Result<T, E>,
    relaxable: impl Fn(&E) -> bool,
) -> Result<T, E> {
    let seed0 = *seed;
    let initial: Vec<bool> = softs.iter().map(|s| s.disabled).collect();
    let mut s = seed0;
    match attempt(&initial, &mut s) {
        Ok(t) => {
            *seed = s;
            return Ok(t);
        }
        Err(e) if !relaxable(&e) => return Err(e),
        Err(_) => {}
    }
    // Hard constraints alone: no solution here is the call's failure.
    let mut best_seed = seed0;
    let mut best = attempt(&vec![true; softs.len()], &mut best_seed)?;
    let mut order: Vec<usize> = (0..softs.len()).filter(|&k| !softs[k].disabled).collect();
    order.sort_by(|&a, &b| softs[b].priority.cmp(&softs[a].priority).then(b.cmp(&a)));
    let mut keep = vec![false; softs.len()];
    for k in order {
        keep[k] = true;
        let discard: Vec<bool> = keep.iter().map(|&b| !b).collect();
        let mut s = seed0;
        match attempt(&discard, &mut s) {
            Ok(t) => {
                best = t;
                best_seed = s;
            }
            Err(e) if relaxable(&e) => keep[k] = false,
            Err(e) => return Err(e),
        }
    }
    *seed = best_seed;
    Ok(best)
}

/// The objects reached from `root` through `contained` (its `rand` object
/// handles in declaration order), lowest priority first and `root` last
/// (18.5.13.1): every container above the objects it contains, a later handle
/// above an earlier one, and an object contained several times at the rank of
/// the highest, i.e. last, occurrence. An object graph can be cyclic, so an
/// object is not entered again while it is being expanded.
pub fn object_order<K: Clone + PartialEq>(root: &K, contained: &impl Fn(&K) -> Vec<K>) -> Vec<K> {
    fn walk<K: Clone + PartialEq>(
        o: &K,
        path: &mut Vec<K>,
        contained: &impl Fn(&K) -> Vec<K>,
        out: &mut Vec<K>,
    ) {
        if path.contains(o) {
            return;
        }
        path.push(o.clone());
        for c in contained(o) {
            walk(&c, path, contained, out);
        }
        path.pop();
        out.push(o.clone());
    }
    let mut all = Vec::new();
    walk(root, &mut Vec::new(), contained, &mut all);
    let mut kept: Vec<K> = Vec::new();
    for o in all.into_iter().rev() {
        if !kept.contains(&o) {
            kept.push(o);
        }
    }
    kept.reverse();
    kept
}

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

    fn s(object: u32, level: u32, seq: u32) -> Priority {
        Priority { object, level, seq }
    }

    fn soft(priority: Priority, keys: Vec<SoftKey>) -> Soft {
        Soft {
            priority,
            dists: Vec::new(),
            keys,
            disabled: false,
            disabled_by: Vec::new(),
        }
    }

    // ---- 18.5.13.1's worked example -------------------------------------

    /// A class of the example: its base, and per class level its soft
    /// constraints in declaration order.
    struct Class {
        base: Option<&'static Class>,
        /// This level's constraint blocks in the order that fixes their
        /// priority (an `extern` block at its prototype), each with the
        /// labels of its soft constraints in order.
        blocks: &'static [&'static [&'static str]],
        /// `rand` handle members: (name, class).
        handles: &'static [(&'static str, &'static Class)],
    }

    static B1: Class = Class {
        base: None,
        blocks: &[&["a1", "a2"]],
        handles: &[],
    };
    static D1: Class = Class {
        base: Some(&B1),
        blocks: &[&["b1"]],
        handles: &[],
    };
    static B2: Class = Class {
        base: None,
        blocks: &[&["c1"]],
        handles: &[],
    };
    static D2: Class = Class {
        base: Some(&B2),
        // d, the extern block e (at its prototype, so before f), then f
        blocks: &[&["d1"], &["e1"], &["f1"]],
        handles: &[("p1", &D1), ("p2", &B1), ("p3", &D1)],
    };

    /// An object of the example: its class and a name for its identity.
    #[derive(Clone, Copy)]
    struct Obj {
        class: &'static Class,
        id: u32,
    }

    /// Objects are the same when their ids are.
    impl PartialEq for Obj {
        fn eq(&self, other: &Obj) -> bool {
            self.id == other.id
        }
    }

    /// The soft constraints of one object with their priorities: base levels
    /// first (lower), then the class's, blocks and constraints in order.
    fn object_softs(o: Obj, rank: u32, prefix: &str) -> Vec<(Priority, String)> {
        let mut chain: Vec<&'static Class> = Vec::new();
        let mut c = Some(o.class);
        while let Some(k) = c {
            chain.push(k);
            c = k.base;
        }
        chain.reverse(); // root base first
        let mut seq = 0;
        let mut out = Vec::new();
        for (level, class) in chain.iter().enumerate() {
            for block in class.blocks {
                for label in *block {
                    out.push((s(rank, level as u32, seq), format!("{prefix}{label}")));
                    seq += 1;
                }
            }
        }
        out
    }

    /// Objects of the example: id 0 is `d`; `p1`, `p2`, `p3` are 1, 2, 3 (or
    /// `p3` is `p1` itself when `same`).
    fn contained(same: bool) -> impl Fn(&Obj) -> Vec<Obj> {
        move |o: &Obj| {
            o.class
                .handles
                .iter()
                .enumerate()
                .map(|(i, (_, class))| Obj {
                    class,
                    id: if same && i == 2 { 1 } else { i as u32 + 1 },
                })
                .collect()
        }
    }

    /// The whole call's soft constraints, highest priority first, as the
    /// subclause writes them: the inline constraints on top.
    fn chain(same: bool) -> String {
        let root = Obj { class: &D2, id: 0 };
        let objects = object_order(&root, &contained(same));
        let mut all: Vec<(Priority, String)> = Vec::new();
        for (rank, o) in objects.iter().enumerate() {
            let prefix = match (o.class as *const Class, o.id) {
                (_, 0) => String::new(),
                (_, 1) if same => "p3.".to_string(),
                (_, 1) => "p1.".to_string(),
                (_, 2) => "p2.".to_string(),
                _ => "p3.".to_string(),
            };
            all.extend(object_softs(*o, rank as u32, &prefix));
        }
        // inline constraints: above every class level of every object
        all.push((s(u32::MAX, u32::MAX, 0), "i1".to_string()));
        all.push((s(u32::MAX, u32::MAX, 1), "i2".to_string()));
        all.sort_by_key(|b| std::cmp::Reverse(b.0));
        all.into_iter()
            .map(|(_, l)| l)
            .collect::<Vec<_>>()
            .join(" ")
    }

    #[test]
    fn the_priority_chain_is_the_one_18_5_13_1_gives() {
        assert_eq!(
            chain(false),
            "i2 i1 f1 e1 d1 c1 p3.b1 p3.a2 p3.a1 p2.a2 p2.a1 p1.b1 p1.a2 p1.a1"
        );
    }

    #[test]
    fn the_same_object_through_two_handles_takes_the_higher_rank() {
        assert_eq!(
            chain(true),
            "i2 i1 f1 e1 d1 c1 p3.b1 p3.a2 p3.a1 p2.a2 p2.a1"
        );
    }

    #[test]
    fn object_order_expands_nested_objects_and_survives_cycles() {
        // 0 holds 1 then 2; 1 holds 3; 3 holds 0 again (a cycle).
        let contained = |o: &u32| match *o {
            0 => vec![1, 2],
            1 => vec![3],
            3 => vec![0],
            _ => Vec::new(),
        };
        // 3 is contained in 1, 1 before 2; the cycle back to 0 is not entered.
        assert_eq!(object_order(&0, &contained), vec![3, 1, 2, 0]);
        // An object reached again later takes the later rank.
        let again = |o: &u32| match *o {
            0 => vec![1, 2, 1],
            _ => Vec::new(),
        };
        assert_eq!(object_order(&0, &again), vec![2, 1, 0]);
    }

    // ---- disable soft ---------------------------------------------------

    #[test]
    fn a_disable_discards_only_lower_priority_softs_that_mention_the_variable() {
        let x = SoftKey::Var(0);
        let y = SoftKey::Var(1);
        let mut softs = vec![
            soft(s(0, 0, 0), vec![x]),    // A1: below the disable
            soft(s(0, 0, 2), vec![x]),    // A3: above it
            soft(s(0, 0, 0), vec![y]),    // another variable
            soft(s(0, 0, 0), vec![x, y]), // mentions x directly
        ];
        let disables = [Disable {
            priority: s(0, 0, 1),
            key: x,
        }];
        apply_disables(&mut softs, &disables);
        let flags: Vec<bool> = softs.iter().map(|s| s.disabled).collect();
        assert_eq!(flags, vec![true, false, false, true]);
    }

    #[test]
    fn a_disable_on_a_whole_array_reaches_its_size_and_a_size_disable_only_the_size() {
        let size = SoftKey::ArraySize(3);
        let elem = SoftKey::Array(3);
        assert!(elem.disables(size) && elem.disables(elem));
        assert!(size.disables(size) && !size.disables(elem));
        assert!(!SoftKey::Array(4).disables(size));
        assert!(!SoftKey::Var(3).disables(elem));
    }

    // ---- resolve_softs ---------------------------------------------------

    fn var(v: u32) -> Expr {
        Expr::Var(v, Ty::unsigned(8))
    }

    fn lt(a: Expr, b: Expr) -> Expr {
        Expr::Binary(BinOp::Lt, Box::new(a), Box::new(b), Ty::BIT)
    }

    #[test]
    fn a_discarded_soft_is_true_and_a_kept_one_is_its_body() {
        let body = lt(var(0), Expr::Const(Ty::unsigned(8), BigUint::from(3u32)));
        let e = Expr::Binary(
            BinOp::LogAnd,
            Box::new(wrap_soft(0, body.clone())),
            Box::new(wrap_soft(1, body.clone())),
            Ty::BIT,
        );
        assert!(has_soft(&e));
        let one = Expr::Const(Ty::BIT, BigUint::one());
        let r = resolve_softs(&e, &[false, true]).unwrap();
        assert_eq!(
            r,
            Expr::Binary(BinOp::LogAnd, Box::new(body), Box::new(one), Ty::BIT)
        );
        assert!(!has_soft(&r));
        assert_eq!(resolve_softs(&e, &[false]), Err(1));
    }

    #[test]
    fn the_variables_a_soft_mentions_exclude_the_flags_and_selectors() {
        let e = wrap_soft(0, lt(var(2), var(5)));
        assert_eq!(
            keys_of(&e, &[]),
            vec![SoftKey::Var(2), SoftKey::Var(5)],
            "the discard flag is not a mention"
        );
        let arr = Expr::Array(ArrayNode::Size(7));
        assert_eq!(keys_of(&arr, &[]), vec![SoftKey::ArraySize(7)]);
    }

    // ---- solve_relaxing --------------------------------------------------

    /// A toy problem: a solution is a set of kept softs; softs `a` and `b`
    /// conflict. `Err(())` when the kept set contains both.
    fn toy(
        softs: &[Soft],
        conflict: (usize, usize),
        hard_fails: bool,
    ) -> Result<Vec<usize>, &'static str> {
        let mut seed = 5u64;
        let mut calls = 0;
        let r = solve_relaxing(
            softs,
            &mut seed,
            |discard, s| {
                calls += 1;
                *s += 1;
                if hard_fails {
                    return Err("unsat");
                }
                let kept: Vec<usize> = (0..discard.len()).filter(|&k| !discard[k]).collect();
                if kept.contains(&conflict.0) && kept.contains(&conflict.1) {
                    return Err("unsat");
                }
                Ok(kept)
            },
            |e| *e == "unsat",
        );
        let _ = calls;
        r
    }

    #[test]
    fn without_a_contradiction_every_soft_is_kept() {
        let softs = vec![soft(s(0, 0, 0), vec![]), soft(s(0, 0, 1), vec![])];
        assert_eq!(toy(&softs, (7, 8), false), Ok(vec![0, 1]));
    }

    #[test]
    fn the_higher_priority_of_two_conflicting_softs_wins() {
        let softs = vec![
            soft(s(0, 0, 0), vec![]),
            soft(s(0, 0, 1), vec![]),
            soft(s(0, 0, 2), vec![]),
        ];
        // 0 and 2 conflict: 2 is later, so higher; 0 is dropped, 1 stays.
        assert_eq!(toy(&softs, (0, 2), false), Ok(vec![1, 2]));
        // 1 and 2 conflict: 1 is dropped.
        assert_eq!(toy(&softs, (1, 2), false), Ok(vec![0, 2]));
    }

    #[test]
    fn a_higher_derived_level_beats_a_later_base_declaration() {
        // soft 0 is in the derived class (level 1) though declared first.
        let softs = vec![soft(s(0, 1, 0), vec![]), soft(s(0, 0, 1), vec![])];
        assert_eq!(toy(&softs, (0, 1), false), Ok(vec![0]));
    }

    #[test]
    fn a_hard_failure_is_reported_and_disabled_softs_never_return() {
        let softs = vec![soft(s(0, 0, 0), vec![])];
        assert_eq!(toy(&softs, (7, 8), true), Err("unsat"));
        let mut d = soft(s(0, 0, 1), vec![]);
        d.disabled = true;
        let softs = vec![soft(s(0, 0, 0), vec![]), d];
        assert_eq!(toy(&softs, (0, 1), false), Ok(vec![0]));
    }

    #[test]
    fn the_seed_is_the_kept_attempts_not_the_probes() {
        let softs = vec![soft(s(0, 0, 0), vec![]), soft(s(0, 0, 1), vec![])];
        let mut seed = 100u64;
        let r = solve_relaxing(
            &softs,
            &mut seed,
            |discard, s| {
                *s += 1;
                if discard.iter().all(|d| !d) {
                    Err("unsat")
                } else {
                    Ok(discard.to_vec())
                }
            },
            |e: &&str| *e == "unsat",
        );
        assert_eq!(r, Ok(vec![true, false]));
        assert_eq!(seed, 101, "one draw on top of the call's own seed");
    }
}