vj-sim 0.3.0

//! The SAT fallback for problems whose BDD exceeds its node budget (plan
//! section 3.4, ledger XC-063).
//!
//! Some constraint sets have no small BDD in any variable order (a wide
//! multiplier is the canonical one, [`super::bdd`]). For those, [`SatSolver`]
//! blasts the same IR with the same gadgets ([`super::blast`]) into a CNF with
//! a Tseitin encoding and searches it with `splr`, a pure-Rust CDCL solver.
//! [`super::solve::Engine`] picks between the two: the exact BDD solver when
//! it fits, this one when it does not.
//!
//! # What is the same and what is not
//!
//! * The set of legal assignments is identical: both back ends run the same
//!   gadgets, so division by zero is "not legal" here too (see
//!   [`super::blast`]). Every sample is also checked with the IR's own
//!   evaluator ([`super::ir::eval`]) before it is returned; a model the
//!   evaluator rejects is the typed [`SolveError::SatModelRejected`], never a
//!   silent wrong answer.
//! * The distribution is NOT exactly uniform. A CDCL solver finds *a* model,
//!   not a uniformly random one, and a model count is not available. The
//!   sampler fixes the bits of each solve-order group one at a time in a
//!   random order, choosing each bit by a fair coin and falling back to the
//!   other value only when the coin's choice leaves no model (a feasibility
//!   solve, skipped when the last model already agrees). Every legal
//!   assignment is reachable and every draw is legal, but the probabilities
//!   are only approximately uniform. The exact path keeps its exactness; this
//!   is the price of a problem the exact path cannot hold.
//! * `solve ... before` groups, `randc` variables, and `dist` selectors follow
//!   the exact sampler's structure ([`super::solve::Solver::sample_plan`]): the
//!   `randc` variables first, then each group with its weighted selectors
//!   before its uniform variables, each later choice constrained by the
//!   earlier ones. What it cannot do it refuses, loudly:
//!   [`SolveError::SatUnsupported`] for a `dist` over an expression with more
//!   than [`MAX_COUNT_MODELS`] feasible assignments (fewer are counted by
//!   enumeration), a `dist` whose subject is wider than 32 bits or has more
//!   than `1 << MAX_ENUM_BITS` feasible values under one weighted range, and a
//!   `randc` variable wider than [`MAX_ENUM_BITS`] bits (LRM 18.4.2 lets an
//!   implementation limit `randc` size; all of those are enumerated value by
//!   value).
//! * `randc`: the cycle restarts whenever the CNF changes at all, which is
//!   more often than the exact path's "the variable's own legal set changed"
//!   (a state value read by an unrelated constraint restarts the cycle).
//!
//! # Limits
//!
//! The CNF is capped at [`SatLimits::clause_limit`] clauses (a typed
//! [`BddBuildError::CnfTooLarge`], reported as the `TooLarge` outcome) and a
//! sampling call at [`SatLimits::timeout`] of search (a typed
//! [`SolveError::SatTimeout`]). `splr` takes its timeout in whole seconds and
//! checks it between restarts, so a call can overrun by a fraction of a second.

use super::bdd::BddBuildError;
use super::blast::{Blaster, Gates};
use super::ir::{Assignment, Expr, Ty, VarId, eval};
use super::randc::{self, RandcState, fingerprint_of};
use super::rng;
use super::solve::{
    Group, Measure, Plan, ResolvedWeighted, SolveError, draw_weighted, resolve_plan,
};
use core::fmt;
use num_bigint::BigUint;
use num_traits::Zero;
use splr::{Certificate, Config, SolveIF, Solver as Cdcl, SolverError};
use std::cell::RefCell;
use std::collections::HashMap;
use std::time::{Duration, Instant};

/// The most bits of one variable the sampler enumerates value by value: a
/// `randc` variable's legal set, a `dist` subject's feasible values.
pub const MAX_ENUM_BITS: usize = 10;

/// The most distinct assignments [`SatSolver::count_models`] enumerates before
/// it refuses.
pub const MAX_COUNT_MODELS: usize = 512;

/// The default cap on CNF clauses (about 100 bytes each with the solver's own
/// copy; a 64-bit multiplier is about 100,000).
pub const CLAUSE_LIMIT: usize = 4_000_000;

/// The default search time allowed to one sampling call.
pub const TIMEOUT: Duration = Duration::from_secs(30);

/// Resource limits of the SAT fallback.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SatLimits {
    /// The most clauses the CNF may hold.
    pub clause_limit: usize,
    /// The search time one sampling call may use.
    pub timeout: Duration,
}

