Skip to main content

sv_sim_sim/constraints/
solve.rs

1//! The solver orchestrator: constraint IR to BDD to exact model count to an
2//! exactly uniform sample (plan C3; LRM 18.4.1).
3//!
4//! # What "uniform" means here
5//!
6//! [`Solver::sample`] draws one assignment uniformly at random from *all*
7//! legal combinations of the random variables: every model of the constraint
8//! set has probability `1 / count`. It walks the variables in a caller-supplied
9//! order, most significant bit first, and fixes each bit with probability
10//! `models-with-that-bit / models-remaining`, so the product along any path
11//! is exactly `1 / count` whatever the order is. (This is the probability
12//! LRM 18.5.9's Table 18-1 states for the *unordered* case.) `solve ...
13//! before` (C4) changes which *marginal* is uniform:
14//! [`Solver::sample_groups`] samples each ordered set uniformly over its own
15//! feasible values, built on the same per-bit machinery.
16//!
17//! # Counts
18//!
19//! The plan sketched a `u128` count. Two 64-bit random variables already have
20//! up to 2^128 combinations, so [`SatCount`] is an arbitrary-precision
21//! integer; nothing in the sampler can silently wrap.
22//!
23//! # Randomness
24//!
25//! Draws come from the `&mut u64` the caller passes (for `randomize()`, the
26//! object's own PRNG seed; see [`super::rng`]). An unsatisfiable problem
27//! consumes no draws, so a failed `randomize()` does not perturb the object's
28//! random stream (LRM 18.6.3).
29
30use super::bdd::{BddBuildError, Compiled, NODE_BUDGET, compile};
31use super::ir::{Expr, Ty, VarId};
32use super::randc::{self, RandcState, fingerprint_of};
33use super::rng;
34use super::sat::{SatLimits, SatSolver};
35use biodivine_lib_bdd::{Bdd, BddValuation};
36use core::fmt;
37use num_bigint::BigUint;
38use num_traits::{One, Zero};
39use std::collections::{BTreeMap, BTreeSet, HashMap};
40use std::sync::{Arc, Mutex};
41
42/// A number of models (assignments of every random-variable bit).
43pub type SatCount = BigUint;
44
45/// Why a solve failed.
46#[derive(Clone, Debug, PartialEq, Eq)]
47pub enum SolveError {
48    /// The constraints could not be turned into a BDD or, for an [`Engine`]
49    /// that fell back, into a CNF (including [`BddBuildError::TooLarge`] from
50    /// [`Solver`], the SAT-fallback trigger).
51    Build(BddBuildError),
52    /// No assignment satisfies the constraints (LRM 18.6.3: `randomize()`
53    /// returns 0 and changes nothing).
54    Unsatisfiable,
55    /// The caller's variable order names an unknown or repeated variable.
56    BadOrder {
57        /// The offending entry.
58        var: VarId,
59    },
60    /// The [`Plan`] is malformed in a way other than naming a bad variable.
61    BadPlan {
62        /// What is wrong.
63        what: &'static str,
64    },
65    /// The SAT fallback ([`super::sat`]) did not finish within its time limit.
66    SatTimeout,
67    /// The plan asks the SAT fallback for something it cannot sample
68    /// faithfully (a `dist` weighted by a count of models or of distinct
69    /// values, a `randc` variable too wide to enumerate); refused instead of
70    /// silently mis-weighted.
71    SatUnsupported {
72        /// What is not supported.
73        what: &'static str,
74    },
75    /// A model the SAT solver returned failed the IR's own evaluator: a bug
76    /// in the CNF encoding, never a legal outcome.
77    SatModelRejected,
78    /// The SAT solver itself failed (out of memory, an internal error) or
79    /// contradicted itself; never a legal outcome.
80    SatSolverFailure {
81        /// What went wrong.
82        what: &'static str,
83    },
84}
85
86/// An exact non-negative rational weight `num / den` (`den` is never 0).
87#[derive(Clone, Debug, PartialEq, Eq)]
88pub struct Weight {
89    /// The numerator.
90    pub num: BigUint,
91    /// The denominator.
92    pub den: BigUint,
93}
94
95impl Weight {
96    /// The integer weight `n`.
97    pub fn int(n: impl Into<BigUint>) -> Weight {
98        Weight {
99            num: n.into(),
100            den: BigUint::one(),
101        }
102    }
103}
104
105/// How much of the *rest of the problem* a code of a [`Weighted`] variable
106/// stands for, which scales the code's weight.
107#[derive(Clone, Copy, Debug, PartialEq, Eq)]
108pub enum Measure {
109    /// Every feasible code counts once.
110    Flat,
111    /// A code counts once per feasible *tuple* (model) it leaves.
112    Models,
113    /// A code counts once per distinct feasible value of this variable (a
114    /// member of the same [`Group`]) it leaves, and once the code is drawn that
115    /// variable is itself drawn, uniformly over those values.
116    Distinct(VarId),
117}
118
119/// A variable drawn by weight instead of uniformly (a `dist` selector, C6):
120/// code `c` is chosen with probability proportional to
121/// `weights[c] * measure(c)` among the codes that are still feasible; a code
122/// with weight 0, or beyond `weights.len()`, is never chosen.
123///
124/// With [`Measure::Distinct`] this is exact for LRM 18.5.3: the selector's
125/// code is an item of the `dist`, its weight the item's per-value weight, and
126/// the measure the number of values of the subject that are both in the item
127/// and still feasible, so the probability of a feasible value is proportional
128/// to the sum of the per-value weights of the items that contain it, however
129/// the items overlap or the other constraints exclude values.
130#[derive(Clone, Debug, PartialEq, Eq)]
131pub struct Weighted {
132    /// The variable.
133    pub var: VarId,
134    /// The weight of each value of the variable, indexed by value.
135    pub weights: Vec<Weight>,
136    /// What a code stands for.
137    pub measure: Measure,
138}
139
140/// One set of variables solved together ([`Plan`]).
141#[derive(Clone, Debug, Default, PartialEq, Eq)]
142pub struct Group {
143    /// Variables drawn uniformly over the group's feasible tuples.
144    pub vars: Vec<VarId>,
145    /// Variables drawn by weight, before `vars`.
146    pub weighted: Vec<Weighted>,
147}
148
149/// How [`Solver::sample_plan`] solves: the `randc` variables first (LRM
150/// 18.4.2), then each ordered [`Group`] in turn (18.5.9).
151#[derive(Clone, Debug, Default, PartialEq, Eq)]
152pub struct Plan {
153    /// The ordered groups.
154    pub groups: Vec<Group>,
155    /// The `randc` variables, in the order they are drawn.
156    pub randc: Vec<VarId>,
157}
158
159/// A [`Weighted`] resolved to bit-vector indices.
160pub(crate) struct ResolvedWeighted<'a> {
161    pub(crate) var: usize,
162    pub(crate) weights: &'a [Weight],
163    pub(crate) measure: Measure,
164}
165
166/// A [`Group`] resolved to bit-vector indices: its uniform variables and its
167/// weighted ones.
168pub(crate) type Resolved<'a> = (Vec<usize>, Vec<ResolvedWeighted<'a>>);
169
170impl fmt::Display for SolveError {
171    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
172        match self {
173            SolveError::Build(e) => fmt::Display::fmt(e, f),
174            SolveError::Unsatisfiable => f.write_str("the constraints have no solution"),
175            SolveError::BadOrder { var } => write!(
176                f,
177                "solve order names random variable #{var} which is unknown or listed twice"
178            ),
179            SolveError::BadPlan { what } => write!(f, "malformed solve plan: {what}"),
180            SolveError::SatTimeout => {
181                f.write_str("the SAT fallback did not finish within its time limit")
182            }
183            SolveError::SatUnsupported { what } => {
184                write!(f, "not supported by the SAT fallback: {what}")
185            }
186            SolveError::SatModelRejected => f.write_str(
187                "the SAT fallback returned an assignment the constraint evaluator rejects \
188                 (an encoding bug)",
189            ),
190            SolveError::SatSolverFailure { what } => write!(f, "the SAT fallback failed: {what}"),
191        }
192    }
193}
194
195impl std::error::Error for SolveError {}
196
197impl From<BddBuildError> for SolveError {
198    fn from(e: BddBuildError) -> Self {
199        SolveError::Build(e)
200    }
201}
202
203/// Compiled problems by [`problem_digest`]. Compiling is a pure function of
204/// its arguments, and a `randomize()` in a loop (or a soft-constraint ladder,
205/// LRM 18.5.13.1, which retries the same hard problem) re-asks the same one
206/// over and over; the build was most of a call. Bounded in entries and in the
207/// size of what is kept; cleared when full.
208static SOLVER_CACHE: Mutex<BTreeMap<(u64, u64), Solver>> = Mutex::new(BTreeMap::new());
209const SOLVER_CACHE_CAP: usize = 512;
210const SOLVER_CACHE_MAX_NODES: usize = 4096;
211
212/// A compiled, counted constraint problem, ready to sample.
213#[derive(Clone, Debug)]
214pub struct Solver {
215    compiled: Compiled,
216    count: SatCount,
217    /// Per-bit draw steps already computed, shared by every clone (see [`StepMemo`]).
218    steps: Arc<Mutex<StepMemo>>,
219}
220
221/// One bit decision of the uniform tuple draw: the models with the bit set
222/// and their count, and (once a draw needed it) the models with it clear.
223#[derive(Debug)]
224struct Step {
225    ones: Bdd,
226    ones_count: BigUint,
227    zeros: Option<Bdd>,
228}
229
230/// Memo of [`Solver::sample_plan`]'s per-bit BDD conjunctions. The BDD a draw
231/// holds before a bit is a pure function of the plan and of every decision
232/// taken earlier in the draw (its `trace`), so a repeated prefix reuses the
233/// conjunction and the cardinality that an apply and a model count would
234/// recompute. Only the random numbers differ between draws, so the output is
235/// unchanged. Bounded: cleared when it holds [`STEP_MEMO_CAP`] steps.
236#[derive(Debug, Default)]
237struct StepMemo {
238    plans: Vec<(Plan, HashMap<Vec<u32>, Step>)>,
239}
240const STEP_MEMO_CAP: usize = 100_000;
241const STEP_MEMO_PLANS: usize = 4;
242
243impl Solver {
244    /// Compiles the conjunction of `conjuncts` with the default node budget.
245    pub fn new(
246        var_tys: &[Ty],
247        states: &[Option<BigUint>],
248        conjuncts: &[Expr],
249    ) -> Result<Solver, SolveError> {
250        Solver::with_budget(var_tys, states, conjuncts, NODE_BUDGET)
251    }
252
253    /// As [`Solver::new`] with an explicit per-BDD node budget.
254    pub fn with_budget(
255        var_tys: &[Ty],
256        states: &[Option<BigUint>],
257        conjuncts: &[Expr],
258        budget: usize,
259    ) -> Result<Solver, SolveError> {
260        let digest = problem_digest(var_tys, states, conjuncts, budget);
261        Solver::with_digest(digest, var_tys, states, conjuncts, budget)
262    }
263
264    /// [`Solver::with_budget`] for a caller that already hashed the problem
265    /// (hashing formats the whole expression tree, so it is done once per call).
266    fn with_digest(
267        digest: (u64, u64),
268        var_tys: &[Ty],
269        states: &[Option<BigUint>],
270        conjuncts: &[Expr],
271        budget: usize,
272    ) -> Result<Solver, SolveError> {
273        if let Some(s) = SOLVER_CACHE
274            .lock()
275            .unwrap_or_else(|p| p.into_inner())
276            .get(&digest)
277        {
278            return Ok(s.clone());
279        }
280        let compiled = compile(var_tys, states, conjuncts, budget)?;
281        let count = compiled.bdd.exact_cardinality();
282        let solver = Solver {
283            compiled,
284            count,
285            steps: Arc::default(),
286        };
287        if solver.compiled.bdd.size() <= SOLVER_CACHE_MAX_NODES {
288            let mut cache = SOLVER_CACHE.lock().unwrap_or_else(|p| p.into_inner());
289            if cache.len() >= SOLVER_CACHE_CAP {
290                cache.clear();
291            }
292            cache.insert(digest, solver.clone());
293        }
294        Ok(solver)
295    }
296
297    /// The exact number of legal assignments.
298    pub fn count(&self) -> &SatCount {
299        &self.count
300    }
301
302    /// The underlying BDD and variable layout.
303    pub fn compiled(&self) -> &Compiled {
304        &self.compiled
305    }
306
307    /// Draws one assignment uniformly from all models (module docs), fixing
308    /// variables in `order`, then any variable `order` omits in id order.
309    /// Returns one value per variable, indexed by [`VarId`].
310    pub fn sample(&self, order: &[VarId], seed: &mut u64) -> Result<Vec<BigUint>, SolveError> {
311        // One group: no later group to project away, so this is exactly the
312        // per-bit walk over all models.
313        self.sample_groups(&[order.to_vec()], seed)
314    }
315
316    /// Draws one assignment solving the variables group by group, in the
317    /// order of `groups` (`solve ... before`, LRM 18.5.9; [`super::order`]
318    /// builds them).
319    ///
320    /// Each group is sampled *uniformly over its own distinct feasible value
321    /// tuples*: the variables of every later group are existentially
322    /// projected away, so a group value is feasible when some assignment of
323    /// the later groups completes it, and its probability does not depend on
324    /// how many completions it has. That is the LRM's `solve s before d`
325    /// (Table 18-2: `s` is 50/50 although `s == 1` has a single completion
326    /// and `s == 0` has 2^32). The chosen values are then fixed in the full
327    /// constraint set and the next group is drawn from what remains. Inside a
328    /// group the per-bit walk of [`Solver::sample`] applies, so a group is
329    /// uniform over its tuples whatever its internal order.
330    ///
331    /// Any variable `groups` omits is solved with the *last* group (LRM
332    /// 18.5.9: unordered variables are solved with the last set of ordered
333    /// variables), in id order. A variable listed twice is
334    /// [`SolveError::BadOrder`].
335    pub fn sample_groups(
336        &self,
337        groups: &[Vec<VarId>],
338        seed: &mut u64,
339    ) -> Result<Vec<BigUint>, SolveError> {
340        let plan = Plan {
341            groups: groups
342                .iter()
343                .map(|g| Group {
344                    vars: g.clone(),
345                    weighted: Vec::new(),
346                })
347                .collect(),
348            randc: Vec::new(),
349        };
350        self.sample_plan(&plan, &mut [], seed)
351    }
352
353    /// Draws one assignment following `plan` (see [`Plan`]): first the `randc`
354    /// variables, each from its own cycle state (`randc_states[i]` belongs to
355    /// `plan.randc[i]`; [`super::randc`]), then the ordered groups, each with
356    /// its weighted variables before its uniform ones. A state is only
357    /// advanced for a variable that was drawn, so a caller that must not
358    /// commit a failed call clones the states first.
359    pub fn sample_plan(
360        &self,
361        plan: &Plan,
362        randc_states: &mut [RandcState],
363        seed: &mut u64,
364    ) -> Result<Vec<BigUint>, SolveError> {
365        let layout = &self.compiled.layout;
366        let n = layout.tys.len();
367        let widths: Vec<usize> = layout.bits.iter().map(Vec::len).collect();
368        let full = resolve_plan(&widths, plan, randc_states.len())?;
369        let last = full.len() - 1;
370
371        if self.count.is_zero() {
372            return Err(SolveError::Unsatisfiable);
373        }
374        let mut cur = self.compiled.bdd.clone();
375        let mut out = vec![BigUint::zero(); n];
376        let mut memo_guard = self.steps.lock().unwrap_or_else(|p| p.into_inner());
377        let memo_all = &mut *memo_guard;
378        if memo_all.plans.iter().map(|(_, m)| m.len()).sum::<usize>() >= STEP_MEMO_CAP {
379            *memo_all = StepMemo::default();
380        }
381        let slot = match memo_all.plans.iter().position(|(p, _)| p == plan) {
382            Some(i) => i,
383            None => {
384                if memo_all.plans.len() >= STEP_MEMO_PLANS {
385                    *memo_all = StepMemo::default();
386                }
387                memo_all.plans.push((plan.clone(), HashMap::new()));
388                memo_all.plans.len() - 1
389            }
390        };
391        let memo = &mut memo_all.plans[slot].1;
392        // Every decision of this draw so far: the memo key. A tagged word is a
393        // randc or weighted code (tag 1) or a bit (tag 0, `var << 1 | bit`).
394        let mut trace: Vec<u32> = Vec::new();
395
396        // The `randc` variables come first, before every `rand` one (LRM
397        // 18.4.2), in id order.
398        for (k, &v) in plan.randc.iter().enumerate() {
399            let v = v as usize;
400            let others: Vec<_> = (0..n)
401                .filter(|&o| o != v)
402                .flat_map(|o| layout.bits[o].iter().copied())
403                .collect();
404            let base = self.compiled.bdd.exists(&others);
405            let now = cur.exists(&others);
406            let w = layout.bits[v].len();
407            let holds = |f: &Bdd, x: u32| -> bool {
408                let mut val = BddValuation::all_false(layout.num_bits() as u16);
409                for (i, bv) in layout.bits[v].iter().enumerate() {
410                    val.set_value(*bv, (x >> i) & 1 == 1);
411                }
412                f.eval_in(&val)
413            };
414            let fp = fingerprint_of(&base.to_bytes());
415            let picked = randc::pick(
416                &mut randc_states[k],
417                fp,
418                &mut || (0..1u32 << w).filter(|&x| holds(&base, x)).collect(),
419                &mut |x| holds(&now, x),
420                seed,
421            )
422            .ok_or(SolveError::Unsatisfiable)?;
423            trace.push((1 << 31) | picked);
424            for (i, bv) in layout.bits[v].iter().enumerate() {
425                let bit = (picked >> i) & 1 == 1;
426                out[v].set_bit(i as u64, bit);
427                cur = cur.var_select(*bv, bit);
428            }
429        }
430
431        for (gi, (group, weighted)) in full.iter().enumerate() {
432            trace.push(u32::MAX - gi as u32);
433            // Existentially project the later groups away. Variables that
434            // become free multiply every count by the same power of two, which
435            // cancels in the ratios below.
436            let later: Vec<_> = full[gi + 1..]
437                .iter()
438                .flat_map(|(vars, wtd)| vars.iter().chain(wtd.iter().map(|w| &w.var)))
439                .flat_map(|&v| layout.bits[v].iter().copied())
440                .collect();
441            let mut proj = if later.is_empty() {
442                cur.clone()
443            } else {
444                cur.exists(&later)
445            };
446            // Weighted variables: a code is chosen with probability
447            // proportional to its weight, scaled by its measure, among the
448            // codes still feasible.
449            let mut drawn: Vec<usize> = Vec::new();
450            for wv in weighted {
451                let v = wv.var;
452                let bits = &layout.bits[v];
453                let fixing = |code: usize| -> Vec<_> {
454                    bits.iter()
455                        .enumerate()
456                        .map(|(i, bv)| (*bv, (code >> i) & 1 == 1))
457                        .collect()
458                };
459                // Every bit but the subject's, to project the rest away.
460                let all_but = |keep: usize| -> Vec<_> {
461                    (0..n)
462                        .filter(|&o| o != keep)
463                        .flat_map(|o| layout.bits[o].iter().copied())
464                        .collect()
465                };
466                let subject_rest = match wv.measure {
467                    Measure::Distinct(s) => Some(all_but(s as usize)),
468                    Measure::Flat | Measure::Models => None,
469                };
470                // (code, numerator * measure, denominator)
471                let mut feasible: Vec<(usize, BigUint, &BigUint)> = Vec::new();
472                for (c, wt) in wv.weights.iter().enumerate() {
473                    if wt.num.is_zero() {
474                        continue;
475                    }
476                    let pc = proj.restrict(&fixing(c));
477                    if pc.is_false() {
478                        continue;
479                    }
480                    let m = match (&wv.measure, &subject_rest) {
481                        (Measure::Models, _) => pc.exact_cardinality(),
482                        (_, Some(rest)) => pc.exists(rest).exact_cardinality(),
483                        _ => BigUint::one(),
484                    };
485                    feasible.push((c, &wt.num * m, &wt.den));
486                }
487                let code = draw_weighted(&feasible, seed)?;
488                trace.push((1 << 31) | code as u32);
489                let fix = fixing(code);
490                proj = proj.restrict(&fix);
491                if gi < last {
492                    cur = cur.restrict(&fix);
493                }
494                for (i, _) in bits.iter().enumerate() {
495                    out[v].set_bit(i as u64, (code >> i) & 1 == 1);
496                }
497                // The subject is drawn now, uniformly over its distinct
498                // feasible values (they are all in the chosen item).
499                if let (Measure::Distinct(s), Some(rest)) = (wv.measure, &subject_rest) {
500                    let s = s as usize;
501                    let mut pj = proj.exists(rest);
502                    let mut remaining = pj.exact_cardinality();
503                    for i in (0..layout.bits[s].len()).rev() {
504                        let bv = layout.bits[s][i];
505                        let ones = pj.var_select(bv, true);
506                        let c1 = ones.exact_cardinality();
507                        let bit = rng::below(seed, &remaining) < c1;
508                        if bit {
509                            pj = ones;
510                            remaining = c1;
511                        } else {
512                            remaining -= &c1;
513                            pj = pj.var_select(bv, false);
514                        }
515                        proj = proj.var_select(bv, bit);
516                        trace.push((bv.to_index() as u32) << 1 | u32::from(bit));
517                        if gi < last {
518                            cur = cur.var_select(bv, bit);
519                        }
520                        out[s].set_bit(i as u64, bit);
521                    }
522                    drawn.push(s);
523                }
524            }
525            let mut remaining = proj.exact_cardinality();
526            // A bit outside `proj`'s support is unconstrained: exactly half of
527            // the `remaining` models have it set, so it is drawn from that
528            // count without conjoining a literal to `proj` (one full BDD
529            // apply per bit is what dominated a randomize call). `proj` then
530            // counts every skipped bit as free, so the models with a later
531            // bit fixed are `skipped` doublings too many; both are undone
532            // exactly below. The last group only: an earlier one must also
533            // keep `cur` fixed for the groups after it.
534            let support = (gi == last).then(|| proj.support_set());
535            let mut skipped = 0u64;
536            for &v in group {
537                if drawn.contains(&v) {
538                    continue;
539                }
540                for i in (0..layout.bits[v].len()).rev() {
541                    let bv = layout.bits[v][i];
542                    if support.as_ref().is_some_and(|s| !s.contains(&bv)) {
543                        let half = &remaining >> 1u32;
544                        let bit = rng::below(seed, &remaining) < half;
545                        remaining = half;
546                        skipped += 1;
547                        if bit {
548                            out[v].set_bit(i as u64, true);
549                        }
550                        continue;
551                    }
552                    if !memo.contains_key(trace.as_slice()) {
553                        let ones = proj.var_select(bv, true);
554                        let ones_count = ones.exact_cardinality();
555                        memo.insert(
556                            trace.clone(),
557                            Step {
558                                ones,
559                                ones_count,
560                                zeros: None,
561                            },
562                        );
563                    }
564                    let step = memo.get_mut(trace.as_slice()).expect("just inserted");
565                    let c1 = &step.ones_count >> skipped;
566                    let zeros_count = &remaining - &c1;
567                    let pick = rng::below(seed, &remaining);
568                    let bit = pick < c1;
569                    if bit {
570                        out[v].set_bit(i as u64, true);
571                        proj = step.ones.clone();
572                        remaining = c1;
573                    } else {
574                        proj = step
575                            .zeros
576                            .get_or_insert_with(|| proj.var_select(bv, false))
577                            .clone();
578                        remaining = zeros_count;
579                    }
580                    trace.push((bv.to_index() as u32) << 1 | u32::from(bit));
581                    // The last group's `proj` is `cur`; nothing follows it.
582                    if gi < last {
583                        cur = cur.var_select(bv, bit);
584                    }
585                }
586            }
587        }
588        debug_assert!(self.compiled.accepts(&out), "sampled a non-model");
589        Ok(out)
590    }
591}
592
593/// Validates `plan` against the variables' bit `widths` and resolves it to
594/// bit-vector indices: one entry per group, the variables the plan omits
595/// appended to the last group (LRM 18.5.9), and a single empty group when
596/// the plan has none. `randc_states` is how many cycle states the caller
597/// holds. Shared by the exact and the SAT sampler so both reject exactly the
598/// same malformed plans.
599pub(crate) fn resolve_plan<'a>(
600    widths: &[usize],
601    plan: &'a Plan,
602    randc_states: usize,
603) -> Result<Vec<Resolved<'a>>, SolveError> {
604    let n = widths.len();
605    if randc_states != plan.randc.len() {
606        return Err(SolveError::BadPlan {
607            what: "one randc state per randc variable",
608        });
609    }
610    for &v in &plan.randc {
611        let w = widths.get(v as usize).copied().unwrap_or(0);
612        if w > randc::MAX_RANDC_WIDTH as usize {
613            return Err(SolveError::BadPlan {
614                what: "a randc variable is wider than the permutation engine supports",
615            });
616        }
617    }
618    let mut seen = vec![false; n];
619    let mut mark = |v: VarId| -> Result<usize, SolveError> {
620        let i = v as usize;
621        if i >= n || seen[i] {
622            return Err(SolveError::BadOrder { var: v });
623        }
624        seen[i] = true;
625        Ok(i)
626    };
627    for &v in &plan.randc {
628        mark(v)?;
629    }
630    let mut full: Vec<Resolved<'a>> = Vec::with_capacity(plan.groups.len().max(1));
631    for g in &plan.groups {
632        let mut weighted = Vec::with_capacity(g.weighted.len());
633        for w in &g.weighted {
634            let i = mark(w.var)?;
635            let width = widths[i];
636            if width >= 32 || w.weights.len() > 1usize << width {
637                return Err(SolveError::BadPlan {
638                    what: "a weighted variable has more weights than values",
639                });
640            }
641            if w.weights.iter().any(|x| x.den.is_zero()) {
642                return Err(SolveError::BadPlan {
643                    what: "a weight has a zero denominator",
644                });
645            }
646            weighted.push(ResolvedWeighted {
647                var: i,
648                weights: w.weights.as_slice(),
649                measure: w.measure,
650            });
651        }
652        let mut members = Vec::with_capacity(g.vars.len());
653        for &v in &g.vars {
654            members.push(mark(v)?);
655        }
656        for w in &weighted {
657            if let Measure::Distinct(s) = w.measure
658                && !members.contains(&(s as usize))
659            {
660                return Err(SolveError::BadPlan {
661                    what: "a weighted variable's subject is not a uniform member of its group",
662                });
663            }
664        }
665        full.push((members, weighted));
666    }
667    if full.is_empty() {
668        full.push((Vec::new(), Vec::new()));
669    }
670    let last = full.len() - 1;
671    full[last].0.extend((0..n).filter(|i| !seen[*i]));
672    Ok(full)
673}
674
675/// Draws one code from `feasible` (`(code, numerator * measure,
676/// denominator)`) with probability proportional to its weight, in exact
677/// rational arithmetic over the common denominator; no feasible weight is
678/// [`SolveError::Unsatisfiable`] and consumes no draw.
679pub(crate) fn draw_weighted(
680    feasible: &[(usize, BigUint, &BigUint)],
681    seed: &mut u64,
682) -> Result<usize, SolveError> {
683    let common = feasible
684        .iter()
685        .fold(BigUint::one(), |acc, (_, _, den)| acc * *den);
686    let scaled: Vec<(usize, BigUint)> = feasible
687        .iter()
688        .map(|(c, num, den)| (*c, num * (&common / *den)))
689        .collect();
690    let total: BigUint = scaled.iter().map(|(_, w)| w.clone()).sum();
691    if total.is_zero() {
692        return Err(SolveError::Unsatisfiable);
693    }
694    let mut pick = rng::below(seed, &total);
695    let mut code = scaled[scaled.len() - 1].0;
696    for (c, w) in &scaled {
697        if pick < *w {
698            code = *c;
699            break;
700        }
701        pick -= w;
702    }
703    Ok(code)
704}
705
706/// The solver `randomize()` uses: the exact BDD [`Solver`] when the problem
707/// fits its node budget, the approximate [`SatSolver`] when it does not (plan
708/// section 3.4, ledger XC-063). Both take the same [`Plan`] and obey the same
709/// contracts; what differs is exactness of the distribution and what the SAT
710/// side refuses ([`SolveError::SatUnsupported`]; [`super::sat`] module docs).
711#[derive(Debug)]
712pub enum Engine {
713    /// The problem fit the BDD node budget: exactly uniform.
714    Exact(Solver),
715    /// It did not: approximately uniform.
716    Sat(SatSolver),
717}
718
719/// Problems (as a 128-bit digest) whose BDD build already exceeded its budget
720/// in this process, so the next call skips straight to SAT instead of paying
721/// for another failing build. Bounded; cleared when full.
722static BDD_TOO_LARGE: Mutex<BTreeSet<(u64, u64)>> = Mutex::new(BTreeSet::new());
723const BDD_TOO_LARGE_CAP: usize = 4096;
724
725/// The digest identifying one build problem: two independent fixed-key SipHash
726/// streams over the structural `Hash` of the inputs (no text is formatted; the
727/// digest is only a cache key, never persisted).
728fn problem_digest(
729    var_tys: &[Ty],
730    states: &[Option<BigUint>],
731    conjuncts: &[Expr],
732    budget: usize,
733) -> (u64, u64) {
734    use std::hash::{Hash as _, Hasher as _};
735    let mut lo = std::collections::hash_map::DefaultHasher::new();
736    let mut hi = std::collections::hash_map::DefaultHasher::new();
737    0x1234_5678_9abc_def1_u64.hash(&mut hi);
738    for h in [&mut lo, &mut hi] {
739        budget.hash(h);
740        var_tys.hash(h);
741        states.hash(h);
742        conjuncts.hash(h);
743    }
744    (lo.finish(), hi.finish())
745}
746
747impl Engine {
748    /// Compiles the conjunction of `conjuncts` with the default BDD budget and
749    /// SAT limits.
750    pub fn new(
751        var_tys: &[Ty],
752        states: &[Option<BigUint>],
753        conjuncts: &[Expr],
754    ) -> Result<Engine, SolveError> {
755        Engine::with_limits(
756            var_tys,
757            states,
758            conjuncts,
759            NODE_BUDGET,
760            SatLimits::default(),
761        )
762    }
763
764    /// As [`Engine::new`] with an explicit BDD node budget and SAT limits.
765    pub fn with_limits(
766        var_tys: &[Ty],
767        states: &[Option<BigUint>],
768        conjuncts: &[Expr],
769        budget: usize,
770        limits: SatLimits,
771    ) -> Result<Engine, SolveError> {
772        let digest = problem_digest(var_tys, states, conjuncts, budget);
773        let known_large = BDD_TOO_LARGE
774            .lock()
775            .unwrap_or_else(|p| p.into_inner())
776            .contains(&digest);
777        if !known_large {
778            match Solver::with_digest(digest, var_tys, states, conjuncts, budget) {
779                Ok(s) => return Ok(Engine::Exact(s)),
780                Err(SolveError::Build(BddBuildError::TooLarge { .. })) => {
781                    let mut seen = BDD_TOO_LARGE.lock().unwrap_or_else(|p| p.into_inner());
782                    if seen.len() >= BDD_TOO_LARGE_CAP {
783                        seen.clear();
784                    }
785                    seen.insert(digest);
786                }
787                Err(e) => return Err(e),
788            }
789        }
790        SatSolver::with_limits(var_tys, states, conjuncts, limits).map(Engine::Sat)
791    }
792
793    /// Whether this problem fell back to SAT.
794    pub fn is_sat(&self) -> bool {
795        matches!(self, Engine::Sat(_))
796    }
797
798    /// As [`Solver::sample_plan`] (exact) or [`SatSolver::sample_plan`].
799    pub fn sample_plan(
800        &self,
801        plan: &Plan,
802        randc_states: &mut [RandcState],
803        seed: &mut u64,
804    ) -> Result<Vec<BigUint>, SolveError> {
805        match self {
806            Engine::Exact(s) => s.sample_plan(plan, randc_states, seed),
807            Engine::Sat(s) => s.sample_plan(plan, randc_states, seed),
808        }
809    }
810}
811
812/// One-shot: compile, count and draw a single sample.
813pub fn solve(
814    var_tys: &[Ty],
815    states: &[Option<BigUint>],
816    conjuncts: &[Expr],
817    order: &[VarId],
818    seed: &mut u64,
819) -> Result<Vec<BigUint>, SolveError> {
820    Solver::new(var_tys, states, conjuncts)?.sample(order, seed)
821}
822
823/// One-shot: compile, count and draw a single sample solving `groups` in
824/// order ([`Solver::sample_groups`]).
825pub fn solve_grouped(
826    var_tys: &[Ty],
827    states: &[Option<BigUint>],
828    conjuncts: &[Expr],
829    groups: &[Vec<VarId>],
830    seed: &mut u64,
831) -> Result<Vec<BigUint>, SolveError> {
832    Solver::new(var_tys, states, conjuncts)?.sample_groups(groups, seed)
833}
834
835#[cfg(test)]
836mod tests {
837    use super::*;
838    use crate::constraints::ir::BinOp;
839
840    fn bin(op: BinOp, a: Expr, b: Expr, t: Ty) -> Expr {
841        Expr::Binary(op, Box::new(a), Box::new(b), t)
842    }
843
844    /// LRM 18.5.9 Table 18-1/18-2 with an 8-bit `d`: `s -> d == 0`.
845    fn table_18_2() -> Solver {
846        let s = Expr::Var(0, Ty::BIT);
847        let d = Expr::Var(1, Ty::unsigned(8));
848        let c = bin(
849            BinOp::LogImpl,
850            s,
851            bin(BinOp::Eq, d, Expr::constant(Ty::unsigned(8), 0u8), Ty::BIT),
852            Ty::BIT,
853        );
854        Solver::new(&[Ty::BIT, Ty::unsigned(8)], &[], &[c]).expect("compiles")
855    }
856
857    /// `(count of s == 1, count of d == 0)` over `n` draws.
858    fn tally(sv: &Solver, groups: &[Vec<VarId>], n: u32, seed0: u64) -> (u32, u32) {
859        let mut seed = seed0;
860        let (mut s1, mut d0) = (0, 0);
861        for _ in 0..n {
862            let v = sv.sample_groups(groups, &mut seed).expect("sat");
863            s1 += u32::from(!v[0].is_zero());
864            d0 += u32::from(v[1].is_zero());
865        }
866        (s1, d0)
867    }
868
869    #[test]
870    fn solve_s_before_d_is_table_18_2() {
871        let sv = table_18_2();
872        assert_eq!(sv.count(), &BigUint::from(257u32));
873        let n = 40_000;
874        let (s1, d0) = tally(&sv, &[vec![0], vec![1]], n, 11);
875        // s is 50/50 (sigma 100): the single completion of s == 1 does not
876        // make it unlikely.
877        assert!((19_400..20_600).contains(&s1), "s==1: {s1}");
878        // P(d == 0) = 1/2 + 1/2 * 2^-8 = 0.50195 (sigma 100).
879        assert!((19_500..20_700).contains(&d0), "d==0: {d0}");
880    }
881
882    #[test]
883    fn unordered_is_table_18_1() {
884        let sv = table_18_2();
885        let n = 40_000;
886        // One group, and no group at all, both mean "unordered".
887        for groups in [vec![vec![0, 1]], vec![]] {
888            let (s1, d0) = tally(&sv, &groups, n, 12);
889            // P(s == 1) = 1/257 = 155.6 (sigma 12.5); P(d == 0) = 2/257.
890            assert!((80..235).contains(&s1), "s==1: {s1}");
891            assert!((235..390).contains(&d0), "d==0: {d0}");
892        }
893    }
894
895    #[test]
896    fn solving_d_first_makes_d_uniform_over_its_feasible_values() {
897        let sv = table_18_2();
898        let n = 51_200;
899        let (s1, d0) = tally(&sv, &[vec![1], vec![0]], n, 13);
900        // Every d is feasible (s == 0 always completes it): P(d == 0) = 1/256
901        // = 200 (sigma 14). Given d == 0, s is 50/50, so P(s == 1) = 1/512.
902        assert!((130..270).contains(&d0), "d==0: {d0}");
903        assert!((60..140).contains(&s1), "s==1: {s1}");
904    }
905
906    #[test]
907    fn a_variable_the_groups_omit_is_solved_with_the_last_group() {
908        let sv = table_18_2();
909        // Only `s` is listed: `d` joins the last (only) group, so this is the
910        // unordered distribution.
911        let (s1, _) = tally(&sv, &[vec![0]], 40_000, 14);
912        assert!((80..235).contains(&s1), "s==1: {s1}");
913    }
914
915    #[test]
916    fn one_group_is_exactly_the_per_bit_sample() {
917        let sv = table_18_2();
918        let (mut a, mut b) = (5u64, 5u64);
919        for _ in 0..50 {
920            assert_eq!(
921                sv.sample(&[1, 0], &mut a).unwrap(),
922                sv.sample_groups(&[vec![1, 0]], &mut b).unwrap()
923            );
924        }
925        assert_eq!(a, b, "the same draws were consumed");
926    }
927
928    #[test]
929    fn a_repeated_or_unknown_variable_is_a_bad_order() {
930        let sv = table_18_2();
931        let mut seed = 1u64;
932        assert_eq!(
933            sv.sample_groups(&[vec![0], vec![0]], &mut seed),
934            Err(SolveError::BadOrder { var: 0 })
935        );
936        assert_eq!(
937            sv.sample_groups(&[vec![7]], &mut seed),
938            Err(SolveError::BadOrder { var: 7 })
939        );
940    }
941
942    #[test]
943    fn an_unsatisfiable_problem_consumes_no_draws_when_grouped() {
944        let x = Expr::Var(0, Ty::unsigned(4));
945        let c = bin(BinOp::Lt, x, Expr::constant(Ty::unsigned(4), 0u8), Ty::BIT);
946        let sv = Solver::new(&[Ty::unsigned(4)], &[], &[c]).unwrap();
947        let mut seed = 9u64;
948        assert_eq!(
949            sv.sample_groups(&[vec![0]], &mut seed),
950            Err(SolveError::Unsatisfiable)
951        );
952        assert_eq!(seed, 9);
953    }
954
955    #[test]
956    fn three_levels_are_each_uniform_over_their_feasible_values() {
957        // a < b < c over 3-bit values: a in 0..=5, c in 2..=7, b in 1..=6.
958        // Solving a, then b, then c: a is uniform over its 6 feasible values
959        // (unordered it would favour small a), and every draw is a model.
960        let t = Ty::unsigned(3);
961        let (a, b, c) = (Expr::Var(0, t), Expr::Var(1, t), Expr::Var(2, t));
962        let conj = vec![
963            bin(BinOp::Lt, a, b.clone(), Ty::BIT),
964            bin(BinOp::Lt, b, c, Ty::BIT),
965        ];
966        let sv = Solver::new(&[t, t, t], &[], &conj).unwrap();
967        let n = 36_000u32;
968        let mut seed = 21u64;
969        let mut hist = [0u32; 8];
970        for _ in 0..n {
971            let v = sv
972                .sample_groups(&[vec![0], vec![1], vec![2]], &mut seed)
973                .unwrap();
974            assert!(v[0] < v[1] && v[1] < v[2], "{v:?}");
975            hist[usize::try_from(v[0].clone()).unwrap()] += 1;
976        }
977        // 6 values, 6000 each (sigma ~71).
978        for (i, h) in hist.iter().enumerate() {
979            if i < 6 {
980                assert!((5_500..6_500).contains(h), "a=={i}: {hist:?}");
981            } else {
982                assert_eq!(*h, 0, "a=={i} is infeasible: {hist:?}");
983            }
984        }
985    }
986
987    fn plan(randc: Vec<VarId>, groups: Vec<Group>) -> Plan {
988        Plan { groups, randc }
989    }
990
991    fn uniform(vars: Vec<VarId>) -> Group {
992        Group {
993            vars,
994            weighted: Vec::new(),
995        }
996    }
997
998    #[test]
999    fn a_randc_variable_cycles_through_every_value_across_calls() {
1000        // LRM 18.4.2 / Figure 18-1: v0 is a 2-bit randc, v1 a free rand bit.
1001        let t = Ty::unsigned(2);
1002        let sv = Solver::new(&[t, Ty::BIT], &[], &[]).unwrap();
1003        let p = plan(vec![0], vec![uniform(vec![1])]);
1004        let mut states = vec![RandcState::default()];
1005        let mut seed = 31u64;
1006        for cycle in 0..40 {
1007            let mut got: Vec<u32> = (0..4)
1008                .map(|_| {
1009                    let v = sv.sample_plan(&p, &mut states, &mut seed).unwrap();
1010                    u32::try_from(v[0].clone()).unwrap()
1011                })
1012                .collect();
1013            got.sort_unstable();
1014            assert_eq!(got, vec![0, 1, 2, 3], "cycle {cycle}");
1015        }
1016    }
1017
1018    #[test]
1019    fn randc_variables_are_solved_before_every_rand_group() {
1020        // v0 randc (2 bits), v1 rand: `v1 -> v0 == 0`. v0 is drawn first
1021        // (uniformly, every value has a completion), so v1 can only be 1 on
1022        // the quarter of calls where v0 == 0, and there it is 50/50. Had v1
1023        // been drawn first, P(v1 == 1) would be 1/2.
1024        let t = Ty::unsigned(2);
1025        let c = bin(
1026            BinOp::LogImpl,
1027            Expr::Var(1, Ty::BIT),
1028            bin(BinOp::Eq, Expr::Var(0, t), Expr::constant(t, 0u8), Ty::BIT),
1029            Ty::BIT,
1030        );
1031        let sv = Solver::new(&[t, Ty::BIT], &[], &[c]).unwrap();
1032        // Even when `solve v1 before ...` style grouping lists v1 first.
1033        let p = plan(vec![0], vec![uniform(vec![1])]);
1034        let mut states = vec![RandcState::default()];
1035        let mut seed = 32u64;
1036        let (mut v1_ones, mut v0_zero) = (0u32, 0u32);
1037        for _ in 0..4000 {
1038            let v = sv.sample_plan(&p, &mut states, &mut seed).unwrap();
1039            v0_zero += u32::from(v[0].is_zero());
1040            if !v[1].is_zero() {
1041                v1_ones += 1;
1042                assert!(v[0].is_zero(), "{v:?}");
1043            }
1044        }
1045        assert_eq!(v0_zero, 1000, "a randc value occurs once per cycle");
1046        assert!((400..600).contains(&v1_ones), "v1==1: {v1_ones}");
1047    }
1048
1049    #[test]
1050    fn a_randc_variable_whose_legal_set_changes_restarts_its_cycle() {
1051        // v0 randc (3 bits), `v0 < s` with s a state: the legal set is the
1052        // first s values; changing s starts a new permutation (18.4.2).
1053        let t = Ty::unsigned(3);
1054        let c = bin(BinOp::Lt, Expr::Var(0, t), Expr::State(0, t), Ty::BIT);
1055        let p = plan(vec![0], vec![]);
1056        let mut states = vec![RandcState::default()];
1057        let mut seed = 33u64;
1058        let wide =
1059            Solver::new(&[t], &[Some(BigUint::from(4u8))], std::slice::from_ref(&c)).unwrap();
1060        let first = wide.sample_plan(&p, &mut states, &mut seed).unwrap();
1061        assert_eq!(states[0].remaining.len(), 3);
1062        let narrow = Solver::new(&[t], &[Some(BigUint::from(2u8))], &[c]).unwrap();
1063        let mut seen = std::collections::BTreeSet::new();
1064        for _ in 0..2 {
1065            let v = narrow.sample_plan(&p, &mut states, &mut seed).unwrap();
1066            seen.insert(u32::try_from(v[0].clone()).unwrap());
1067        }
1068        assert_eq!(seen, [0, 1].into_iter().collect());
1069        assert!(first[0] < BigUint::from(4u8));
1070    }
1071
1072    #[test]
1073    fn equal_functions_serialise_to_equal_bytes() {
1074        // The randc fingerprint hashes the BDD: two spellings of one
1075        // constraint must agree, a different one must not.
1076        let t = Ty::unsigned(3);
1077        let x = || Expr::Var(0, t);
1078        let k = |n: u8| Expr::constant(t, n);
1079        let a = bin(BinOp::Lt, x(), k(2), Ty::BIT);
1080        let b = bin(
1081            BinOp::LogOr,
1082            bin(BinOp::Eq, x(), k(0), Ty::BIT),
1083            bin(BinOp::Eq, x(), k(1), Ty::BIT),
1084            Ty::BIT,
1085        );
1086        let c = bin(BinOp::Lt, x(), k(3), Ty::BIT);
1087        let bytes = |e: Expr| {
1088            Solver::new(&[t], &[], &[e])
1089                .unwrap()
1090                .compiled()
1091                .bdd
1092                .to_bytes()
1093        };
1094        assert_eq!(bytes(a.clone()), bytes(b));
1095        assert_ne!(bytes(a), bytes(c));
1096    }
1097
1098    #[test]
1099    fn a_weighted_variable_is_drawn_in_proportion_to_its_weights() {
1100        let t = Ty::unsigned(2);
1101        let sv = Solver::new(&[t], &[], &[]).unwrap();
1102        // Weights 1 : 3 : 0, and code 3 has none: only 0 and 1 appear.
1103        let p = plan(
1104            vec![],
1105            vec![Group {
1106                vars: vec![],
1107                weighted: vec![Weighted {
1108                    var: 0,
1109                    weights: ints(&[1, 3, 0]),
1110                    measure: Measure::Flat,
1111                }],
1112            }],
1113        );
1114        let mut seed = 34u64;
1115        let mut hist = [0u32; 4];
1116        for _ in 0..40_000 {
1117            let v = sv.sample_plan(&p, &mut [], &mut seed).unwrap();
1118            hist[usize::try_from(v[0].clone()).unwrap()] += 1;
1119        }
1120        assert_eq!(hist[2] + hist[3], 0, "{hist:?}");
1121        // 10000 / 30000 (sigma 87).
1122        assert!((9_600..10_400).contains(&hist[0]), "{hist:?}");
1123        assert!((29_600..30_400).contains(&hist[1]), "{hist:?}");
1124    }
1125
1126    #[test]
1127    fn a_weighted_variable_only_takes_feasible_codes() {
1128        // Code 1 carries all the weight but is excluded: weights are
1129        // renormalised over what is left (18.5.3: unlisted values excluded).
1130        let t = Ty::unsigned(2);
1131        let c = bin(BinOp::Ne, Expr::Var(0, t), Expr::constant(t, 1u8), Ty::BIT);
1132        let sv = Solver::new(&[t], &[], &[c]).unwrap();
1133        let p = plan(
1134            vec![],
1135            vec![Group {
1136                vars: vec![],
1137                weighted: vec![Weighted {
1138                    var: 0,
1139                    weights: ints(&[1, 100, 1, 0]),
1140                    measure: Measure::Flat,
1141                }],
1142            }],
1143        );
1144        let mut seed = 35u64;
1145        let mut hist = [0u32; 4];
1146        for _ in 0..4_000 {
1147            let v = sv.sample_plan(&p, &mut [], &mut seed).unwrap();
1148            hist[usize::try_from(v[0].clone()).unwrap()] += 1;
1149        }
1150        assert_eq!(hist[1] + hist[3], 0, "{hist:?}");
1151        assert!((1_800..2_200).contains(&hist[0]), "{hist:?}");
1152    }
1153
1154    #[test]
1155    fn weights_with_no_feasible_code_are_unsatisfiable() {
1156        let t = Ty::unsigned(2);
1157        let sv = Solver::new(&[t], &[], &[]).unwrap();
1158        let p = plan(
1159            vec![],
1160            vec![Group {
1161                vars: vec![],
1162                weighted: vec![Weighted {
1163                    var: 0,
1164                    weights: ints(&[0; 4]),
1165                    measure: Measure::Flat,
1166                }],
1167            }],
1168        );
1169        let mut seed = 36u64;
1170        assert_eq!(
1171            sv.sample_plan(&p, &mut [], &mut seed),
1172            Err(SolveError::Unsatisfiable)
1173        );
1174    }
1175
1176    #[test]
1177    fn malformed_plans_are_rejected_loudly() {
1178        let t = Ty::unsigned(2);
1179        let sv = Solver::new(&[t, Ty::unsigned(17)], &[], &[]).unwrap();
1180        let mut seed = 37u64;
1181        // One randc state is missing.
1182        assert!(matches!(
1183            sv.sample_plan(&plan(vec![0], vec![]), &mut [], &mut seed),
1184            Err(SolveError::BadPlan { .. })
1185        ));
1186        // A 17-bit randc exceeds the engine's width.
1187        assert!(matches!(
1188            sv.sample_plan(
1189                &plan(vec![1], vec![]),
1190                &mut [RandcState::default()],
1191                &mut seed
1192            ),
1193            Err(SolveError::BadPlan { .. })
1194        ));
1195        // More weights than the 2-bit variable has values.
1196        let too_many = Group {
1197            vars: vec![],
1198            weighted: vec![Weighted {
1199                var: 0,
1200                weights: ints(&[1; 5]),
1201                measure: Measure::Flat,
1202            }],
1203        };
1204        assert!(matches!(
1205            sv.sample_plan(&plan(vec![], vec![too_many]), &mut [], &mut seed),
1206            Err(SolveError::BadPlan { .. })
1207        ));
1208        // A variable both randc and in a group is a repeat.
1209        assert_eq!(
1210            sv.sample_plan(
1211                &plan(vec![0], vec![uniform(vec![0])]),
1212                &mut [RandcState::default()],
1213                &mut seed
1214            ),
1215            Err(SolveError::BadOrder { var: 0 })
1216        );
1217    }
1218
1219    fn ints(ws: &[u32]) -> Vec<Weight> {
1220        ws.iter().map(|&w| Weight::int(w)).collect()
1221    }
1222
1223    /// `K == i` over a selector variable `k` of type `kt`.
1224    fn is_code(k: u32, kt: Ty, i: u64) -> Expr {
1225        bin(BinOp::Eq, Expr::Var(k, kt), Expr::constant(kt, i), Ty::BIT)
1226    }
1227
1228    /// `lo <= x && x <= hi` for the variable `x`.
1229    fn in_range(x: &Expr, lo: u64, hi: u64) -> Expr {
1230        let t = x.ty();
1231        bin(
1232            BinOp::LogAnd,
1233            bin(BinOp::Ge, x.clone(), Expr::constant(t, lo), Ty::BIT),
1234            bin(BinOp::Le, x.clone(), Expr::constant(t, hi), Ty::BIT),
1235            Ty::BIT,
1236        )
1237    }
1238
1239    /// The conjunct of a `dist` whose item `i` is `items[i]` (inclusive), over
1240    /// the subject `x` and the selector `k` (id 1, 3 bits).
1241    fn dist_conjunct(x: &Expr, items: &[(u64, u64)]) -> (Expr, Ty) {
1242        let kt = Ty::unsigned(3);
1243        let mut acc: Option<Expr> = None;
1244        for (i, &(lo, hi)) in items.iter().enumerate() {
1245            let term = bin(
1246                BinOp::LogAnd,
1247                is_code(1, kt, i as u64),
1248                in_range(x, lo, hi),
1249                Ty::BIT,
1250            );
1251            acc = Some(match acc {
1252                None => term,
1253                Some(p) => bin(BinOp::LogOr, p, term, Ty::BIT),
1254            });
1255        }
1256        (acc.expect("an item"), kt)
1257    }
1258
1259    /// Draws `n` samples of `x` (var 0) from the plan "selector 1 weighted
1260    /// by `weights`, subject 0 distinct" and tallies each value.
1261    fn dist_hist(
1262        xt: Ty,
1263        items: &[(u64, u64)],
1264        weights: Vec<Weight>,
1265        extra: &[Expr],
1266        n: u32,
1267        seed0: u64,
1268    ) -> std::collections::BTreeMap<u32, u32> {
1269        let x = Expr::Var(0, xt);
1270        let (c, kt) = dist_conjunct(&x, items);
1271        let mut conj = vec![c];
1272        conj.extend_from_slice(extra);
1273        let sv = Solver::new(&[xt, kt], &[], &conj).unwrap();
1274        let p = plan(
1275            vec![],
1276            vec![Group {
1277                vars: vec![0],
1278                weighted: vec![Weighted {
1279                    var: 1,
1280                    weights,
1281                    measure: Measure::Distinct(0),
1282                }],
1283            }],
1284        );
1285        let mut seed = seed0;
1286        let mut hist = std::collections::BTreeMap::new();
1287        for _ in 0..n {
1288            let v = sv.sample_plan(&p, &mut [], &mut seed).unwrap();
1289            let x = u32::try_from(v[0].clone()).unwrap();
1290            *hist.entry(x).or_insert(0) += 1;
1291        }
1292        hist
1293    }
1294
1295    #[test]
1296    fn a_dist_with_an_excluded_value_renormalises_over_the_rest() {
1297        // LRM 18.5.3: `x != 200` with `{100 := 1, 200 := 2, 300 := 5}` leaves
1298        // 100 and 300 in the ratio 1 : 5 (N = 36000, 100 expected 6000,
1299        // sigma 71).
1300        let t = Ty::unsigned(9);
1301        let ne = bin(
1302            BinOp::Ne,
1303            Expr::Var(0, t),
1304            Expr::constant(t, 200u32),
1305            Ty::BIT,
1306        );
1307        let h = dist_hist(
1308            t,
1309            &[(100, 100), (200, 200), (300, 300)],
1310            ints(&[1, 2, 5]),
1311            &[ne],
1312            36_000,
1313            41,
1314        );
1315        assert_eq!(
1316            h.keys().copied().collect::<Vec<_>>(),
1317            vec![100, 300],
1318            "{h:?}"
1319        );
1320        assert!((5_600..6_400).contains(&h[&100]), "{h:?}");
1321    }
1322
1323    #[test]
1324    fn a_partly_excluded_range_item_keeps_only_its_feasible_values() {
1325        // `[100:102] := 3, 103 := 1` with `x > 101`: the range keeps only 102
1326        // (weight 3 per value), against 103 (weight 1): 3 : 1. An item-level
1327        // weighting (3 * 3 values vs 1) would give 9 : 1.
1328        let t = Ty::unsigned(8);
1329        let gt = bin(
1330            BinOp::Gt,
1331            Expr::Var(0, t),
1332            Expr::constant(t, 101u32),
1333            Ty::BIT,
1334        );
1335        let h = dist_hist(
1336            t,
1337            &[(100, 102), (103, 103)],
1338            ints(&[3, 1]),
1339            &[gt],
1340            40_000,
1341            42,
1342        );
1343        assert_eq!(
1344            h.keys().copied().collect::<Vec<_>>(),
1345            vec![102, 103],
1346            "{h:?}"
1347        );
1348        // 30000 expected, sigma 87.
1349        assert!((29_600..30_400).contains(&h[&102]), "{h:?}");
1350    }
1351
1352    #[test]
1353    fn overlapping_items_add_their_per_value_weights() {
1354        // `[0:3] := 1, [2:5] := 1`: 2 and 3 are in both, so they weigh 2
1355        // against 1 for 0, 1, 4, 5 (LRM 18.5.3: weights of a value in several
1356        // items add). Total weight 8; N = 48000, expected 6000 and 12000.
1357        let t = Ty::unsigned(4);
1358        let h = dist_hist(t, &[(0, 3), (2, 5)], ints(&[1, 1]), &[], 48_000, 43);
1359        for v in [0, 1, 4, 5] {
1360            assert!((5_600..6_400).contains(&h[&v]), "{v}: {h:?}");
1361        }
1362        for v in [2, 3] {
1363            assert!((11_500..12_500).contains(&h[&v]), "{v}: {h:?}");
1364        }
1365        assert_eq!(h.len(), 6, "{h:?}");
1366    }
1367
1368    #[test]
1369    fn a_range_weight_spread_with_a_denominator_is_per_value() {
1370        // `[0:3] :/ 4` gives each of its four values 4/4 = 1; with `7 := 1`
1371        // all five values are equally likely (N = 40000, 8000 each, sigma 80).
1372        let t = Ty::unsigned(4);
1373        let w = vec![
1374            Weight {
1375                num: 4u32.into(),
1376                den: 4u32.into(),
1377            },
1378            Weight::int(1u32),
1379        ];
1380        let h = dist_hist(t, &[(0, 3), (7, 7)], w, &[], 40_000, 44);
1381        assert_eq!(h.len(), 5, "{h:?}");
1382        for (v, c) in &h {
1383            assert!((7_550..8_450).contains(c), "{v}: {h:?}");
1384        }
1385    }
1386
1387    #[test]
1388    fn the_dist_subject_is_uniform_over_its_values_not_over_completions() {
1389        // x and y are 2 bits with `y <= x`; x = 3 has four completions and
1390        // x = 0 only one. The dist `[0:3] := 1` makes x uniform over its four
1391        // feasible values regardless (N = 40000, 10000 each, sigma 87).
1392        let t = Ty::unsigned(2);
1393        let (x, y) = (Expr::Var(0, t), Expr::Var(2, t));
1394        let (c, kt) = dist_conjunct(&x, &[(0, 3)]);
1395        let le = bin(BinOp::Le, y, x, Ty::BIT);
1396        let sv = Solver::new(&[t, kt, t], &[], &[c, le]).unwrap();
1397        let p = plan(
1398            vec![],
1399            vec![Group {
1400                vars: vec![0, 2],
1401                weighted: vec![Weighted {
1402                    var: 1,
1403                    weights: ints(&[1]),
1404                    measure: Measure::Distinct(0),
1405                }],
1406            }],
1407        );
1408        let mut seed = 45u64;
1409        let (mut hx, mut ymax) = ([0u32; 4], 0u32);
1410        for _ in 0..40_000 {
1411            let v = sv.sample_plan(&p, &mut [], &mut seed).unwrap();
1412            assert!(v[2] <= v[0], "{v:?}");
1413            hx[usize::try_from(v[0].clone()).unwrap()] += 1;
1414            ymax += u32::from(v[0] == BigUint::from(3u8) && v[2] == BigUint::from(3u8));
1415        }
1416        for h in hx {
1417            assert!((9_550..10_450).contains(&h), "{hx:?}");
1418        }
1419        // Given x == 3, y is uniform over 0..=3: a quarter of 10000.
1420        assert!((2_250..2_750).contains(&ymax), "{ymax}");
1421    }
1422
1423    #[test]
1424    fn the_models_measure_weights_a_selector_by_its_completions() {
1425        // The same problem with `Measure::Models`: code 0 stands for each
1426        // (x, y) tuple, so x is weighted by its completions 1 : 2 : 3 : 4
1427        // (N = 40000, x == 3 expected 16000, sigma 98).
1428        let t = Ty::unsigned(2);
1429        let (x, y) = (Expr::Var(0, t), Expr::Var(2, t));
1430        let (c, kt) = dist_conjunct(&x, &[(0, 3)]);
1431        let le = bin(BinOp::Le, y, x, Ty::BIT);
1432        let sv = Solver::new(&[t, kt, t], &[], &[c, le]).unwrap();
1433        let p = plan(
1434            vec![],
1435            vec![Group {
1436                vars: vec![0, 2],
1437                weighted: vec![Weighted {
1438                    var: 1,
1439                    weights: ints(&[1]),
1440                    measure: Measure::Models,
1441                }],
1442            }],
1443        );
1444        let mut seed = 46u64;
1445        let mut x3 = 0u32;
1446        for _ in 0..40_000 {
1447            let v = sv.sample_plan(&p, &mut [], &mut seed).unwrap();
1448            x3 += u32::from(v[0] == BigUint::from(3u8));
1449        }
1450        assert!((15_500..16_500).contains(&x3), "{x3}");
1451    }
1452
1453    #[test]
1454    fn a_selector_with_no_positive_feasible_weight_is_unsatisfiable() {
1455        // Every item has weight 0: nothing can be chosen (LRM 18.5.3: a
1456        // zero-weight value is never selected).
1457        let t = Ty::unsigned(4);
1458        let x = Expr::Var(0, t);
1459        let (c, kt) = dist_conjunct(&x, &[(0, 3), (5, 5)]);
1460        let sv = Solver::new(&[t, kt], &[], &[c]).unwrap();
1461        let p = plan(
1462            vec![],
1463            vec![Group {
1464                vars: vec![0],
1465                weighted: vec![Weighted {
1466                    var: 1,
1467                    weights: ints(&[0, 0]),
1468                    measure: Measure::Distinct(0),
1469                }],
1470            }],
1471        );
1472        let mut seed = 47u64;
1473        assert_eq!(
1474            sv.sample_plan(&p, &mut [], &mut seed),
1475            Err(SolveError::Unsatisfiable)
1476        );
1477        assert_eq!(seed, 47, "no draws consumed");
1478    }
1479
1480    #[test]
1481    fn a_distinct_subject_outside_the_group_or_a_zero_denominator_is_a_bad_plan() {
1482        let t = Ty::unsigned(4);
1483        let x = Expr::Var(0, t);
1484        let (c, kt) = dist_conjunct(&x, &[(0, 3)]);
1485        let sv = Solver::new(&[t, kt], &[], &[c]).unwrap();
1486        let mut seed = 48u64;
1487        let weighted = |w: Vec<Weight>| Weighted {
1488            var: 1,
1489            weights: w,
1490            measure: Measure::Distinct(0),
1491        };
1492        // The subject is not listed in the selector's group.
1493        let p = plan(
1494            vec![],
1495            vec![Group {
1496                vars: vec![],
1497                weighted: vec![weighted(ints(&[1]))],
1498            }],
1499        );
1500        assert!(matches!(
1501            sv.sample_plan(&p, &mut [], &mut seed),
1502            Err(SolveError::BadPlan { .. })
1503        ));
1504        // A weight whose denominator is 0.
1505        let bad = vec![Weight {
1506            num: 1u32.into(),
1507            den: 0u32.into(),
1508        }];
1509        let p = plan(
1510            vec![],
1511            vec![Group {
1512                vars: vec![0],
1513                weighted: vec![weighted(bad)],
1514            }],
1515        );
1516        assert!(matches!(
1517            sv.sample_plan(&p, &mut [], &mut seed),
1518            Err(SolveError::BadPlan { .. })
1519        ));
1520    }
1521
1522    /// Bits outside the constraints' support are drawn without conjoining a
1523    /// literal (`sample_plan`); the draw must be exactly the one the per-bit
1524    /// walk over every bit makes. The expected values were produced by the
1525    /// walk before that shortcut existed.
1526    #[test]
1527    fn unconstrained_bits_draw_as_the_full_walk_does() {
1528        let x = Expr::Var(0, Ty::unsigned(8));
1529        let c = bin(
1530            BinOp::Lt,
1531            x,
1532            Expr::constant(Ty::unsigned(8), 100u8),
1533            Ty::BIT,
1534        );
1535        let tys = [Ty::unsigned(8), Ty::unsigned(8), Ty::unsigned(4)];
1536        let sv = Solver::new(&tys, &[], &[c]).expect("compiles");
1537        let mut seed = 7;
1538        let mut got = Vec::new();
1539        for _ in 0..6 {
1540            let v = sv.sample(&[1, 0, 2], &mut seed).expect("sat");
1541            assert!(v[0] < BigUint::from(100u8));
1542            got.push(
1543                v.iter()
1544                    .map(ToString::to_string)
1545                    .collect::<Vec<_>>()
1546                    .join(","),
1547            );
1548        }
1549        assert_eq!(
1550            got.join(" "),
1551            "21,216,3 52,172,1 61,93,5 42,49,7 18,1,0 31,165,6"
1552        );
1553    }
1554
1555    /// A repeated compile is served from the cache with the same count.
1556    #[test]
1557    fn repeated_compile_is_cached_and_equal() {
1558        let x = Expr::Var(0, Ty::unsigned(6));
1559        let c = bin(BinOp::Lt, x, Expr::constant(Ty::unsigned(6), 37u8), Ty::BIT);
1560        let tys = [Ty::unsigned(6)];
1561        let a = Solver::new(&tys, &[], std::slice::from_ref(&c)).unwrap();
1562        let b = Solver::new(&tys, &[], std::slice::from_ref(&c)).unwrap();
1563        assert_eq!(a.count(), &BigUint::from(37u32));
1564        assert_eq!(a.count(), b.count());
1565        let digest = problem_digest(&tys, &[], &[c], NODE_BUDGET);
1566        assert!(SOLVER_CACHE.lock().unwrap().contains_key(&digest));
1567    }
1568}