Skip to main content

vj_sim/constraints/
sat.rs

1//! The SAT fallback for problems whose BDD exceeds its node budget (plan
2//! section 3.4, ledger XC-063).
3//!
4//! Some constraint sets have no small BDD in any variable order (a wide
5//! multiplier is the canonical one, [`super::bdd`]). For those, [`SatSolver`]
6//! blasts the same IR with the same gadgets ([`super::blast`]) into a CNF with
7//! a Tseitin encoding and searches it with `splr`, a pure-Rust CDCL solver.
8//! [`super::solve::Engine`] picks between the two: the exact BDD solver when
9//! it fits, this one when it does not.
10//!
11//! # What is the same and what is not
12//!
13//! * The set of legal assignments is identical: both back ends run the same
14//!   gadgets, so division by zero is "not legal" here too (see
15//!   [`super::blast`]). Every sample is also checked with the IR's own
16//!   evaluator ([`super::ir::eval`]) before it is returned; a model the
17//!   evaluator rejects is the typed [`SolveError::SatModelRejected`], never a
18//!   silent wrong answer.
19//! * The distribution is NOT exactly uniform. A CDCL solver finds *a* model,
20//!   not a uniformly random one, and a model count is not available. The
21//!   sampler fixes the bits of each solve-order group one at a time in a
22//!   random order, choosing each bit by a fair coin and falling back to the
23//!   other value only when the coin's choice leaves no model (a feasibility
24//!   solve, skipped when the last model already agrees). Every legal
25//!   assignment is reachable and every draw is legal, but the probabilities
26//!   are only approximately uniform. The exact path keeps its exactness; this
27//!   is the price of a problem the exact path cannot hold.
28//! * `solve ... before` groups, `randc` variables, and `dist` selectors follow
29//!   the exact sampler's structure ([`super::solve::Solver::sample_plan`]): the
30//!   `randc` variables first, then each group with its weighted selectors
31//!   before its uniform variables, each later choice constrained by the
32//!   earlier ones. What it cannot do it refuses, loudly:
33//!   [`SolveError::SatUnsupported`] for a `dist` over an expression with more
34//!   than [`MAX_COUNT_MODELS`] feasible assignments (fewer are counted by
35//!   enumeration), a `dist` whose subject is wider than 32 bits or has more
36//!   than `1 << MAX_ENUM_BITS` feasible values under one weighted range, and a
37//!   `randc` variable wider than [`MAX_ENUM_BITS`] bits (LRM 18.4.2 lets an
38//!   implementation limit `randc` size; all of those are enumerated value by
39//!   value).
40//! * `randc`: the cycle restarts whenever the CNF changes at all, which is
41//!   more often than the exact path's "the variable's own legal set changed"
42//!   (a state value read by an unrelated constraint restarts the cycle).
43//!
44//! # Limits
45//!
46//! The CNF is capped at [`SatLimits::clause_limit`] clauses (a typed
47//! [`BddBuildError::CnfTooLarge`], reported as the `TooLarge` outcome) and a
48//! sampling call at [`SatLimits::timeout`] of search (a typed
49//! [`SolveError::SatTimeout`]). `splr` takes its timeout in whole seconds and
50//! checks it between restarts, so a call can overrun by a fraction of a second.
51
52use super::bdd::BddBuildError;
53use super::blast::{Blaster, Gates};
54use super::ir::{Assignment, Expr, Ty, VarId, eval};
55use super::randc::{self, RandcState, fingerprint_of};
56use super::rng;
57use super::solve::{
58    Group, Measure, Plan, ResolvedWeighted, SolveError, draw_weighted, resolve_plan,
59};
60use core::fmt;
61use num_bigint::BigUint;
62use num_traits::Zero;
63use splr::{Certificate, Config, SolveIF, Solver as Cdcl, SolverError};
64use std::cell::RefCell;
65use std::collections::HashMap;
66use std::time::{Duration, Instant};
67
68/// The most bits of one variable the sampler enumerates value by value: a
69/// `randc` variable's legal set, a `dist` subject's feasible values.
70pub const MAX_ENUM_BITS: usize = 10;
71
72/// The most distinct assignments [`SatSolver::count_models`] enumerates before
73/// it refuses.
74pub const MAX_COUNT_MODELS: usize = 512;
75
76/// The default cap on CNF clauses (about 100 bytes each with the solver's own
77/// copy; a 64-bit multiplier is about 100,000).
78pub const CLAUSE_LIMIT: usize = 4_000_000;
79
80/// The default search time allowed to one sampling call.
81pub const TIMEOUT: Duration = Duration::from_secs(30);
82
83/// Resource limits of the SAT fallback.
84#[derive(Clone, Copy, Debug, PartialEq, Eq)]
85pub struct SatLimits {
86    /// The most clauses the CNF may hold.
87    pub clause_limit: usize,
88    /// The search time one sampling call may use.
89    pub timeout: Duration,
90}
91
92impl Default for SatLimits {
93    fn default() -> Self {
94        SatLimits {
95            clause_limit: CLAUSE_LIMIT,
96            timeout: TIMEOUT,
97        }
98    }
99}
100
101// ---------- the CNF back end of the shared gadgets ----------
102
103/// A bit of the blasted formula: a constant or a (signed, 1-based) CNF literal.
104#[derive(Clone, Copy, Debug, PartialEq, Eq)]
105pub(crate) enum SatBit {
106    Const(bool),
107    Lit(i32),
108}
109
110/// The growing CNF: a flat literal pool with clause end offsets, and a
111/// structural cache so a gate over the same inputs is encoded once.
112struct Cnf {
113    lits: Vec<i32>,
114    ends: Vec<usize>,
115    /// The last CNF variable allocated.
116    last_var: i32,
117    /// `(operation, a, b, c)` to the gate's output variable.
118    cache: HashMap<(u8, i32, i32, i32), i32>,
119}
120
121/// The Tseitin encoder implementing [`Gates`].
122struct SatGates<'a> {
123    cnf: RefCell<Cnf>,
124    /// First (0-based) CNF variable of each random variable's bits.
125    offsets: &'a [usize],
126    limit: usize,
127}
128
129const OP_AND: u8 = 0;
130const OP_XOR: u8 = 1;
131const OP_MUX: u8 = 2;
132
133impl SatGates<'_> {
134    fn clause(&self, during: &'static str, c: &[i32]) -> Result<(), BddBuildError> {
135        let mut cnf = self.cnf.borrow_mut();
136        if cnf.ends.len() >= self.limit {
137            return Err(BddBuildError::CnfTooLarge {
138                limit: self.limit,
139                during,
140            });
141        }
142        cnf.lits.extend_from_slice(c);
143        let end = cnf.lits.len();
144        cnf.ends.push(end);
145        Ok(())
146    }
147
148    fn cached(&self, key: (u8, i32, i32, i32)) -> Option<i32> {
149        self.cnf.borrow().cache.get(&key).copied()
150    }
151
152    fn fresh(&self, key: (u8, i32, i32, i32)) -> i32 {
153        let mut cnf = self.cnf.borrow_mut();
154        cnf.last_var += 1;
155        let g = cnf.last_var;
156        cnf.cache.insert(key, g);
157        g
158    }
159}
160
161impl Gates for SatGates<'_> {
162    type Bit = SatBit;
163
164    fn constant(&self, v: bool) -> SatBit {
165        SatBit::Const(v)
166    }
167
168    fn as_const(&self, a: &SatBit) -> Option<bool> {
169        match a {
170            SatBit::Const(v) => Some(*v),
171            SatBit::Lit(_) => None,
172        }
173    }
174
175    fn input(&self, var: VarId, i: usize) -> SatBit {
176        SatBit::Lit(lit_of(self.offsets[var as usize] + i, true))
177    }
178
179    fn not(&self, a: &SatBit) -> SatBit {
180        match a {
181            SatBit::Const(v) => SatBit::Const(!v),
182            SatBit::Lit(l) => SatBit::Lit(-l),
183        }
184    }
185
186    fn and(&self, d: &'static str, a: &SatBit, b: &SatBit) -> Result<SatBit, BddBuildError> {
187        let (x, y) = match (*a, *b) {
188            (SatBit::Const(false), _) | (_, SatBit::Const(false)) => {
189                return Ok(SatBit::Const(false));
190            }
191            (SatBit::Const(true), v) | (v, SatBit::Const(true)) => return Ok(v),
192            (SatBit::Lit(x), SatBit::Lit(y)) => (x, y),
193        };
194        if x == y {
195            return Ok(SatBit::Lit(x));
196        }
197        if x == -y {
198            return Ok(SatBit::Const(false));
199        }
200        let key = (OP_AND, x.min(y), x.max(y), 0);
201        if let Some(g) = self.cached(key) {
202            return Ok(SatBit::Lit(g));
203        }
204        let g = self.fresh(key);
205        self.clause(d, &[-g, x])?;
206        self.clause(d, &[-g, y])?;
207        self.clause(d, &[g, -x, -y])?;
208        Ok(SatBit::Lit(g))
209    }
210
211    fn or(&self, d: &'static str, a: &SatBit, b: &SatBit) -> Result<SatBit, BddBuildError> {
212        // De Morgan over the shared AND gate.
213        let both_off = self.and(d, &self.not(a), &self.not(b))?;
214        Ok(self.not(&both_off))
215    }
216
217    fn xor(&self, d: &'static str, a: &SatBit, b: &SatBit) -> Result<SatBit, BddBuildError> {
218        let (x, y) = match (*a, *b) {
219            (SatBit::Const(p), SatBit::Const(q)) => return Ok(SatBit::Const(p ^ q)),
220            (SatBit::Const(false), v) | (v, SatBit::Const(false)) => return Ok(v),
221            (SatBit::Const(true), v) | (v, SatBit::Const(true)) => return Ok(self.not(&v)),
222            (SatBit::Lit(x), SatBit::Lit(y)) => (x, y),
223        };
224        if x == y {
225            return Ok(SatBit::Const(false));
226        }
227        if x == -y {
228            return Ok(SatBit::Const(true));
229        }
230        // `x ^ y` over signs is `|x| ^ |y|` flipped when exactly one is negative.
231        let flip = (x < 0) != (y < 0);
232        let (p, q) = (x.abs(), y.abs());
233        let key = (OP_XOR, p.min(q), p.max(q), 0);
234        let g = match self.cached(key) {
235            Some(g) => g,
236            None => {
237                let g = self.fresh(key);
238                self.clause(d, &[-g, p, q])?;
239                self.clause(d, &[-g, -p, -q])?;
240                self.clause(d, &[g, -p, q])?;
241                self.clause(d, &[g, p, -q])?;
242                g
243            }
244        };
245        Ok(SatBit::Lit(if flip { -g } else { g }))
246    }
247
248    fn mux(
249        &self,
250        d: &'static str,
251        s: &SatBit,
252        a: &SatBit,
253        b: &SatBit,
254    ) -> Result<SatBit, BddBuildError> {
255        let sl = match *s {
256            SatBit::Const(true) => return Ok(*a),
257            SatBit::Const(false) => return Ok(*b),
258            SatBit::Lit(l) => l,
259        };
260        if a == b {
261            return Ok(*a);
262        }
263        let (x, y) = match (*a, *b) {
264            (SatBit::Const(true), _) => return self.or(d, s, b),
265            (SatBit::Const(false), _) => return self.and(d, &self.not(s), b),
266            (_, SatBit::Const(true)) => return self.or(d, &self.not(s), a),
267            (_, SatBit::Const(false)) => return self.and(d, s, a),
268            (SatBit::Lit(x), SatBit::Lit(y)) => (x, y),
269        };
270        // Data inputs that are the select itself or its complement reduce to
271        // a two-input gate (and the general encoding below would need
272        // tautological or repeated-literal clauses).
273        if x == sl {
274            return self.or(d, s, b);
275        }
276        if x == -sl {
277            return self.and(d, &self.not(s), b);
278        }
279        if y == sl {
280            return self.and(d, s, a);
281        }
282        if y == -sl {
283            return self.or(d, &self.not(s), a);
284        }
285        if x == -y {
286            // `s ? x : !x` is `!(s ^ x)`.
287            return Ok(self.not(&self.xor(d, s, a)?));
288        }
289        // Normalise a negative select by swapping the data inputs.
290        let (sl, x, y) = if sl < 0 { (-sl, y, x) } else { (sl, x, y) };
291        let key = (OP_MUX, sl, x, y);
292        if let Some(g) = self.cached(key) {
293            return Ok(SatBit::Lit(g));
294        }
295        let g = self.fresh(key);
296        self.clause(d, &[-sl, -x, g])?;
297        self.clause(d, &[-sl, x, -g])?;
298        self.clause(d, &[sl, -y, g])?;
299        self.clause(d, &[sl, y, -g])?;
300        // Redundant, but they let unit propagation see through equal inputs.
301        self.clause(d, &[-x, -y, g])?;
302        self.clause(d, &[x, y, -g])?;
303        Ok(SatBit::Lit(g))
304    }
305}
306
307/// The CNF literal for 0-based variable `idx` taking value `val`.
308fn lit_of(idx: usize, val: bool) -> i32 {
309    let v = i32::try_from(idx + 1).expect("CNF variable count is below the clause limit");
310    if val { v } else { -v }
311}
312
313// ---------- the solver ----------
314
315/// A constraint problem as a CNF, ready to be searched and sampled.
316pub struct SatSolver {
317    /// The declared width of each random variable.
318    widths: Vec<usize>,
319    /// First 0-based CNF variable of each random variable's bits (bit `i` of
320    /// variable `v` is CNF variable `offsets[v] + i`).
321    offsets: Vec<usize>,
322    /// The clauses: a flat literal pool and the end of each clause in it.
323    lits: Vec<i32>,
324    ends: Vec<usize>,
325    /// Number of CNF variables (inputs first, then gate outputs).
326    num_vars: usize,
327    /// Whether the random-variable bit (a CNF variable) occurs in any clause;
328    /// a bit that does not is unconstrained and is drawn without a solve.
329    occurs: Vec<bool>,
330    /// A conjunct blasted to the constant false: no assignment is legal.
331    contradiction: bool,
332    states: Vec<Option<BigUint>>,
333    /// Kept to check every returned model with the IR's own evaluator.
334    conjuncts: Vec<Expr>,
335    limits: SatLimits,
336}
337
338impl fmt::Debug for SatSolver {
339    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
340        f.debug_struct("SatSolver")
341            .field("vars", &self.widths.len())
342            .field("cnf_vars", &self.num_vars)
343            .field("clauses", &self.ends.len())
344            .field("contradiction", &self.contradiction)
345            .field("limits", &self.limits)
346            .finish()
347    }
348}
349
350impl SatSolver {
351    /// Blasts the conjunction of `conjuncts` over variables of types
352    /// `var_tys` (state variables fixed to `states`, as [`super::bdd::compile`])
353    /// with the default [`SatLimits`].
354    pub fn new(
355        var_tys: &[Ty],
356        states: &[Option<BigUint>],
357        conjuncts: &[Expr],
358    ) -> Result<SatSolver, SolveError> {
359        SatSolver::with_limits(var_tys, states, conjuncts, SatLimits::default())
360    }
361
362    /// As [`SatSolver::new`] with explicit limits.
363    pub fn with_limits(
364        var_tys: &[Ty],
365        states: &[Option<BigUint>],
366        conjuncts: &[Expr],
367        limits: SatLimits,
368    ) -> Result<SatSolver, SolveError> {
369        let widths: Vec<usize> = var_tys.iter().map(|t| t.width as usize).collect();
370        let mut offsets = Vec::with_capacity(widths.len());
371        let mut total = 0usize;
372        for w in &widths {
373            offsets.push(total);
374            total = total.saturating_add(*w);
375        }
376        if total > limits.clause_limit {
377            return Err(SolveError::Build(BddBuildError::CnfTooLarge {
378                limit: limits.clause_limit,
379                during: "random variable bits",
380            }));
381        }
382        let gates = SatGates {
383            cnf: RefCell::new(Cnf {
384                lits: Vec::new(),
385                ends: Vec::new(),
386                last_var: i32::try_from(total).expect("below the clause limit"),
387                cache: HashMap::new(),
388            }),
389            offsets: &offsets,
390            limit: limits.clause_limit,
391        };
392        let mut contradiction = false;
393        {
394            let blaster = Blaster {
395                g: &gates,
396                widths: &widths,
397                states,
398            };
399            for c in conjuncts {
400                c.validate().map_err(BddBuildError::from)?;
401                match blaster.conjunct(c)? {
402                    SatBit::Const(true) => {}
403                    SatBit::Const(false) => contradiction = true,
404                    SatBit::Lit(l) => gates.clause("constraint conjunction", &[l])?,
405                }
406            }
407        }
408        let Cnf {
409            lits,
410            ends,
411            last_var,
412            ..
413        } = gates.cnf.into_inner();
414        let mut occurs = vec![false; total];
415        for l in &lits {
416            let idx = l.unsigned_abs() as usize - 1;
417            if idx < total {
418                occurs[idx] = true;
419            }
420        }
421        Ok(SatSolver {
422            widths,
423            offsets,
424            lits,
425            ends,
426            num_vars: usize::try_from(last_var).expect("a count"),
427            occurs,
428            contradiction,
429            states: states.to_vec(),
430            conjuncts: conjuncts.to_vec(),
431            limits,
432        })
433    }
434
435    /// The number of clauses of the CNF.
436    pub fn num_clauses(&self) -> usize {
437        self.ends.len()
438    }
439
440    /// Whether the CNF is trivially contradictory (a conjunct folded to false).
441    pub fn is_contradiction(&self) -> bool {
442        self.contradiction
443    }
444
445    /// Whether the assignment `values` (one per variable) is a model, decided
446    /// by the CNF and the solver (so it tests the encoding, unlike
447    /// [`eval`](super::ir::eval)).
448    pub fn accepts(&self, values: &[BigUint]) -> Result<bool, SolveError> {
449        assert_eq!(values.len(), self.widths.len(), "one value per variable");
450        let mut units = Vec::new();
451        for (v, w) in self.widths.iter().enumerate() {
452            for i in 0..*w {
453                units.push(lit_of(self.offsets[v] + i, values[v].bit(i as u64)));
454            }
455        }
456        let deadline = Instant::now() + self.limits.timeout;
457        Ok(self.run(&units, &[], &[], deadline)?.is_some())
458    }
459
460    /// Draws one legal assignment, fixing variables in `order`, then any
461    /// variable `order` omits in id order. Approximately uniform (module docs).
462    pub fn sample(&self, order: &[VarId], seed: &mut u64) -> Result<Vec<BigUint>, SolveError> {
463        self.sample_groups(&[order.to_vec()], seed)
464    }
465
466    /// As [`super::solve::Solver::sample_groups`].
467    pub fn sample_groups(
468        &self,
469        groups: &[Vec<VarId>],
470        seed: &mut u64,
471    ) -> Result<Vec<BigUint>, SolveError> {
472        let plan = Plan {
473            groups: groups
474                .iter()
475                .map(|g| Group {
476                    vars: g.clone(),
477                    weighted: Vec::new(),
478                })
479                .collect(),
480            randc: Vec::new(),
481        };
482        self.sample_plan(&plan, &mut [], seed)
483    }
484
485    /// Draws one assignment following `plan`, as
486    /// [`super::solve::Solver::sample_plan`] but with the SAT sampler's
487    /// approximate uniformity and limits (module docs). `randc_states` is
488    /// advanced only for variables that were drawn.
489    pub fn sample_plan(
490        &self,
491        plan: &Plan,
492        randc_states: &mut [RandcState],
493        seed: &mut u64,
494    ) -> Result<Vec<BigUint>, SolveError> {
495        let full = resolve_plan(&self.widths, plan, randc_states.len())?;
496        self.check_supported(plan)?;
497        let deadline = Instant::now() + self.limits.timeout;
498        // An unsatisfiable problem consumes no draws (LRM 18.6.3).
499        let mut walk = Walk::start(self, deadline)?;
500        let mut out = vec![BigUint::zero(); self.widths.len()];
501
502        // The `randc` variables come first (LRM 18.4.2).
503        for (k, &v) in plan.randc.iter().enumerate() {
504            let v = v as usize;
505            let fp = self.fingerprint() ^ (v as u64 + 1).wrapping_mul(0x9e37_79b9_7f4a_7c15);
506            let failure: RefCell<Option<SolveError>> = RefCell::new(None);
507            let stash = |e: SolveError| {
508                failure.borrow_mut().get_or_insert(e);
509            };
510            let mut legal = || -> Vec<u32> {
511                self.enumerate(&[], &[], v, deadline).unwrap_or_else(|e| {
512                    stash(e);
513                    Vec::new()
514                })
515            };
516            let mut feasible = |x: u32| -> bool {
517                walk.agrees_or_probe(&self.value_units(v, x))
518                    .unwrap_or_else(|e| {
519                        stash(e);
520                        false
521                    })
522            };
523            let picked = randc::pick(&mut randc_states[k], fp, &mut legal, &mut feasible, seed);
524            if let Some(e) = failure.into_inner() {
525                return Err(e);
526            }
527            let picked = picked.ok_or(SolveError::Unsatisfiable)?;
528            walk.commit(&self.value_units(v, picked))?;
529            for i in 0..self.widths[v] {
530                out[v].set_bit(i as u64, (picked >> i) & 1 == 1);
531            }
532        }
533
534        for (members, weighted) in &full {
535            let mut drawn: Vec<usize> = Vec::new();
536            for (k, wv) in weighted.iter().enumerate() {
537                // The variables a model of this group is counted over: the
538                // group's members and the selectors still to be drawn.
539                let mut proj: Vec<usize> = members.clone();
540                proj.extend(weighted[k + 1..].iter().map(|w| w.var));
541                self.draw_selector(&mut walk, wv, &proj, &mut out, &mut drawn, seed)?;
542            }
543            // The group's remaining bits, in a random order (module docs).
544            let mut bits: Vec<(usize, usize)> = members
545                .iter()
546                .filter(|v| !drawn.contains(v))
547                .flat_map(|&v| (0..self.widths[v]).map(move |i| (v, i)))
548                .collect();
549            for i in (1..bits.len()).rev() {
550                let j = rng::below(seed, &BigUint::from(i as u64 + 1));
551                let j = usize::try_from(j).expect("below a usize");
552                bits.swap(i, j);
553            }
554            // One coin per bit, in that order. A bit no clause mentions is
555            // legal either way and takes its coin; the others are decided
556            // together by `Walk::decide`.
557            let mut idxs = Vec::with_capacity(bits.len());
558            let mut coins = Vec::with_capacity(bits.len());
559            for &(v, i) in &bits {
560                let coin = rng::draw(seed) >> 63 == 1;
561                let idx = self.offsets[v] + i;
562                if self.occurs[idx] {
563                    idxs.push(idx);
564                    coins.push(coin);
565                } else {
566                    out[v].set_bit(i as u64, coin);
567                }
568            }
569            let decided = walk.decide(&idxs, &coins)?;
570            for (&idx, &bit) in idxs.iter().zip(&decided) {
571                let v = self.offsets.partition_point(|&o| o <= idx) - 1;
572                out[v].set_bit((idx - self.offsets[v]) as u64, bit);
573            }
574        }
575
576        // Every sample is checked against the IR's own evaluator.
577        let asg = Assignment {
578            vars: out.clone(),
579            states: self.states.clone(),
580        };
581        for c in &self.conjuncts {
582            match eval(c, &asg) {
583                Ok(v) if !v.is_zero() => {}
584                _ => return Err(SolveError::SatModelRejected),
585            }
586        }
587        Ok(out)
588    }
589
590    /// Refuses what the SAT sampler cannot do faithfully (module docs).
591    fn check_supported(&self, plan: &Plan) -> Result<(), SolveError> {
592        for &v in &plan.randc {
593            if self.widths[v as usize] > MAX_ENUM_BITS {
594                return Err(SolveError::SatUnsupported {
595                    what: "a `randc` variable wider than 10 bits",
596                });
597            }
598        }
599        for g in &plan.groups {
600            for w in &g.weighted {
601                match w.measure {
602                    Measure::Flat | Measure::Models => {}
603                    Measure::Distinct(s) => {
604                        if self.widths[s as usize] > 32 {
605                            return Err(SolveError::SatUnsupported {
606                                what: "a `dist` over a variable wider than 32 bits",
607                            });
608                        }
609                    }
610                }
611            }
612        }
613        Ok(())
614    }
615
616    /// Draws the code of the weighted selector `wv` (and, for a
617    /// [`Measure::Distinct`] selector, its subject) and fixes them.
618    fn draw_selector(
619        &self,
620        walk: &mut Walk<'_>,
621        wv: &ResolvedWeighted<'_>,
622        proj: &[usize],
623        out: &mut [BigUint],
624        drawn: &mut Vec<usize>,
625        seed: &mut u64,
626    ) -> Result<(), SolveError> {
627        let v = wv.var;
628        // The feasible values of the subject under each feasible code.
629        let mut sets: HashMap<usize, Vec<u32>> = HashMap::new();
630        let mut feasible: Vec<(usize, BigUint, &BigUint)> = Vec::new();
631        for (c, wt) in wv.weights.iter().enumerate() {
632            if wt.num.is_zero() {
633                continue;
634            }
635            let units = self.value_units(v, c as u32);
636            match wv.measure {
637                Measure::Distinct(s) => {
638                    let values = self.enumerate(&walk.fixed, &units, s as usize, walk.deadline)?;
639                    if values.is_empty() {
640                        continue;
641                    }
642                    feasible.push((c, &wt.num * BigUint::from(values.len()), &wt.den));
643                    sets.insert(c, values);
644                }
645                Measure::Flat => {
646                    if walk.agrees_or_probe(&units)? {
647                        feasible.push((c, wt.num.clone(), &wt.den));
648                    }
649                }
650                Measure::Models => {
651                    let n = self.count_models(&walk.fixed, &units, proj, walk.deadline)?;
652                    if n > 0 {
653                        feasible.push((c, &wt.num * BigUint::from(n), &wt.den));
654                    }
655                }
656            }
657        }
658        let code = draw_weighted(&feasible, seed)?;
659        walk.commit(&self.value_units(v, code as u32))?;
660        for i in 0..self.widths[v] {
661            out[v].set_bit(i as u64, (code >> i) & 1 == 1);
662        }
663        if let Measure::Distinct(s) = wv.measure {
664            let s = s as usize;
665            let values = &sets[&code];
666            let pick = rng::below(seed, &BigUint::from(values.len()));
667            let x = values[usize::try_from(pick).expect("below a usize")];
668            walk.commit(&self.value_units(s, x))?;
669            for i in 0..self.widths[s] {
670                out[s].set_bit(i as u64, (x >> i) & 1 == 1);
671            }
672            drawn.push(s);
673        }
674        Ok(())
675    }
676
677    /// The number of distinct assignments of the variables `proj` that extend
678    /// to a model satisfying `fixed` and `extra`: the weight of a `dist` over
679    /// an expression ([`Measure::Models`], LRM 18.5.4), which the exact path
680    /// reads off its BDD. Counted by enumeration with blocking clauses, so it
681    /// is exact up to [`MAX_COUNT_MODELS`] and refused by name beyond that
682    /// (no approximate count is substituted: a wrong ratio would be a silent
683    /// wrong distribution).
684    fn count_models(
685        &self,
686        fixed: &[i32],
687        extra: &[i32],
688        proj: &[usize],
689        deadline: Instant,
690    ) -> Result<usize, SolveError> {
691        let mut blocks: Vec<Vec<i32>> = Vec::new();
692        while let Some(model) = self.run(fixed, extra, &blocks, deadline)? {
693            if blocks.len() == MAX_COUNT_MODELS {
694                return Err(SolveError::SatUnsupported {
695                    what: "a `dist` over an expression with more than 512 feasible \
696                           assignments of the variables it weights (the SAT fallback \
697                           counts them by enumeration)",
698                });
699            }
700            let block: Vec<i32> = proj
701                .iter()
702                .flat_map(|&v| (0..self.widths[v]).map(move |i| (v, i)))
703                .map(|(v, i)| {
704                    let idx = self.offsets[v] + i;
705                    lit_of(idx, !model[idx])
706                })
707                .collect();
708            if block.is_empty() {
709                return Ok(1);
710            }
711            blocks.push(block);
712        }
713        Ok(blocks.len())
714    }
715
716    /// The unit literals that give variable `v` the value `x`.
717    fn value_units(&self, v: usize, x: u32) -> Vec<i32> {
718        (0..self.widths[v])
719            .map(|i| lit_of(self.offsets[v] + i, (x >> i) & 1 == 1))
720            .collect()
721    }
722
723    /// Every value variable `v` takes in some model that also satisfies the
724    /// unit literals `fixed` and `extra`: one solve per value, each excluding
725    /// the values found so far.
726    fn enumerate(
727        &self,
728        fixed: &[i32],
729        extra: &[i32],
730        v: usize,
731        deadline: Instant,
732    ) -> Result<Vec<u32>, SolveError> {
733        let mut values = Vec::new();
734        let mut blocks: Vec<Vec<i32>> = Vec::new();
735        while let Some(model) = self.run(fixed, extra, &blocks, deadline)? {
736            if values.len() == 1 << MAX_ENUM_BITS {
737                return Err(SolveError::SatUnsupported {
738                    what: "a `dist` whose subject has more than 1024 feasible values under \
739                           one weighted range (the SAT fallback counts them by enumeration)",
740                });
741            }
742            let mut x = 0u32;
743            let mut block = Vec::with_capacity(self.widths[v]);
744            for i in 0..self.widths[v] {
745                let idx = self.offsets[v] + i;
746                x |= u32::from(model[idx]) << i;
747                block.push(lit_of(idx, !model[idx]));
748            }
749            values.push(x);
750            blocks.push(block);
751        }
752        Ok(values)
753    }
754
755    /// A fingerprint of the whole CNF (the `randc` cycle's identity: module
756    /// docs on when it restarts).
757    fn fingerprint(&self) -> u64 {
758        let mut h = fingerprint_of(&(self.ends.len() as u64).to_le_bytes());
759        for chunk in self.lits.chunks(4096) {
760            let mut bytes = Vec::with_capacity(8 + chunk.len() * 4);
761            bytes.extend_from_slice(&h.to_le_bytes());
762            for l in chunk {
763                bytes.extend_from_slice(&l.to_le_bytes());
764            }
765            h = fingerprint_of(&bytes);
766        }
767        for chunk in self.ends.chunks(4096) {
768            let mut bytes = Vec::with_capacity(8 + chunk.len() * 8);
769            bytes.extend_from_slice(&h.to_le_bytes());
770            for e in chunk {
771                bytes.extend_from_slice(&(*e as u64).to_le_bytes());
772            }
773            h = fingerprint_of(&bytes);
774        }
775        h
776    }
777
778    /// One solve of the CNF plus the unit literals `fixed` and `extra` plus
779    /// the clauses `blocks`: a model (one bool per CNF variable) or `None`
780    /// when there is none. `splr` has no assumptions in its default build, so
781    /// each call is a fresh solver; the clause list is shared, not copied.
782    fn run(
783        &self,
784        fixed: &[i32],
785        extra: &[i32],
786        blocks: &[Vec<i32>],
787        deadline: Instant,
788    ) -> Result<Option<Vec<bool>>, SolveError> {
789        if self.contradiction {
790            return Ok(None);
791        }
792        let remaining = deadline
793            .checked_duration_since(Instant::now())
794            .filter(|d| !d.is_zero())
795            .ok_or(SolveError::SatTimeout)?;
796        let mut clauses: Vec<&[i32]> =
797            Vec::with_capacity(self.ends.len() + fixed.len() + extra.len() + blocks.len());
798        let mut start = 0;
799        for &end in &self.ends {
800            clauses.push(&self.lits[start..end]);
801            start = end;
802        }
803        for l in fixed.iter().chain(extra) {
804            clauses.push(std::slice::from_ref(l));
805        }
806        for b in blocks {
807            clauses.push(b.as_slice());
808        }
809        if clauses.is_empty() {
810            return Ok(Some(vec![false; self.num_vars]));
811        }
812        let config = Config {
813            quiet_mode: true,
814            c_timeout: remaining.as_secs_f64().ceil().max(1.0),
815            ..Config::default()
816        };
817        let answer = match Cdcl::try_from((config, clauses.as_slice())) {
818            Ok(mut solver) => solver.solve(),
819            Err(Ok(cert)) => Ok(cert),
820            Err(Err(SolverError::EmptyClause)) => Ok(Certificate::UNSAT),
821            Err(Err(e)) => Err(e),
822        };
823        match answer {
824            Ok(Certificate::SAT(vals)) => {
825                let mut model = vec![false; self.num_vars];
826                for l in vals {
827                    let idx = l.unsigned_abs() as usize - 1;
828                    if idx < model.len() {
829                        model[idx] = l > 0;
830                    }
831                }
832                Ok(Some(model))
833            }
834            Ok(Certificate::UNSAT)
835            | Err(
836                SolverError::EmptyClause
837                | SolverError::Inconsistent
838                | SolverError::RootLevelConflict(_),
839            ) => Ok(None),
840            Err(SolverError::TimeOut) => Err(SolveError::SatTimeout),
841            Err(SolverError::OutOfMemory) => Err(SolveError::SatSolverFailure {
842                what: "the solver ran out of memory",
843            }),
844            Err(SolverError::SolverBug) => Err(SolveError::SatSolverFailure {
845                what: "the solver reported an internal error",
846            }),
847            Err(_) => Err(SolveError::SatSolverFailure {
848                what: "the solver rejected the CNF",
849            }),
850        }
851    }
852}
853
854/// The state of one sampling call: the unit literals fixed so far and a model
855/// that satisfies them (the witness that lets most decisions skip a solve).
856struct Walk<'a> {
857    s: &'a SatSolver,
858    deadline: Instant,
859    fixed: Vec<i32>,
860    model: Vec<bool>,
861}
862
863impl<'a> Walk<'a> {
864    /// The first model, or [`SolveError::Unsatisfiable`].
865    fn start(s: &'a SatSolver, deadline: Instant) -> Result<Walk<'a>, SolveError> {
866        match s.run(&[], &[], &[], deadline)? {
867            Some(model) => Ok(Walk {
868                s,
869                deadline,
870                fixed: Vec::new(),
871                model,
872            }),
873            None => Err(SolveError::Unsatisfiable),
874        }
875    }
876
877    /// Whether the current model already satisfies every literal of `units`.
878    fn agrees(&self, units: &[i32]) -> bool {
879        units
880            .iter()
881            .all(|&l| self.model[l.unsigned_abs() as usize - 1] == (l > 0))
882    }
883
884    /// Whether some model satisfies the fixed literals and `units`; on yes the
885    /// current model becomes that one.
886    fn probe(&mut self, units: &[i32]) -> Result<bool, SolveError> {
887        match self.s.run(&self.fixed, units, &[], self.deadline)? {
888            Some(model) => {
889                self.model = model;
890                Ok(true)
891            }
892            None => Ok(false),
893        }
894    }
895
896    fn agrees_or_probe(&mut self, units: &[i32]) -> Result<bool, SolveError> {
897        Ok(self.agrees(units) || self.probe(units)?)
898    }
899
900    /// Fixes `units`, which must be feasible (a bit or code just checked).
901    fn commit(&mut self, units: &[i32]) -> Result<(), SolveError> {
902        if !self.agrees(units) && !self.probe(units)? {
903            return Err(SolveError::SatSolverFailure {
904                what: "a value found feasible became infeasible",
905            });
906        }
907        self.fixed.extend_from_slice(units);
908        Ok(())
909    }
910
911    /// Decides CNF variables `idxs` in order: each takes its coin from
912    /// `coins` unless that leaves no model given the earlier decisions, in
913    /// which case it takes the other value; all are fixed. The result is
914    /// exactly what deciding one bit at a time gives, but a run of coins that
915    /// is feasible together costs one solve (or none, when the current model
916    /// already agrees) instead of one per bit: a prefix of a feasible run is
917    /// feasible, so a feasible whole decides every bit of it the same way, and
918    /// an infeasible run is split in two and each half decided in turn.
919    fn decide(&mut self, idxs: &[usize], coins: &[bool]) -> Result<Vec<bool>, SolveError> {
920        let mut out = Vec::with_capacity(idxs.len());
921        self.decide_run(idxs, coins, &mut out)?;
922        Ok(out)
923    }
924
925    fn decide_run(
926        &mut self,
927        idxs: &[usize],
928        coins: &[bool],
929        out: &mut Vec<bool>,
930    ) -> Result<(), SolveError> {
931        if idxs.is_empty() {
932            return Ok(());
933        }
934        let units: Vec<i32> = idxs
935            .iter()
936            .zip(coins)
937            .map(|(&i, &c)| lit_of(i, c))
938            .collect();
939        if self.agrees_or_probe(&units)? {
940            self.fixed.extend_from_slice(&units);
941            out.extend_from_slice(coins);
942            return Ok(());
943        }
944        if let [idx] = idxs {
945            // The coin has no model, so the other value must.
946            let flipped = [lit_of(*idx, !coins[0])];
947            self.commit(&flipped)?;
948            out.push(!coins[0]);
949            return Ok(());
950        }
951        let mid = idxs.len() / 2;
952        self.decide_run(&idxs[..mid], &coins[..mid], out)?;
953        self.decide_run(&idxs[mid..], &coins[mid..], out)
954    }
955}