impl Default for SatLimits {
    fn default() -> Self {
        SatLimits {
            clause_limit: CLAUSE_LIMIT,
            timeout: TIMEOUT,
        }
    }
}

// ---------- the CNF back end of the shared gadgets ----------

/// A bit of the blasted formula: a constant or a (signed, 1-based) CNF literal.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum SatBit {
    Const(bool),
    Lit(i32),
}

/// The growing CNF: a flat literal pool with clause end offsets, and a
/// structural cache so a gate over the same inputs is encoded once.
struct Cnf {
    lits: Vec<i32>,
    ends: Vec<usize>,
    /// The last CNF variable allocated.
    last_var: i32,
    /// `(operation, a, b, c)` to the gate's output variable.
    cache: HashMap<(u8, i32, i32, i32), i32>,
}

/// The Tseitin encoder implementing [`Gates`].
struct SatGates<'a> {
    cnf: RefCell<Cnf>,
    /// First (0-based) CNF variable of each random variable's bits.
    offsets: &'a [usize],
    limit: usize,
}

const OP_AND: u8 = 0;
const OP_XOR: u8 = 1;
const OP_MUX: u8 = 2;

impl SatGates<'_> {
    fn clause(&self, during: &'static str, c: &[i32]) -> Result<(), BddBuildError> {
        let mut cnf = self.cnf.borrow_mut();
        if cnf.ends.len() >= self.limit {
            return Err(BddBuildError::CnfTooLarge {
                limit: self.limit,
                during,
            });
        }
        cnf.lits.extend_from_slice(c);
        let end = cnf.lits.len();
        cnf.ends.push(end);
        Ok(())
    }

    fn cached(&self, key: (u8, i32, i32, i32)) -> Option<i32> {
        self.cnf.borrow().cache.get(&key).copied()
    }

    fn fresh(&self, key: (u8, i32, i32, i32)) -> i32 {
        let mut cnf = self.cnf.borrow_mut();
        cnf.last_var += 1;
        let g = cnf.last_var;
        cnf.cache.insert(key, g);
        g
    }
}

impl Gates for SatGates<'_> {
    type Bit = SatBit;

    fn constant(&self, v: bool) -> SatBit {
        SatBit::Const(v)
    }

    fn as_const(&self, a: &SatBit) -> Option<bool> {
        match a {
            SatBit::Const(v) => Some(*v),
            SatBit::Lit(_) => None,
        }
    }

    fn input(&self, var: VarId, i: usize) -> SatBit {
        SatBit::Lit(lit_of(self.offsets[var as usize] + i, true))
    }

    fn not(&self, a: &SatBit) -> SatBit {
        match a {
            SatBit::Const(v) => SatBit::Const(!v),
            SatBit::Lit(l) => SatBit::Lit(-l),
        }
    }

    fn and(&self, d: &'static str, a: &SatBit, b: &SatBit) -> Result<SatBit, BddBuildError> {
        let (x, y) = match (*a, *b) {
            (SatBit::Const(false), _) | (_, SatBit::Const(false)) => {
                return Ok(SatBit::Const(false));
            }
            (SatBit::Const(true), v) | (v, SatBit::Const(true)) => return Ok(v),
            (SatBit::Lit(x), SatBit::Lit(y)) => (x, y),
        };
        if x == y {
            return Ok(SatBit::Lit(x));
        }
        if x == -y {
            return Ok(SatBit::Const(false));
        }
        let key = (OP_AND, x.min(y), x.max(y), 0);
        if let Some(g) = self.cached(key) {
            return Ok(SatBit::Lit(g));
        }
        let g = self.fresh(key);
        self.clause(d, &[-g, x])?;
        self.clause(d, &[-g, y])?;
        self.clause(d, &[g, -x, -y])?;
        Ok(SatBit::Lit(g))
    }

    fn or(&self, d: &'static str, a: &SatBit, b: &SatBit) -> Result<SatBit, BddBuildError> {
        // De Morgan over the shared AND gate.
        let both_off = self.and(d, &self.not(a), &self.not(b))?;
        Ok(self.not(&both_off))
    }

    fn xor(&self, d: &'static str, a: &SatBit, b: &SatBit) -> Result<SatBit, BddBuildError> {
        let (x, y) = match (*a, *b) {
            (SatBit::Const(p), SatBit::Const(q)) => return Ok(SatBit::Const(p ^ q)),
            (SatBit::Const(false), v) | (v, SatBit::Const(false)) => return Ok(v),
            (SatBit::Const(true), v) | (v, SatBit::Const(true)) => return Ok(self.not(&v)),
            (SatBit::Lit(x), SatBit::Lit(y)) => (x, y),
        };
        if x == y {
            return Ok(SatBit::Const(false));
        }
        if x == -y {
            return Ok(SatBit::Const(true));
        }
        // `x ^ y` over signs is `|x| ^ |y|` flipped when exactly one is negative.
        let flip = (x < 0) != (y < 0);
        let (p, q) = (x.abs(), y.abs());
        let key = (OP_XOR, p.min(q), p.max(q), 0);
        let g = match self.cached(key) {
            Some(g) => g,
            None => {
                let g = self.fresh(key);
                self.clause(d, &[-g, p, q])?;
                self.clause(d, &[-g, -p, -q])?;
                self.clause(d, &[g, -p, q])?;
                self.clause(d, &[g, p, -q])?;
                g
            }
        };
        Ok(SatBit::Lit(if flip { -g } else { g }))
    }

    fn mux(
        &self,
        d: &'static str,
        s: &SatBit,
        a: &SatBit,
        b: &SatBit,
    ) -> Result<SatBit, BddBuildError> {
        let sl = match *s {
            SatBit::Const(true) => return Ok(*a),
            SatBit::Const(false) => return Ok(*b),
            SatBit::Lit(l) => l,
        };
        if a == b {
            return Ok(*a);
        }
        let (x, y) = match (*a, *b) {
            (SatBit::Const(true), _) => return self.or(d, s, b),
            (SatBit::Const(false), _) => return self.and(d, &self.not(s), b),
            (_, SatBit::Const(true)) => return self.or(d, &self.not(s), a),
            (_, SatBit::Const(false)) => return self.and(d, s, a),
            (SatBit::Lit(x), SatBit::Lit(y)) => (x, y),
        };
        // Data inputs that are the select itself or its complement reduce to
        // a two-input gate (and the general encoding below would need
        // tautological or repeated-literal clauses).
        if x == sl {
            return self.or(d, s, b);
        }
        if x == -sl {
            return self.and(d, &self.not(s), b);
        }
        if y == sl {
            return self.and(d, s, a);
        }
        if y == -sl {
            return self.or(d, &self.not(s), a);
        }
        if x == -y {
            // `s ? x : !x` is `!(s ^ x)`.
            return Ok(self.not(&self.xor(d, s, a)?));
        }
        // Normalise a negative select by swapping the data inputs.
        let (sl, x, y) = if sl < 0 { (-sl, y, x) } else { (sl, x, y) };
        let key = (OP_MUX, sl, x, y);
        if let Some(g) = self.cached(key) {
            return Ok(SatBit::Lit(g));
        }
        let g = self.fresh(key);
        self.clause(d, &[-sl, -x, g])?;
        self.clause(d, &[-sl, x, -g])?;
        self.clause(d, &[sl, -y, g])?;
        self.clause(d, &[sl, y, -g])?;
        // Redundant, but they let unit propagation see through equal inputs.
        self.clause(d, &[-x, -y, g])?;
        self.clause(d, &[x, y, -g])?;
        Ok(SatBit::Lit(g))
    }
}

/// The CNF literal for 0-based variable `idx` taking value `val`.
fn lit_of(idx: usize, val: bool) -> i32 {
    let v = i32::try_from(idx + 1).expect("CNF variable count is below the clause limit");
    if val { v } else { -v }
}

// ---------- the solver ----------

/// A constraint problem as a CNF, ready to be searched and sampled.
pub struct SatSolver {
    /// The declared width of each random variable.
    widths: Vec<usize>,
    /// First 0-based CNF variable of each random variable's bits (bit `i` of
    /// variable `v` is CNF variable `offsets[v] + i`).
    offsets: Vec<usize>,
    /// The clauses: a flat literal pool and the end of each clause in it.
    lits: Vec<i32>,
    ends: Vec<usize>,
    /// Number of CNF variables (inputs first, then gate outputs).
    num_vars: usize,
    /// Whether the random-variable bit (a CNF variable) occurs in any clause;
    /// a bit that does not is unconstrained and is drawn without a solve.
    occurs: Vec<bool>,
    /// A conjunct blasted to the constant false: no assignment is legal.
    contradiction: bool,
    states: Vec<Option<BigUint>>,
    /// Kept to check every returned model with the IR's own evaluator.
    conjuncts: Vec<Expr>,
    limits: SatLimits,
}

impl fmt::Debug for SatSolver {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SatSolver")
            .field("vars", &self.widths.len())
            .field("cnf_vars", &self.num_vars)
            .field("clauses", &self.ends.len())
            .field("contradiction", &self.contradiction)
            .field("limits", &self.limits)
            .finish()
    }
}

impl SatSolver {
    /// Blasts the conjunction of `conjuncts` over variables of types
    /// `var_tys` (state variables fixed to `states`, as [`super::bdd::compile`])
    /// with the default [`SatLimits`].
    pub fn new(
        var_tys: &[Ty],
        states: &[Option<BigUint>],
        conjuncts: &[Expr],
    ) -> Result<SatSolver, SolveError> {
        SatSolver::with_limits(var_tys, states, conjuncts, SatLimits::default())
    }

    /// As [`SatSolver::new`] with explicit limits.
    pub fn with_limits(
        var_tys: &[Ty],
        states: &[Option<BigUint>],
        conjuncts: &[Expr],
        limits: SatLimits,
    ) -> Result<SatSolver, SolveError> {
        let widths: Vec<usize> = var_tys.iter().map(|t| t.width as usize).collect();
        let mut offsets = Vec::with_capacity(widths.len());
        let mut total = 0usize;
        for w in &widths {
            offsets.push(total);
            total = total.saturating_add(*w);
        }
        if total > limits.clause_limit {
            return Err(SolveError::Build(BddBuildError::CnfTooLarge {
                limit: limits.clause_limit,
                during: "random variable bits",
            }));
        }
        let gates = SatGates {
            cnf: RefCell::new(Cnf {
                lits: Vec::new(),
                ends: Vec::new(),
                last_var: i32::try_from(total).expect("below the clause limit"),
                cache: HashMap::new(),
            }),
            offsets: &offsets,
            limit: limits.clause_limit,
        };
        let mut contradiction = false;
        {
            let blaster = Blaster {
                g: &gates,
                widths: &widths,
                states,
            };
            for c in conjuncts {
                c.validate().map_err(BddBuildError::from)?;
                match blaster.conjunct(c)? {
                    SatBit::Const(true) => {}
                    SatBit::Const(false) => contradiction = true,
                    SatBit::Lit(l) => gates.clause("constraint conjunction", &[l])?,
                }
            }
        }
        let Cnf {
            lits,
            ends,
            last_var,
            ..
        } = gates.cnf.into_inner();
        let mut occurs = vec![false; total];
        for l in &lits {
            let idx = l.unsigned_abs() as usize - 1;
            if idx < total {
                occurs[idx] = true;
            }
        }
        Ok(SatSolver {
            widths,
            offsets,
            lits,
            ends,
            num_vars: usize::try_from(last_var).expect("a count"),
            occurs,
            contradiction,
            states: states.to_vec(),
            conjuncts: conjuncts.to_vec(),
            limits,
        })
    }

    /// The number of clauses of the CNF.
    pub fn num_clauses(&self) -> usize {
        self.ends.len()
    }

    /// Whether the CNF is trivially contradictory (a conjunct folded to false).
    pub fn is_contradiction(&self) -> bool {
        self.contradiction
    }

    /// Whether the assignment `values` (one per variable) is a model, decided
    /// by the CNF and the solver (so it tests the encoding, unlike
    /// [`eval`](super::ir::eval)).
    pub fn accepts(&self, values: &[BigUint]) -> Result<bool, SolveError> {
        assert_eq!(values.len(), self.widths.len(), "one value per variable");
        let mut units = Vec::new();
        for (v, w) in self.widths.iter().enumerate() {
            for i in 0..*w {
                units.push(lit_of(self.offsets[v] + i, values[v].bit(i as u64)));
            }
        }
        let deadline = Instant::now() + self.limits.timeout;
        Ok(self.run(&units, &[], &[], deadline)?.is_some())
    }

    /// Draws one legal assignment, fixing variables in `order`, then any
    /// variable `order` omits in id order. Approximately uniform (module docs).
    pub fn sample(&self, order: &[VarId], seed: &mut u64) -> Result<Vec<BigUint>, SolveError> {
        self.sample_groups(&[order.to_vec()], seed)
    }

    /// As [`super::solve::Solver::sample_groups`].
    pub fn sample_groups(
        &self,
        groups: &[Vec<VarId>],
        seed: &mut u64,
    ) -> Result<Vec<BigUint>, SolveError> {
        let plan = Plan {
            groups: groups
                .iter()
                .map(|g| Group {
                    vars: g.clone(),
                    weighted: Vec::new(),
                })
                .collect(),
            randc: Vec::new(),
        };
        self.sample_plan(&plan, &mut [], seed)
    }

    /// Draws one assignment following `plan`, as
    /// [`super::solve::Solver::sample_plan`] but with the SAT sampler's
    /// approximate uniformity and limits (module docs). `randc_states` is
    /// advanced only for variables that were drawn.
    pub fn sample_plan(
        &self,
        plan: &Plan,
        randc_states: &mut [RandcState],
        seed: &mut u64,
    ) -> Result<Vec<BigUint>, SolveError> {
        let full = resolve_plan(&self.widths, plan, randc_states.len())?;
        self.check_supported(plan)?;
        let deadline = Instant::now() + self.limits.timeout;
        // An unsatisfiable problem consumes no draws (LRM 18.6.3).
        let mut walk = Walk::start(self, deadline)?;
        let mut out = vec![BigUint::zero(); self.widths.len()];

        // The `randc` variables come first (LRM 18.4.2).
        for (k, &v) in plan.randc.iter().enumerate() {
            let v = v as usize;
            let fp = self.fingerprint() ^ (v as u64 + 1).wrapping_mul(0x9e37_79b9_7f4a_7c15);
            let failure: RefCell<Option<SolveError>> = RefCell::new(None);
            let stash = |e: SolveError| {
                failure.borrow_mut().get_or_insert(e);
            };
            let mut legal = || -> Vec<u32> {
                self.enumerate(&[], &[], v, deadline).unwrap_or_else(|e| {
                    stash(e);
                    Vec::new()
                })
            };
            let mut feasible = |x: u32| -> bool {
                walk.agrees_or_probe(&self.value_units(v, x))
                    .unwrap_or_else(|e| {
                        stash(e);
                        false
                    })
            };
            let picked = randc::pick(&mut randc_states[k], fp, &mut legal, &mut feasible, seed);
            if let Some(e) = failure.into_inner() {
                return Err(e);
            }
            let picked = picked.ok_or(SolveError::Unsatisfiable)?;
            walk.commit(&self.value_units(v, picked))?;
            for i in 0..self.widths[v] {
                out[v].set_bit(i as u64, (picked >> i) & 1 == 1);
            }
        }

        for (members, weighted) in &full {
            let mut drawn: Vec<usize> = Vec::new();
            for (k, wv) in weighted.iter().enumerate() {
                // The variables a model of this group is counted over: the
                // group's members and the selectors still to be drawn.
                let mut proj: Vec<usize> = members.clone();
                proj.extend(weighted[k + 1..].iter().map(|w| w.var));
                self.draw_selector(&mut walk, wv, &proj, &mut out, &mut drawn, seed)?;
            }
            // The group's remaining bits, in a random order (module docs).
            let mut bits: Vec<(usize, usize)> = members
                .iter()
                .filter(|v| !drawn.contains(v))
                .flat_map(|&v| (0..self.widths[v]).map(move |i| (v, i)))
                .collect();
            for i in (1..bits.len()).rev() {
                let j = rng::below(seed, &BigUint::from(i as u64 + 1));
                let j = usize::try_from(j).expect("below a usize");
                bits.swap(i, j);
            }
            // One coin per bit, in that order. A bit no clause mentions is
            // legal either way and takes its coin; the others are decided
            // together by `Walk::decide`.
            let mut idxs = Vec::with_capacity(bits.len());
            let mut coins = Vec::with_capacity(bits.len());
            for &(v, i) in &bits {
                let coin = rng::draw(seed) >> 63 == 1;
                let idx = self.offsets[v] + i;
                if self.occurs[idx] {
                    idxs.push(idx);
                    coins.push(coin);
                } else {
                    out[v].set_bit(i as u64, coin);
                }
            }
            let decided = walk.decide(&idxs, &coins)?;
            for (&idx, &bit) in idxs.iter().zip(&decided) {
                let v = self.offsets.partition_point(|&o| o <= idx) - 1;
                out[v].set_bit((idx - self.offsets[v]) as u64, bit);
            }
        }

        // Every sample is checked against the IR's own evaluator.
        let asg = Assignment {
            vars: out.clone(),
            states: self.states.clone(),
        };
        for c in &self.conjuncts {
            match eval(c, &asg) {
                Ok(v) if !v.is_zero() => {}
                _ => return Err(SolveError::SatModelRejected),
            }
        }
        Ok(out)
    }

    /// Refuses what the SAT sampler cannot do faithfully (module docs).
    fn check_supported(&self, plan: &Plan) -> Result<(), SolveError> {
        for &v in &plan.randc {
            if self.widths[v as usize] > MAX_ENUM_BITS {
                return Err(SolveError::SatUnsupported {
                    what: "a `randc` variable wider than 10 bits",
                });
            }
        }
        for g in &plan.groups {
            for w in &g.weighted {
                match w.measure {
                    Measure::Flat | Measure::Models => {}
                    Measure::Distinct(s) => {
                        if self.widths[s as usize] > 32 {
                            return Err(SolveError::SatUnsupported {
                                what: "a `dist` over a variable wider than 32 bits",
                            });
                        }
                    }
                }
            }
        }
        Ok(())
    }

    /// Draws the code of the weighted selector `wv` (and, for a
    /// [`Measure::Distinct`] selector, its subject) and fixes them.
    fn draw_selector(
        &self,
        walk: &mut Walk<'_>,
        wv: &ResolvedWeighted<'_>,
        proj: &[usize],
        out: &mut [BigUint],
        drawn: &mut Vec<usize>,
        seed: &mut u64,
    ) -> Result<(), SolveError> {
        let v = wv.var;
        // The feasible values of the subject under each feasible code.
        let mut sets: HashMap<usize, Vec<u32>> = HashMap::new();
        let mut feasible: Vec<(usize, BigUint, &BigUint)> = Vec::new();
        for (c, wt) in wv.weights.iter().enumerate() {
            if wt.num.is_zero() {
                continue;
            }
            let units = self.value_units(v, c as u32);
            match wv.measure {
                Measure::Distinct(s) => {
                    let values = self.enumerate(&walk.fixed, &units, s as usize, walk.deadline)?;
                    if values.is_empty() {
                        continue;
                    }
                    feasible.push((c, &wt.num * BigUint::from(values.len()), &wt.den));
                    sets.insert(c, values);
                }
                Measure::Flat => {
                    if walk.agrees_or_probe(&units)? {
                        feasible.push((c, wt.num.clone(), &wt.den));
                    }
                }
                Measure::Models => {
                    let n = self.count_models(&walk.fixed, &units, proj, walk.deadline)?;
                    if n > 0 {
                        feasible.push((c, &wt.num * BigUint::from(n), &wt.den));
                    }
                }
            }
        }
        let code = draw_weighted(&feasible, seed)?;
        walk.commit(&self.value_units(v, code as u32))?;
        for i in 0..self.widths[v] {
            out[v].set_bit(i as u64, (code >> i) & 1 == 1);
        }
        if let Measure::Distinct(s) = wv.measure {
            let s = s as usize;
            let values = &sets[&code];
            let pick = rng::below(seed, &BigUint::from(values.len()));
            let x = values[usize::try_from(pick).expect("below a usize")];
            walk.commit(&self.value_units(s, x))?;
            for i in 0..self.widths[s] {
                out[s].set_bit(i as u64, (x >> i) & 1 == 1);
            }
            drawn.push(s);
        }
        Ok(())
    }

    /// The number of distinct assignments of the variables `proj` that extend
    /// to a model satisfying `fixed` and `extra`: the weight of a `dist` over
    /// an expression ([`Measure::Models`], LRM 18.5.4), which the exact path
    /// reads off its BDD. Counted by enumeration with blocking clauses, so it
    /// is exact up to [`MAX_COUNT_MODELS`] and refused by name beyond that
    /// (no approximate count is substituted: a wrong ratio would be a silent
    /// wrong distribution).
    fn count_models(
        &self,
        fixed: &[i32],
        extra: &[i32],
        proj: &[usize],
        deadline: Instant,
    ) -> Result<usize, SolveError> {
        let mut blocks: Vec<Vec<i32>> = Vec::new();
        while let Some(model) = self.run(fixed, extra, &blocks, deadline)? {
            if blocks.len() == MAX_COUNT_MODELS {
                return Err(SolveError::SatUnsupported {
                    what: "a `dist` over an expression with more than 512 feasible \
                           assignments of the variables it weights (the SAT fallback \
                           counts them by enumeration)",
                });
            }
            let block: Vec<i32> = proj
                .iter()
                .flat_map(|&v| (0..self.widths[v]).map(move |i| (v, i)))
                .map(|(v, i)| {
                    let idx = self.offsets[v] + i;
                    lit_of(idx, !model[idx])
                })
                .collect();
            if block.is_empty() {
                return Ok(1);
            }
            blocks.push(block);
        }
        Ok(blocks.len())
    }

    /// The unit literals that give variable `v` the value `x`.
    fn value_units(&self, v: usize, x: u32) -> Vec<i32> {
        (0..self.widths[v])
            .map(|i| lit_of(self.offsets[v] + i, (x >> i) & 1 == 1))
            .collect()
    }

    /// Every value variable `v` takes in some model that also satisfies the
    /// unit literals `fixed` and `extra`: one solve per value, each excluding
    /// the values found so far.
    fn enumerate(
        &self,
        fixed: &[i32],
        extra: &[i32],
        v: usize,
        deadline: Instant,
    ) -> Result<Vec<u32>, SolveError> {
        let mut values = Vec::new();
        let mut blocks: Vec<Vec<i32>> = Vec::new();
        while let Some(model) = self.run(fixed, extra, &blocks, deadline)? {
            if values.len() == 1 << MAX_ENUM_BITS {
                return Err(SolveError::SatUnsupported {
                    what: "a `dist` whose subject has more than 1024 feasible values under \
                           one weighted range (the SAT fallback counts them by enumeration)",
                });
            }
            let mut x = 0u32;
            let mut block = Vec::with_capacity(self.widths[v]);
            for i in 0..self.widths[v] {
                let idx = self.offsets[v] + i;
                x |= u32::from(model[idx]) << i;
                block.push(lit_of(idx, !model[idx]));
            }
            values.push(x);
            blocks.push(block);
        }
        Ok(values)
    }

    /// A fingerprint of the whole CNF (the `randc` cycle's identity: module
    /// docs on when it restarts).
    fn fingerprint(&self) -> u64 {
        let mut h = fingerprint_of(&(self.ends.len() as u64).to_le_bytes());
        for chunk in self.lits.chunks(4096) {
            let mut bytes = Vec::with_capacity(8 + chunk.len() * 4);
            bytes.extend_from_slice(&h.to_le_bytes());
            for l in chunk {
                bytes.extend_from_slice(&l.to_le_bytes());
            }
            h = fingerprint_of(&bytes);
        }
        for chunk in self.ends.chunks(4096) {
            let mut bytes = Vec::with_capacity(8 + chunk.len() * 8);
            bytes.extend_from_slice(&h.to_le_bytes());
            for e in chunk {
                bytes.extend_from_slice(&(*e as u64).to_le_bytes());
            }
            h = fingerprint_of(&bytes);
        }
        h
    }

    /// One solve of the CNF plus the unit literals `fixed` and `extra` plus
    /// the clauses `blocks`: a model (one bool per CNF variable) or `None`
    /// when there is none. `splr` has no assumptions in its default build, so
    /// each call is a fresh solver; the clause list is shared, not copied.
    fn run(
        &self,
        fixed: &[i32],
        extra: &[i32],
        blocks: &[Vec<i32>],
        deadline: Instant,
    ) -> Result<Option<Vec<bool>>, SolveError> {
        if self.contradiction {
            return Ok(None);
        }
        let remaining = deadline
            .checked_duration_since(Instant::now())
            .filter(|d| !d.is_zero())
            .ok_or(SolveError::SatTimeout)?;
        let mut clauses: Vec<&[i32]> =
            Vec::with_capacity(self.ends.len() + fixed.len() + extra.len() + blocks.len());
        let mut start = 0;
        for &end in &self.ends {
            clauses.push(&self.lits[start..end]);
            start = end;
        }
        for l in fixed.iter().chain(extra) {
            clauses.push(std::slice::from_ref(l));
        }
        for b in blocks {
            clauses.push(b.as_slice());
        }
        if clauses.is_empty() {
            return Ok(Some(vec![false; self.num_vars]));
        }
        let config = Config {
            quiet_mode: true,
            c_timeout: remaining.as_secs_f64().ceil().max(1.0),
            ..Config::default()
        };
        let answer = match Cdcl::try_from((config, clauses.as_slice())) {
            Ok(mut solver) => solver.solve(),
            Err(Ok(cert)) => Ok(cert),
            Err(Err(SolverError::EmptyClause)) => Ok(Certificate::UNSAT),
            Err(Err(e)) => Err(e),
        };
        match answer {
            Ok(Certificate::SAT(vals)) => {
                let mut model = vec![false; self.num_vars];
                for l in vals {
                    let idx = l.unsigned_abs() as usize - 1;
                    if idx < model.len() {
                        model[idx] = l > 0;
                    }
                }
                Ok(Some(model))
            }
            Ok(Certificate::UNSAT)
            | Err(
                SolverError::EmptyClause
                | SolverError::Inconsistent
                | SolverError::RootLevelConflict(_),
            ) => Ok(None),
            Err(SolverError::TimeOut) => Err(SolveError::SatTimeout),
            Err(SolverError::OutOfMemory) => Err(SolveError::SatSolverFailure {
                what: "the solver ran out of memory",
            }),
            Err(SolverError::SolverBug) => Err(SolveError::SatSolverFailure {
                what: "the solver reported an internal error",
            }),
            Err(_) => Err(SolveError::SatSolverFailure {
                what: "the solver rejected the CNF",
            }),
        }
    }
}

/// The state of one sampling call: the unit literals fixed so far and a model
/// that satisfies them (the witness that lets most decisions skip a solve).
struct Walk<'a> {
    s: &'a SatSolver,
    deadline: Instant,
    fixed: Vec<i32>,
    model: Vec<bool>,
}

impl<'a> Walk<'a> {
    /// The first model, or [`SolveError::Unsatisfiable`].
    fn start(s: &'a SatSolver, deadline: Instant) -> Result<Walk<'a>, SolveError> {
        match s.run(&[], &[], &[], deadline)? {
            Some(model) => Ok(Walk {
                s,
                deadline,
                fixed: Vec::new(),
                model,
            }),
            None => Err(SolveError::Unsatisfiable),
        }
    }

    /// Whether the current model already satisfies every literal of `units`.
    fn agrees(&self, units: &[i32]) -> bool {
        units
            .iter()
            .all(|&l| self.model[l.unsigned_abs() as usize - 1] == (l > 0))
    }

    /// Whether some model satisfies the fixed literals and `units`; on yes the
    /// current model becomes that one.
    fn probe(&mut self, units: &[i32]) -> Result<bool, SolveError> {
        match self.s.run(&self.fixed, units, &[], self.deadline)? {
            Some(model) => {
                self.model = model;
                Ok(true)
            }
            None => Ok(false),
        }
    }

    fn agrees_or_probe(&mut self, units: &[i32]) -> Result<bool, SolveError> {
        Ok(self.agrees(units) || self.probe(units)?)
    }

    /// Fixes `units`, which must be feasible (a bit or code just checked).
    fn commit(&mut self, units: &[i32]) -> Result<(), SolveError> {
        if !self.agrees(units) && !self.probe(units)? {
            return Err(SolveError::SatSolverFailure {
                what: "a value found feasible became infeasible",
            });
        }
        self.fixed.extend_from_slice(units);
        Ok(())
    }

    /// Decides CNF variables `idxs` in order: each takes its coin from
    /// `coins` unless that leaves no model given the earlier decisions, in
    /// which case it takes the other value; all are fixed. The result is
    /// exactly what deciding one bit at a time gives, but a run of coins that
    /// is feasible together costs one solve (or none, when the current model
    /// already agrees) instead of one per bit: a prefix of a feasible run is
    /// feasible, so a feasible whole decides every bit of it the same way, and
    /// an infeasible run is split in two and each half decided in turn.
    fn decide(&mut self, idxs: &[usize], coins: &[bool]) -> Result<Vec<bool>, SolveError> {
        let mut out = Vec::with_capacity(idxs.len());
        self.decide_run(idxs, coins, &mut out)?;
        Ok(out)
    }

    fn decide_run(
        &mut self,
        idxs: &[usize],
        coins: &[bool],
        out: &mut Vec<bool>,
    ) -> Result<(), SolveError> {
        if idxs.is_empty() {
            return Ok(());
        }
        let units: Vec<i32> = idxs
            .iter()
            .zip(coins)
            .map(|(&i, &c)| lit_of(i, c))
            .collect();
        if self.agrees_or_probe(&units)? {
            self.fixed.extend_from_slice(&units);
            out.extend_from_slice(coins);
            return Ok(());
        }
        if let [idx] = idxs {
            // The coin has no model, so the other value must.
            let flipped = [lit_of(*idx, !coins[0])];
            self.commit(&flipped)?;
            out.push(!coins[0]);
            return Ok(());
        }
        let mid = idxs.len() / 2;
        self.decide_run(&idxs[..mid], &coins[..mid], out)?;
        self.decide_run(&idxs[mid..], &coins[mid..], out)
    }
}