Skip to main content

probl_engine/
interp.rs

1//! The world-set interpreter.
2//!
3//! Every statement maps a set of worlds to a set of worlds. All the worlds in
4//! a set are at the same point of the program, so control flow is shared and
5//! only the data differs. Splitting happens in statements; expressions are
6//! evaluated in one world at a time. See docs/semantics.md for the rules this
7//! implements.
8
9use crate::analytic::{self, Analytic};
10use crate::builtins;
11use crate::chain::{Chain, Solution};
12use crate::conjugate::{self, Seen};
13use crate::continuous::{Family, Rng};
14use crate::dist::{Budget, Counts, Dist};
15use crate::error::{ErrorKind, Fault, OpError, OpResult, Result, RuntimeError};
16use crate::failure::Failures;
17use crate::ops::{self, Truth};
18use crate::report::Sink;
19use crate::value::{Closure, Delayed, Value, fmt_prob};
20use crate::weight::Weight;
21use crate::world::Returned;
22use crate::world::{Flow, World, clear, clear_dead, live_slots, merge, merge_values, state_hash, total_weight};
23use probl_sema::builtins::Lifting;
24use probl_sema::conjugate::{Conjugacy, Likelihood, Update};
25use probl_sema::effects::EvidenceOrder;
26use probl_sema::ir::*;
27use probl_sema::{Builtin, Liveness};
28use probl_syntax::Span;
29use probl_syntax::ast::BinOp;
30use rustc_hash::{FxHashMap, FxHashSet};
31use std::collections::BTreeMap;
32use std::sync::Arc;
33use std::sync::atomic::{AtomicBool, Ordering};
34
35#[derive(Clone, Debug, Default)]
36pub struct Stats {
37    /// Largest number of worlds any statement ran on.
38    pub peak_worlds: usize,
39    /// Statement executions, counting each world separately.
40    pub world_steps: u64,
41    pub calls: u64,
42    pub memo_hits: u64,
43    /// Loops solved as Markov chains, and their states (section 10).
44    pub solved_loops: u64,
45    pub chain_states: u64,
46    /// Calls that came back to themselves, solved by iteration, and the
47    /// rounds that took (section 6).
48    pub solved_calls: u64,
49    pub call_rounds: u64,
50    /// When sampling: what happened to each variable whose draws may be
51    /// delayed (`Conjugacy::variables`).
52    pub updates: Vec<Updates>,
53    /// When the program was answered with decision diagrams instead of
54    /// worlds (`Options::symbolic`): their size.
55    pub diagrams: Option<crate::symbolic::Diagrams>,
56}
57
58/// How often a variable's draws were delayed and updated exactly
59/// (docs/semantics.md, section 14), counted over all runs.
60#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
61pub struct Updates {
62    /// Draws delayed.
63    pub delayed: u64,
64    /// Observations that updated it exactly.
65    pub exact: u64,
66    /// Times it was drawn when first needed, from its updated distribution.
67    pub drawn: u64,
68}
69
70impl Stats {
71    /// Add up another batch's statistics.
72    pub fn absorb(&mut self, other: &Stats) {
73        self.peak_worlds = self.peak_worlds.max(other.peak_worlds);
74        self.world_steps += other.world_steps;
75        self.calls += other.calls;
76        self.memo_hits += other.memo_hits;
77        self.solved_loops += other.solved_loops;
78        self.chain_states += other.chain_states;
79        self.solved_calls += other.solved_calls;
80        self.call_rounds += other.call_rounds;
81        if self.updates.len() < other.updates.len() {
82            self.updates.resize(other.updates.len(), Updates::default());
83        }
84        for (mine, theirs) in self.updates.iter_mut().zip(&other.updates) {
85            mine.delayed += theirs.delayed;
86            mine.exact += theirs.exact;
87            mine.drawn += theirs.drawn;
88        }
89    }
90}
91
92/// How the engine runs, with the host's limits already applied.
93#[derive(Clone, Debug)]
94pub struct Config {
95    pub today: Option<i32>,
96    /// The host's environment snapshot, for `env`.
97    pub env: Arc<BTreeMap<String, String>>,
98    pub epsilon: f64,
99    pub merging: bool,
100    pub memoizing: bool,
101    pub max_worlds: usize,
102    pub max_iterations: u64,
103    pub max_call_depth: usize,
104    pub max_cached_calls: usize,
105    pub max_output: usize,
106    pub budget: Budget,
107    pub cancel: Option<Arc<AtomicBool>>,
108    /// Sample with this seed, instead of enumerating (docs/semantics.md,
109    /// section 14).
110    pub sample_seed: Option<u64>,
111    /// When sampling, delay the draws of conjugate priors and update them
112    /// exactly (section 14).
113    pub conjugate: bool,
114    /// When enumerating, solve loops that cycle as Markov chains (section 10).
115    pub solving: bool,
116    /// The most states a loop's chain may have to be solved.
117    pub max_chain_states: usize,
118    /// A fault ends only the world it happens in (docs/semantics.md,
119    /// section 11), instead of the run.
120    pub partial: bool,
121    /// Compute what has no formula by numerical integration, to this
122    /// relative tolerance (section 13).
123    pub integrate: Option<f64>,
124}
125
126/// The most steps eliminating one loop's chain may take before the loop is
127/// unrolled instead.
128const MAX_ELIMINATION: u64 = 200_000_000;
129
130/// Runs are executed this many at a time, and each batch has a random
131/// stream of its own. The size is fixed: changing it changes the output for
132/// a given seed.
133pub const BATCH: u64 = 1000;
134
135/// What a batch of sampled runs produced, to be combined with the other
136/// batches in order.
137#[derive(Debug)]
138pub struct Batch {
139    pub sinks: Vec<Sink>,
140    pub totals: SampleTotals,
141    pub observed: bool,
142    pub densities: bool,
143    pub unresolved: Weight,
144    pub last_ruling_out: Option<Span>,
145    pub stats: Stats,
146    pub failures: Failures,
147}
148
149/// Lines printed by `print(…)`, with where.
150pub type Printed = Vec<(Span, String)>;
151
152/// The final weights of all the runs, when sampling.
153#[derive(Clone, Copy, Debug)]
154pub struct SampleTotals {
155    pub weight: Weight,
156    /// Sum of the squared weights, for the effective sample size.
157    pub squares: Weight,
158}
159
160/// The distribution of a function's return value, unnormalized.
161#[derive(Clone, Debug)]
162pub struct CallResult {
163    pub outcomes: Vec<Returned>,
164    /// Weight the call left unresolved.
165    pub unresolved: Weight,
166    /// Weight its observations ruled out, when enumerating.
167    pub lost: Weight,
168    /// Whether the call ran `observe` (outside any `simulate`).
169    pub observed: bool,
170    /// The part of `unresolved` waiting on calls that came back to
171    /// themselves and are still being solved (section 6), for each by its
172    /// depth in the stack. A result waiting on one isn't final.
173    pub pending: Vec<(usize, Weight)>,
174    /// In partial mode: its worlds that failed.
175    pub failures: Failures,
176}
177
178impl CallResult {
179    /// The result of a call that came back to itself, before the first
180    /// round: nothing resolved, all of it waiting on that call.
181    fn waiting(depth: usize) -> CallResult {
182        CallResult {
183            outcomes: Vec::new(),
184            unresolved: Weight::ONE,
185            lost: Weight::ZERO,
186            observed: false,
187            pending: vec![(depth, Weight::ONE)],
188            failures: Failures::default(),
189        }
190    }
191
192    /// Take out the weight waiting on the call at `depth`.
193    fn take_pending(&mut self, depth: usize) -> Weight {
194        match self.pending.iter().position(|&(d, _)| d == depth) {
195            Some(i) => self.pending.remove(i).1,
196            None => Weight::ZERO,
197        }
198    }
199}
200
201/// A call's result so far, while the call it waits on is being solved
202/// (section 6).
203struct Approx {
204    result: Arc<CallResult>,
205    /// The depth of that call, and its frame and round when this was worked
206    /// out. It stands for the call as long as that frame runs, and can be
207    /// reused as it is until that frame's next round.
208    head: usize,
209    frame: u64,
210    round: u64,
211}
212
213/// Add `weight` times what's waiting in `from` to `into`.
214fn add_pending(into: &mut Vec<(usize, Weight)>, from: &[(usize, Weight)], weight: Weight) {
215    for &(depth, w) in from {
216        match into.iter_mut().find(|(d, _)| *d == depth) {
217            Some((_, total)) => *total += weight * w,
218            None => into.push((depth, weight * w)),
219        }
220    }
221}
222
223type CallKey = (FnId, Vec<Value>);
224
225pub struct Engine<'p> {
226    prog: &'p Program,
227    live: &'p Liveness,
228    /// Which draws may be delayed, and what must draw them.
229    conj: &'p Conjugacy,
230    config: Config,
231    budget: Budget,
232    memo: FxHashMap<CallKey, Arc<CallResult>>,
233    /// `simulate` results that are distributions whose outcomes own their
234    /// draws: pure values of the block and its captures, which a call's
235    /// result, holding analytic outcomes, can't be cached as.
236    simulated: FxHashMap<CallKey, Value>,
237    /// The calls running, with their depth in the stack.
238    active: FxHashMap<CallKey, usize>,
239    /// Results so far of calls that wait on a call being solved by
240    /// iteration (section 6).
241    approx: FxHashMap<CallKey, Approx>,
242    /// For each depth in the stack: the frame running there, and its round,
243    /// each a number never used before.
244    frames_at: Vec<u64>,
245    rounds_at: Vec<u64>,
246    last_number: u64,
247    /// Weight waiting on calls being solved, in the current call.
248    pending: Vec<(usize, Weight)>,
249    depth: usize,
250    /// Weight left unaccounted for (loops cut short, truncated tails).
251    pub unresolved: Weight,
252    /// When enumerating: weight ruled out by observations, which solving a
253    /// loop counts as a way out of it.
254    lost: Weight,
255    /// For each statement: whether it's an unbounded loop that may be
256    /// solved as a Markov chain (section 10).
257    solvable: Vec<bool>,
258    /// For each loop: the variables it only adds to, which the chain carries
259    /// as rewards rather than in its states.
260    rewards: FxHashMap<StmtId, Vec<SlotId>>,
261    /// Loops whose chains were found too large to solve: they're unrolled
262    /// from then on.
263    too_large: FxHashSet<StmtId>,
264    /// Whether the current inference scope has run `observe`.
265    pub observed: bool,
266    /// Whether it has observed a value with a density, which makes the
267    /// evidence a density too (docs/semantics.md, section 14).
268    pub densities: bool,
269    /// The last observation that ruled out a world, for impossible evidence.
270    pub last_ruling_out: Option<Span>,
271    /// When enumerating, the first observation of a density.
272    pub density_at: Option<Span>,
273    /// How many densities the weights of the worlds that reported or
274    /// finished include. It must be as many for each.
275    pub reported_densities: Option<u32>,
276    /// In the innermost `simulate` running: its first observation of a
277    /// density, and how many its returning worlds' weights include, which
278    /// must be as many for each, since its result is normalized.
279    local_density_at: Option<Span>,
280    returned_densities: Option<u32>,
281    /// The tolerance, if a result is computed by numerical integration.
282    pub integrated: Option<f64>,
283    pub sinks: Vec<Sink>,
284    dice: FxHashMap<(u32, u32), Value>,
285    pools: FxHashMap<(u32, Value), Value>,
286    record_names: Vec<Arc<str>>,
287    enum_values: Vec<Vec<Value>>,
288    print: &'p mut (dyn FnMut(&str) + Send),
289    /// The values of the program's inputs (`read`).
290    inputs: &'p [Value],
291    /// Bytes printed, counting the batches printed before this one.
292    printed: usize,
293    /// When sampling a batch: what it printed, kept to be printed in batch
294    /// order.
295    lines: Option<Printed>,
296    pub stats: Stats,
297    /// When sampling: the random numbers. `None` when enumerating, including
298    /// inside `simulate` while sampling.
299    sampler: Option<Rng>,
300    /// Nesting depth of local `simulate` inference (in either outer mode).
301    nested: usize,
302    next_latent: u64,
303    /// `env` as a map, once it's been read: each engine makes its own, so
304    /// that sampling threads don't share its reference counts.
305    env: Option<Value>,
306    /// Dynamic effect boundary for collection callbacks. Effects inside an
307    /// explicit `simulate` are local to that computation instead.
308    callback: Option<&'static str>,
309    /// In partial mode: the worlds that failed, in the current call.
310    pub failures: Failures,
311    /// The `try`s running, innermost last, each with the depth of the call
312    /// it's in, and its catches.
313    handlers: Vec<(usize, &'p [Catch])>,
314    /// Worlds of the statement running that faulted inside a `try` of this
315    /// call that catches the fault, on their way to it.
316    faulted: Vec<(World, RuntimeError)>,
317    /// Where evidence can still be applied, to tell whether a failed
318    /// world's weight can be added to the finished worlds'.
319    evidence: EvidenceOrder,
320}
321
322/// Where a world was when it failed.
323#[derive(Clone, Copy)]
324struct At {
325    f: FnId,
326    stmt: StmtId,
327    weight: Weight,
328    run: u32,
329    densities: u32,
330}
331
332impl At {
333    fn new(f: FnId, stmt: &Stmt, w: &World) -> At {
334        At {
335            f,
336            stmt: stmt.id,
337            weight: w.weight,
338            run: w.run,
339            densities: w.densities,
340        }
341    }
342}
343
344/// The value of `$e` in one world. If it fails with a fault that a `catch`
345/// or partial mode takes, the world leaves the statement: `$undo` runs, and
346/// the statement goes on with the next world (see `Engine::fail`). A catch
347/// gets the world as its statement began: `world w` when `$e` leaves `w` as
348/// it was, or `saved s` with a copy `s` taken before `$e` changed it.
349macro_rules! each {
350    ($self:ident, $at:expr, world $w:expr, $e:expr $(, $undo:block)?) => {
351        match $e {
352            Ok(v) => v,
353            Err(e) => {
354                let saved = $self.snapshot(&$w);
355                $self.fail(e, $at, saved)?;
356                $($undo)?
357                continue;
358            }
359        }
360    };
361    ($self:ident, $at:expr, saved $saved:ident, $e:expr $(, $undo:block)?) => {
362        match $e {
363            Ok(v) => v,
364            Err(e) => {
365                $self.fail(e, $at, $saved.take())?;
366                $($undo)?
367                continue;
368            }
369        }
370    };
371}
372
373impl<'p> Engine<'p> {
374    pub fn new(
375        prog: &'p Program,
376        live: &'p Liveness,
377        conj: &'p Conjugacy,
378        config: Config,
379        inputs: &'p [Value],
380        print: &'p mut (dyn FnMut(&str) + Send),
381    ) -> Engine<'p> {
382        let sampler = config.sample_seed.map(Rng::new);
383        Engine {
384            prog,
385            live,
386            conj,
387            budget: config.budget.clone(),
388            config,
389            memo: FxHashMap::default(),
390            simulated: FxHashMap::default(),
391            active: FxHashMap::default(),
392            approx: FxHashMap::default(),
393            frames_at: Vec::new(),
394            rounds_at: Vec::new(),
395            last_number: 0,
396            pending: Vec::new(),
397            depth: 0,
398            unresolved: Weight::ZERO,
399            lost: Weight::ZERO,
400            solvable: probl_sema::effects::solvable_loops(prog),
401            rewards: probl_sema::rewards::loop_rewards(prog),
402            too_large: FxHashSet::default(),
403            observed: false,
404            densities: false,
405            last_ruling_out: None,
406            density_at: None,
407            reported_densities: None,
408            local_density_at: None,
409            returned_densities: None,
410            integrated: None,
411            sinks: vec![Sink::default(); prog.reports.len()],
412            dice: FxHashMap::default(),
413            pools: FxHashMap::default(),
414            record_names: prog.records.iter().map(|r| Arc::from(r.name.as_str())).collect(),
415            enum_values: prog
416                .enums
417                .iter()
418                .enumerate()
419                .map(|(t, e)| {
420                    e.variants
421                        .iter()
422                        .enumerate()
423                        .map(|(v, name)| ops::enum_value(t as u32, v as u32, name))
424                        .collect()
425                })
426                .collect(),
427            print,
428            inputs,
429            printed: 0,
430            lines: None,
431            stats: Stats::default(),
432            sampler,
433            nested: 0,
434            next_latent: 0,
435            env: None,
436            callback: None,
437            failures: Failures::default(),
438            evidence: probl_sema::effects::evidence_order(prog),
439            handlers: Vec::new(),
440            faulted: Vec::new(),
441        }
442    }
443
444    /// Run the top level. Returns the total weight of the worlds that finish.
445    pub fn run_main(&mut self) -> Result<Weight> {
446        let prog = self.prog;
447        let main = prog.main();
448        let world = World {
449            slots: vec![Value::Dead; main.n_slots()],
450            constraints: Default::default(),
451            inherited: Default::default(),
452            weight: Weight::ONE,
453            run: 0,
454            densities: 0,
455        };
456        let flow = self.exec_block(MAIN, &main.body, vec![world])?;
457        uncaught(&flow)?;
458        for w in &flow.next {
459            self.same_units(w, main.span)?;
460        }
461        Ok(total_weight(&flow.next))
462    }
463
464    /// Weights with different numbers of densities aren't in the same
465    /// units, so they can't be added up or compared.
466    fn same_units(&mut self, w: &World, span: Span) -> Result<()> {
467        match self.reported_densities {
468            Some(n) if n != w.densities => Err(OpError::unsupported(
469                "worlds that observe different numbers of continuous values aren't supported when enumerating",
470            )
471            .help("a density is per unit of what it observes, so worlds with more of them can't be weighed against the others; observe the same values in every world")
472            .at(self.density_at.unwrap_or(span))),
473            _ => {
474                self.reported_densities = Some(w.densities);
475                Ok(())
476            }
477        }
478    }
479
480    /// Sample the runs `first..first + n` with the random numbers `rng`
481    /// (docs/semantics.md, section 14), after earlier batches printed
482    /// `printed` bytes. A batch starts afresh: its reports and counters are
483    /// its own, so batches can run on any thread; only caches carry over.
484    /// What it prints comes back with its result, even when it fails.
485    pub fn run_batch(&mut self, rng: Rng, first: u64, n: u64, printed: usize) -> (Result<Batch>, Printed) {
486        self.sampler = Some(rng);
487        self.sinks = vec![Sink::default(); self.prog.reports.len()];
488        self.observed = false;
489        self.densities = false;
490        self.unresolved = Weight::ZERO;
491        self.last_ruling_out = None;
492        self.stats = Stats::default();
493        self.failures = Failures::default();
494        self.printed = printed;
495        self.lines = Some(Vec::new());
496        self.budget = self.config.budget.clone();
497        let result = self.sample_runs(first, n);
498        self.budget.give_back();
499        (result, self.lines.take().unwrap_or_default())
500    }
501
502    fn sample_runs(&mut self, first: u64, n: u64) -> Result<Batch> {
503        let main = self.prog.main();
504        let worlds = (first..first + n)
505            .map(|run| World {
506                slots: vec![Value::Dead; main.n_slots()],
507                constraints: Default::default(),
508                inherited: Default::default(),
509                weight: Weight::ONE,
510                run: run as u32,
511                densities: 0,
512            })
513            .collect();
514        let flow = self.exec_block(MAIN, &main.body, worlds)?;
515        uncaught(&flow)?;
516        let mut totals = SampleTotals {
517            weight: Weight::ZERO,
518            squares: Weight::ZERO,
519        };
520        for w in &flow.next {
521            totals.weight += w.weight;
522            totals.squares += w.weight * w.weight;
523        }
524        for sink in &mut self.sinks {
525            sink.end_batch();
526        }
527        Ok(Batch {
528            sinks: std::mem::take(&mut self.sinks),
529            totals,
530            observed: self.observed,
531            densities: self.densities,
532            unresolved: self.unresolved,
533            last_ruling_out: self.last_ruling_out,
534            stats: std::mem::take(&mut self.stats),
535            failures: std::mem::take(&mut self.failures),
536        })
537    }
538
539    /// When sampling: one draw from a distribution. A finite distribution
540    /// gives one of its outcomes, chosen with its probability (missing mass
541    /// is ignored); a continuous one, a number.
542    fn sample(&mut self, v: &Value) -> Value {
543        let rng = self.sampler.as_mut().expect("only called when sampling");
544        match v {
545            Value::Dist(d) if !d.outcomes.is_empty() => {
546                let i = rng.choose(d.outcomes.iter().map(|(_, p)| *p)).unwrap_or(0);
547                match &d.outcomes[i].0 {
548                    Value::Continuous(f) => Value::Float(f.sample(rng)),
549                    Value::Joint(j) => crate::joint::sample(j, rng),
550                    x => x.clone(),
551                }
552            }
553            Value::Continuous(f) => Value::Float(f.sample(rng)),
554            Value::Counts(c) => Value::Int(c.sample(rng).into()),
555            Value::Joint(j) => crate::joint::sample(j, rng),
556            other => other.clone(),
557        }
558    }
559
560    /// When sampling, a value to report: continuous parts are drawn; a
561    /// finite distribution is kept whole, so each outcome counts with its
562    /// probability.
563    #[inline(never)]
564    fn sample_continuous(&mut self, v: Value) -> Value {
565        match &v {
566            Value::Continuous(_) | Value::Joint(_) => self.sample(&v),
567            Value::Dist(d)
568                if d.outcomes
569                    .iter()
570                    .any(|(x, _)| matches!(x, Value::Continuous(_) | Value::Joint(_))) =>
571            {
572                let pairs = d.outcomes.iter().map(|(x, p)| (self.sample(x), *p)).collect();
573                Dist::from_pairs(pairs, d.missing).into_value()
574            }
575            _ => v,
576        }
577    }
578
579    /// When sampling, `binomial(…)`, `poisson(…)` or `geometric(…)` with plain
580    /// arguments is drawn from, or observed, without listing its outcomes:
581    /// the distribution is the same, and it's much faster.
582    fn direct_counts(&mut self, f: FnId, e: &'p Expr, w: &World) -> Result<Option<Counts>> {
583        let ExprKind::Builtin {
584            func: b @ (Builtin::Binomial | Builtin::Poisson | Builtin::Geometric),
585            args,
586            ..
587        } = &e.kind
588        else {
589            return Ok(None);
590        };
591        if self.sampler.is_none() {
592            return Ok(None);
593        }
594        let mut values = Vec::with_capacity(args.len());
595        for a in args {
596            values.push(self.eval(f, a, w)?);
597        }
598        let counts = builtins::counts(*b, &values, &mut self.budget).map_err(|err| err.at(e.span))?;
599        Ok(counts.filter(Counts::direct))
600    }
601
602    /// Whether draws of conjugate priors are delayed now: when sampling,
603    /// but not inside `simulate` (section 14).
604    fn delaying(&self) -> bool {
605        self.config.conjugate && self.sampler.is_some()
606    }
607
608    fn updates(&mut self, variable: u32) -> &mut Updates {
609        let i = variable as usize;
610        if self.stats.updates.len() <= i {
611            let n = self.conj.variables.len().max(i + 1);
612            self.stats.updates.resize(n, Updates::default());
613        }
614        &mut self.stats.updates[i]
615    }
616
617    /// Draw a delayed variable from its distribution, updated by the
618    /// observations so far, now that its value is needed.
619    fn draw_delayed(&mut self, w: &mut World, slot: SlotId) {
620        if let Value::Delayed(d) = &w.slots[slot as usize] {
621            let d = **d;
622            let x = d
623                .family
624                .sample(self.sampler.as_mut().expect("delayed only when sampling"));
625            w.slots[slot as usize] = Value::Float(x);
626            self.updates(d.variable).drawn += 1;
627        }
628    }
629
630    /// `observe` of a delayed variable through a conjugate form: the log of
631    /// the observation's probability (or density), with the variable
632    /// updated to its distribution after it. `None` when it isn't an exact
633    /// update: the variable isn't delayed, or doesn't have the right family,
634    /// or the rest of the observation isn't plain. Then the variable is
635    /// drawn, and the observation is made as usual. `from` is where
636    /// problems with the observed distribution's parameters are reported,
637    /// as when it's evaluated.
638    fn observe_exactly(&mut self, f: FnId, u: &Update<'p>, w: &mut World, from: Span) -> Result<Option<f64>> {
639        let Value::Delayed(d) = &w.slots[u.slot as usize] else {
640            return Ok(None);
641        };
642        let d = **d;
643        let pair = matches!(
644            (d.family, u.likelihood),
645            (
646                Family::Beta { .. },
647                Likelihood::Binomial { .. } | Likelihood::Bernoulli { .. }
648            ) | (Family::Gamma { .. }, Likelihood::Poisson { .. })
649                | (Family::Normal { .. }, Likelihood::Normal { .. })
650        );
651        let seen = if pair { self.seen(f, u, w, from)? } else { None };
652        let Some(seen) = seen else {
653            self.draw_delayed(w, u.slot);
654            return Ok(None);
655        };
656        let (ln, posterior) = conjugate::update(&d.family, seen).expect("a conjugate pair");
657        w.slots[u.slot as usize] = Value::Delayed(Arc::new(Delayed {
658            family: posterior,
659            variable: d.variable,
660        }));
661        self.updates(d.variable).exact += 1;
662        self.densities |= matches!(seen, Seen::Normal { .. });
663        Ok(Some(ln))
664    }
665
666    /// Assigning a discrete outcome, like `floor(x)`, splits its world: one
667    /// for each of its values, with the draw restricted to where it's that
668    /// value, so that the variable is a plain number. The split worlds go
669    /// to `failed`, to be dropped, and theirs to the end. (A place with a
670    /// path keeps the outcome.)
671    #[cold]
672    #[inline(never)]
673    fn split_discrete(
674        &mut self,
675        f: FnId,
676        stmt: &'p Stmt,
677        place: &Place,
678        worlds: &mut Vec<World>,
679        failed: &mut Vec<usize>,
680    ) -> Result<()> {
681        if !place.path.is_empty() {
682            return Ok(());
683        }
684        let mut split = Vec::new();
685        let mut gone = Vec::new();
686        for (i, w) in worlds.iter().enumerate() {
687            let Value::Analytic(a) = &w.slots[place.slot as usize] else {
688                continue;
689            };
690            // `failed` is in order: the worlds that failed to assign.
691            if !a.is_discrete() || failed.binary_search(&i).is_ok() {
692                continue;
693            }
694            if a.form.is_some() {
695                return Err(
696                    analytic::joint_condition("assigning a rounded outcome of several normal draws").at(stmt.span),
697                );
698            }
699            for (v, region, share) in a.atoms() {
700                if share <= 0.0 {
701                    continue;
702                }
703                let mut nw = w.clone().scaled(share);
704                let latent = analytic::Latent {
705                    family: a.family,
706                    domain: a.domain.intersect(&region),
707                    form: None,
708                    weight: a.weight.clone(),
709                };
710                Arc::make_mut(&mut nw.constraints).insert(a.id, latent);
711                if let Err(e) = self.assign(f, place, v, &mut nw, stmt.span) {
712                    self.fail(e, At::new(f, stmt, &nw), None)?;
713                    continue;
714                }
715                split.push(nw);
716            }
717            gone.push(i);
718        }
719        failed.extend(gone);
720        failed.sort_unstable();
721        worlds.extend(split);
722        self.check_worlds(worlds.len() - failed.len(), stmt.span)
723    }
724
725    /// `if p` for a probability `p` drawn from a beta: each branch updates
726    /// it, as observing that side from `bernoulli(p)` would. Whether it was
727    /// such a `p`. (Out of `exec_stmt_kind`, like the next two, which every
728    /// statement runs: inlined, they'd make its frame larger.)
729    #[inline(never)]
730    fn branch_on_draw(
731        &mut self,
732        x: &Analytic,
733        w: &World,
734        yes: &mut Vec<World>,
735        no: &mut Vec<World>,
736        span: Span,
737    ) -> Result<bool> {
738        let sides = [Seen::Bernoulli(true), Seen::Bernoulli(false)];
739        let mut sides = self.outcomes_of(x, w, &sides, span)?.into_iter();
740        let (Some(y), Some(n)) = (sides.next(), sides.next()) else {
741            return Ok(false);
742        };
743        self.check_callback_effect(span)?;
744        yes.extend(y);
745        no.extend(n);
746        Ok(true)
747    }
748
749    /// A draw whose distribution depends on an analytic draw, into `out`:
750    /// from `normal(mean, sd)` with an outcome of normal draws as its mean
751    /// (see `gaussian_draw`), or from `bernoulli(p)` or `binomial(n, p)`
752    /// with a probability drawn from a beta (see `dependent_draw`). Whether
753    /// it was one.
754    #[inline(never)]
755    fn draw_dependent(
756        &mut self,
757        f: FnId,
758        place: &Place,
759        dist: &'p Expr,
760        w: &World,
761        out: &mut Vec<World>,
762    ) -> Result<bool> {
763        if let Some(v) = self.gaussian_draw(f, dist, w)? {
764            let mut w = w.clone();
765            self.assign(f, place, v, &mut w, dist.span)?;
766            out.push(w);
767            return Ok(true);
768        }
769        let Some(outcomes) = self.dependent_draw(f, dist, w)? else {
770            return Ok(false);
771        };
772        let mark = out.len();
773        for (v, mut w) in outcomes {
774            if let Err(e) = self.assign(f, place, v, &mut w, dist.span) {
775                out.truncate(mark);
776                return Err(e);
777            }
778            out.push(w);
779        }
780        Ok(true)
781    }
782
783    /// Update an analytic draw exactly when enumerating: the observation's
784    /// probability, as a logarithm, or `None` if it isn't an update this
785    /// can make.
786    #[inline(never)]
787    fn observe_posterior(&mut self, f: FnId, u: &Update<'p>, w: &mut World, from: Span) -> Result<Option<f64>> {
788        let Some(x) = self.draw_in(w, u.slot, from)? else {
789            return Ok(None);
790        };
791        let Some(seen) = self.seen(f, u, w, from)? else {
792            return Ok(None);
793        };
794        let Some((ln, latent)) = analytic::update(&x, seen).map_err(|e| e.at(from))? else {
795            return Ok(None);
796        };
797        if matches!(seen, Seen::Normal { .. }) {
798            self.check_density(f, from)?;
799            self.densities = true;
800            w.densities += 1;
801        }
802        self.know(w, x.id, latent, from)?;
803        Ok(Some(ln))
804    }
805
806    /// A draw from `normal(mean, sd)` when enumerating, with `mean` an
807    /// outcome of normal draws: the mean plus a fresh normal draw. `None`
808    /// for another draw.
809    fn gaussian_draw(&mut self, f: FnId, dist: &'p Expr, w: &World) -> Result<Option<Value>> {
810        let ExprKind::Builtin {
811            func: Builtin::Normal,
812            args,
813            named,
814        } = &dist.kind
815        else {
816            return Ok(None);
817        };
818        let ([mean, sd], true) = (args.as_slice(), named.is_empty()) else {
819            return Ok(None);
820        };
821        let Value::Analytic(m) = self.eval(f, mean, w)? else {
822            return Ok(None);
823        };
824        let Some(sd) = self.plain_sd(f, sd, w, dist.span)? else {
825            return Ok(None);
826        };
827        self.next_latent += 1;
828        match analytic::normal_draw(&m, sd, self.next_latent).map_err(|e| e.at(dist.span))? {
829            Some(v) => Ok(Some(v)),
830            None => Err(
831                analytic::unsupported("drawing from a normal distribution whose mean is this outcome").at(mean.span),
832            ),
833        }
834    }
835
836    /// `observe y from normal(mean, sd)` when enumerating, with `mean` an
837    /// outcome of normal draws: the observation's density, as a logarithm,
838    /// with the world updated exactly. `None` if the mean isn't an outcome.
839    #[inline(never)]
840    fn observe_gaussian(&mut self, f: FnId, v: &Value, d: &'p Expr, w: &mut World, span: Span) -> Result<Option<f64>> {
841        let ExprKind::Builtin {
842            func: Builtin::Normal,
843            args,
844            named,
845        } = &d.kind
846        else {
847            return Ok(None);
848        };
849        let ([mean, sd], true) = (args.as_slice(), named.is_empty()) else {
850            return Ok(None);
851        };
852        let Value::Analytic(m) = self.eval(f, mean, w)? else {
853            return Ok(None);
854        };
855        let Some(sd) = self.plain_sd(f, sd, w, d.span)? else {
856            return Ok(None);
857        };
858        let y = match v {
859            Value::Int(_) | Value::Float(_) | Value::Prob(_) => v.as_f64(),
860            _ => None,
861        };
862        let Some(y) = y else {
863            return Ok(None);
864        };
865        self.check_density(f, span)?;
866        let observed = analytic::observe_normal(&m, y, sd, &mut self.next_latent).map_err(|e| e.at(span))?;
867        let Some((ln, known)) = observed else {
868            if self.config.integrate.is_some() {
869                return Ok(None);
870            }
871            return Err(analytic::unsupported(
872                "observing a value from a normal distribution whose mean is this outcome",
873            )
874            .at(mean.span));
875        };
876        self.densities = true;
877        w.densities += 1;
878        for (id, latent) in known {
879            Arc::make_mut(&mut w.constraints).insert(id, latent);
880        }
881        self.reread(w, span)?;
882        Ok(Some(ln))
883    }
884
885    /// `observe v from D(…)` with `@integrate`, when parameters of `D` are
886    /// outcomes of one draw: what it says about the draw, as a factor of
887    /// its likelihood, and its probability or density, as a logarithm.
888    /// `None` if it isn't such an observation.
889    #[inline(never)]
890    fn observe_numerically(
891        &mut self,
892        f: FnId,
893        v: &Value,
894        d: &'p Expr,
895        w: &mut World,
896        span: Span,
897    ) -> Result<Option<f64>> {
898        use crate::likelihood::{Factor, Law, Likelihood, Param};
899        let Some(tolerance) = self.config.integrate else {
900            return Ok(None);
901        };
902        let ExprKind::Builtin { func, args, named } = &d.kind else {
903            return Ok(None);
904        };
905        // Each law's parameters, with where each is valid: from, to, and
906        // whether each end is included.
907        type Range = (Option<(f64, bool)>, Option<(f64, bool)>);
908        let positive: Range = (Some((0.0, false)), None);
909        let any: Range = (None, None);
910        let probability: Range = (Some((0.0, true)), Some((1.0, true)));
911        let (law, ranges): (Law, Vec<Range>) = match func {
912            Builtin::Normal => (Law::Normal, vec![any, positive]),
913            Builtin::Lognormal => (Law::Lognormal, vec![any, positive]),
914            Builtin::Uniform => (Law::Uniform, vec![any, any]),
915            Builtin::Beta => (Law::Beta, vec![positive, positive]),
916            Builtin::Gamma => (Law::Gamma, vec![positive, positive]),
917            Builtin::Exponential => (Law::Exponential, vec![positive]),
918            Builtin::Binomial => (Law::Binomial, vec![any, probability]),
919            Builtin::Bernoulli => (Law::Bernoulli, vec![probability]),
920            Builtin::Poisson => (Law::Poisson, vec![(Some((0.0, true)), None)]),
921            Builtin::Geometric => (Law::Geometric, vec![(Some((0.0, false)), Some((1.0, true)))]),
922            _ => return Ok(None),
923        };
924        if !named.is_empty() || args.len() != ranges.len() {
925            return Ok(None);
926        }
927        let mut values = Vec::with_capacity(args.len());
928        for a in args {
929            values.push(self.eval(f, a, w)?);
930        }
931        let mut draw: Option<Arc<Analytic>> = None;
932        for value in &values {
933            match value {
934                Value::Analytic(a) => match &draw {
935                    None => draw = Some(a.clone()),
936                    Some(x) if x.same_axis(a) => {}
937                    Some(_) => {
938                        return Err(analytic::unsupported(
939                            "an observation whose parameters are outcomes of different draws",
940                        )
941                        .at(d.span));
942                    }
943                },
944                v if v.as_f64().is_some() => {}
945                _ => return Ok(None),
946            }
947        }
948        let Some(x) = draw else {
949            return Ok(None);
950        };
951        if x.form.is_some() {
952            return Err(
953                analytic::unsupported("integrating over a sum of draws, which has several dimensions").at(d.span),
954            );
955        }
956        let value = match (law, v) {
957            (Law::Bernoulli, Value::Bool(b)) => f64::from(u8::from(*b)),
958            (_, Value::Int(_) | Value::Float(_) | Value::Prob(_)) => match v.as_f64() {
959                Some(y) => y,
960                None => return Ok(None),
961            },
962            _ => return Ok(None),
963        };
964        // A whole number of trials, as `binomial` takes it.
965        if law == Law::Binomial {
966            let n = &values[0];
967            if !matches!(n, Value::Int(_)) {
968                return Err(analytic::unsupported("a number of trials that's an outcome of a draw").at(d.span));
969            }
970            builtins::call_plain(Builtin::Binomial, &[n.clone(), Value::Prob(0.5)], &mut self.budget)
971                .map_err(|e| e.at(d.span))?;
972        }
973        // Parameters outside their range: everywhere, the fault a number
974        // makes; on part of the draw's range, not supported yet.
975        let mut params = Vec::with_capacity(values.len());
976        for (value, &(lo, hi)) in values.iter().zip(&ranges) {
977            match value {
978                Value::Analytic(a) => {
979                    let bad = analytic::outside(&x, &a.pieces, lo, hi).map_err(|e| e.at(d.span))?;
980                    let mass = bad.mass();
981                    if mass > 0.0 {
982                        let name = law.name();
983                        return Err(if mass >= x.domain.mass() {
984                            RuntimeError::new(d.span, format!("`{name}` isn't defined for these parameters"))
985                                .as_fault(Fault::DomainError)
986                        } else {
987                            analytic::unsupported(&format!(
988                                "`{name}` with a parameter outside its range on part of the draw's range"
989                            ))
990                            .at(d.span)
991                        });
992                    }
993                    params.push(Param::Outcome(a.pieces.clone()));
994                }
995                v => params.push(Param::Plain(v.as_f64().expect("a number"))),
996            }
997        }
998        if law.is_density() {
999            self.check_density(f, span)?;
1000        }
1001        let mut factors = x.weight.as_ref().map(|l| l.factors.clone()).unwrap_or_default();
1002        factors.push(Factor { law, params, value });
1003        let before = match &x.weight {
1004            None => libm::log(x.domain.mass()),
1005            Some(l) => l.shift + libm::log(l.mass(&x.domain)),
1006        };
1007        let likelihood = Likelihood::new(x.family, factors, tolerance, &mut self.budget).map_err(|e| e.at(span))?;
1008        self.integrated = Some(tolerance);
1009        if law.is_density() {
1010            self.densities = true;
1011            w.densities += 1;
1012        }
1013        let Some(likelihood) = likelihood else {
1014            return Ok(Some(f64::NEG_INFINITY));
1015        };
1016        let after = likelihood.shift + libm::log(likelihood.mass(&x.domain));
1017        let latent = analytic::Latent {
1018            family: x.family,
1019            domain: x.domain.clone(),
1020            form: None,
1021            weight: Some(Arc::new(likelihood)),
1022        };
1023        self.know(w, x.id, latent, span)?;
1024        Ok(Some(after - before))
1025    }
1026
1027    /// An observation's distribution that takes outcomes of draws as
1028    /// parameters, without `@integrate`: its error says that `@integrate`
1029    /// would compute it.
1030    #[cold]
1031    fn integration_hint(&self, d: &Expr, e: RuntimeError) -> RuntimeError {
1032        let law = matches!(
1033            &d.kind,
1034            ExprKind::Builtin {
1035                func: Builtin::Normal
1036                    | Builtin::Lognormal
1037                    | Builtin::Uniform
1038                    | Builtin::Beta
1039                    | Builtin::Gamma
1040                    | Builtin::Exponential
1041                    | Builtin::Binomial
1042                    | Builtin::Bernoulli
1043                    | Builtin::Poisson
1044                    | Builtin::Geometric,
1045                ..
1046            }
1047        );
1048        if e.kind != ErrorKind::Unsupported || !law || self.config.integrate.is_some() || self.sampler.is_some() {
1049            return e;
1050        }
1051        e.with_help(
1052            "`@integrate` computes such an observation of one draw by numerical integration, approximately; or use `@mode sample(runs: 10_000)`",
1053        )
1054    }
1055
1056    /// A normal distribution's standard deviation, checked as `normal`
1057    /// checks it; `None` if it isn't a plain value.
1058    fn plain_sd(&mut self, f: FnId, sd: &Expr, w: &World, span: Span) -> Result<Option<f64>> {
1059        let sd = self.eval(f, sd, w)?;
1060        if analytic::contains(&sd) || sd.is_uncertain() {
1061            return Ok(None);
1062        }
1063        let checked = builtins::call_plain(Builtin::Normal, &[Value::Float(0.0), sd], &mut self.budget)
1064            .map_err(|e| e.at(span))?;
1065        let Value::Continuous(family) = checked else {
1066            unreachable!("`normal` gives a continuous distribution")
1067        };
1068        let Family::Normal { sd, .. } = *family else {
1069            unreachable!("`normal` gives a normal distribution")
1070        };
1071        Ok(Some(sd))
1072    }
1073
1074    /// Drawing from `bernoulli(p)` or `binomial(n, p)` when enumerating, with
1075    /// `p` a variable drawn from a beta: each outcome, in a world where `p`
1076    /// is updated exactly by it. `None` for another draw.
1077    fn dependent_draw(&mut self, f: FnId, dist: &'p Expr, w: &World) -> Result<Option<Vec<(Value, World)>>> {
1078        let ExprKind::Builtin { func, args, named } = &dist.kind else {
1079            return Ok(None);
1080        };
1081        if self.sampler.is_some() || !named.is_empty() {
1082            return Ok(None);
1083        }
1084        let (p, trials) = match (func, args.as_slice()) {
1085            (Builtin::Bernoulli, [p]) => (p, None),
1086            (Builtin::Binomial, [n, p]) => (p, Some(n)),
1087            _ => return Ok(None),
1088        };
1089        // The variable, or `prob` of it.
1090        let slot = match &p.kind {
1091            ExprKind::Slot(s) => *s,
1092            ExprKind::Builtin {
1093                func: Builtin::Prob,
1094                args,
1095                ..
1096            } => match args.as_slice() {
1097                [
1098                    Expr {
1099                        kind: ExprKind::Slot(s),
1100                        ..
1101                    },
1102                ] => *s,
1103                _ => return Ok(None),
1104            },
1105            _ => return Ok(None),
1106        };
1107        let Some(x) = self.draw_in(w, slot, p.span)? else {
1108            return Ok(None);
1109        };
1110        let outcomes: Vec<(Value, Seen)> = match trials {
1111            None => [false, true].map(|b| (Value::Bool(b), Seen::Bernoulli(b))).into(),
1112            Some(n) => {
1113                let n = self.eval(f, n, w)?;
1114                if analytic::contains(&n) {
1115                    return Ok(None);
1116                }
1117                let counts = builtins::counts(Builtin::Binomial, &[n, Value::Prob(0.5)], &mut self.budget)
1118                    .map_err(|e| e.at(dist.span))?;
1119                let Some(Counts::Binomial { n, .. }) = counts else {
1120                    return Ok(None);
1121                };
1122                self.budget.outcomes(n as u128 + 1).map_err(|e| e.at(dist.span))?;
1123                (0..=n)
1124                    .map(|k| (Value::Int(k.into()), Seen::Binomial { trials: n, k: k as f64 }))
1125                    .collect()
1126            }
1127        };
1128        let seen: Vec<Seen> = outcomes.iter().map(|(_, s)| *s).collect();
1129        let worlds = self.outcomes_of(&x, w, &seen, dist.span)?;
1130        if worlds.is_empty() {
1131            return Ok(None);
1132        }
1133        Ok(Some(
1134            outcomes
1135                .into_iter()
1136                .zip(worlds)
1137                .filter_map(|((v, _), w)| Some((v, w?)))
1138                .collect(),
1139        ))
1140    }
1141
1142    /// The worlds where a draw `x`, the parameter of a likelihood, saw each
1143    /// of `seen`, updated exactly: an empty list if `x` isn't the prior of
1144    /// that likelihood, and `None` for what `x` can't have seen.
1145    fn outcomes_of(&mut self, x: &Analytic, w: &World, seen: &[Seen], span: Span) -> Result<Vec<Option<World>>> {
1146        let mut out = Vec::with_capacity(seen.len());
1147        for &s in seen {
1148            let Some((ln, latent)) = analytic::update(x, s).map_err(|e| e.at(span))? else {
1149                return Ok(Vec::new());
1150            };
1151            if ln == f64::NEG_INFINITY {
1152                out.push(None);
1153                continue;
1154            }
1155            let mut w = w.clone();
1156            w.weight = w.weight * Weight::from_ln(ln);
1157            self.know(&mut w, x.id, latent, span)?;
1158            out.push(Some(w));
1159        }
1160        Ok(out)
1161    }
1162
1163    /// The analytic outcome in `slot`, as the world reads it.
1164    fn draw_in(&mut self, w: &World, slot: SlotId, span: Span) -> Result<Option<Analytic>> {
1165        let v = &w.slots[slot as usize];
1166        if !matches!(v, Value::Analytic(_)) {
1167            return Ok(None);
1168        }
1169        Ok(
1170            match analytic::resolve(v, &w.constraints, &mut self.budget).map_err(|e| e.at(span))? {
1171                Value::Analytic(x) => Some((*x).clone()),
1172                _ => None,
1173            },
1174        )
1175    }
1176
1177    /// Record what a world now knows about a latent, and read the values
1178    /// that hold it again, as they would be read.
1179    fn know(&mut self, w: &mut World, id: u64, latent: analytic::Latent, span: Span) -> Result<()> {
1180        Arc::make_mut(&mut w.constraints).insert(id, latent);
1181        self.reread(w, span)
1182    }
1183
1184    /// Read the values that hold latents again, as the world reads them.
1185    fn reread(&mut self, w: &mut World, span: Span) -> Result<()> {
1186        for slot in &mut w.slots {
1187            if analytic::contains(slot) {
1188                *slot = analytic::resolve(slot, &w.constraints, &mut self.budget).map_err(|e| e.at(span))?;
1189            }
1190        }
1191        Ok(())
1192    }
1193
1194    /// What a conjugate observation saw, checked as when its variable is
1195    /// drawn; `None` if a part of it isn't a plain value.
1196    fn seen(&mut self, f: FnId, u: &Update<'p>, w: &World, from: Span) -> Result<Option<Seen>> {
1197        if self.sampler.is_none() {
1198            for e in u.others() {
1199                if analytic::contains(&self.eval(f, e, w)?) {
1200                    return Ok(None);
1201                }
1202            }
1203        }
1204        // A count, as observing one from a drawn distribution reads it:
1205        // anything but a whole number of 0 or more is impossible.
1206        let count = |v: &Value| match v {
1207            Value::Bool(_) => f64::NAN,
1208            _ => v.as_f64().unwrap_or(f64::NAN),
1209        };
1210        Ok(match u.likelihood {
1211            Likelihood::Binomial { value, trials } => {
1212                let v = self.eval(f, value, w)?;
1213                let n = self.eval(f, trials, w)?;
1214                match builtins::counts(Builtin::Binomial, &[n, Value::Prob(0.5)], &mut self.budget)
1215                    .map_err(|e| e.at(from))?
1216                {
1217                    Some(Counts::Binomial { n, .. }) => Some(Seen::Binomial {
1218                        trials: n,
1219                        k: count(&v),
1220                    }),
1221                    _ => None,
1222                }
1223            }
1224            Likelihood::Bernoulli { value } => match self.eval(f, value, w)? {
1225                Value::Bool(b) => Some(Seen::Bernoulli(b)),
1226                _ => None,
1227            },
1228            Likelihood::Poisson { value } => {
1229                let v = self.eval(f, value, w)?;
1230                Some(Seen::Poisson { k: count(&v) })
1231            }
1232            Likelihood::Normal { value, sd } => {
1233                let v = self.eval(f, value, w)?;
1234                let sd = self.eval(f, sd, w)?;
1235                if sd.is_uncertain() {
1236                    return Ok(None);
1237                }
1238                let checked = builtins::call_plain(Builtin::Normal, &[Value::Float(0.0), sd], &mut self.budget)
1239                    .map_err(|e| e.at(from))?;
1240                let Value::Continuous(family) = checked else {
1241                    unreachable!("`normal` gives a continuous distribution")
1242                };
1243                let Family::Normal { sd, .. } = *family else {
1244                    unreachable!("`normal` gives a normal distribution")
1245                };
1246                match v {
1247                    Value::Int(_) | Value::Float(_) | Value::Prob(_) if v.as_f64().is_some() => Some(Seen::Normal {
1248                        y: v.as_f64().expect("a number"),
1249                        sd,
1250                    }),
1251                    _ => None,
1252                }
1253            }
1254        })
1255    }
1256
1257    fn merge(&self, worlds: Vec<World>, stmt: StmtId) -> Vec<World> {
1258        merge(worlds, &self.live.after[stmt as usize], self.merging())
1259    }
1260
1261    /// Whether equal worlds merge: never when sampling, since every run
1262    /// counts separately (section 14), but inside `simulate` they do.
1263    fn merging(&self) -> bool {
1264        self.config.merging && self.sampler.is_none()
1265    }
1266
1267    /// Whether a fault ends only its world: in partial mode, except inside
1268    /// `simulate` and collection callbacks, which fail as a whole, in the
1269    /// world that runs them (docs/semantics.md, section 11).
1270    fn partial(&self) -> bool {
1271        self.config.partial && self.nested == 0 && self.callback.is_none()
1272    }
1273
1274    /// Whether a `try` running in this call catches `fault`.
1275    fn catches_here(&self, fault: Option<Fault>) -> bool {
1276        self.handlers
1277            .iter()
1278            .any(|(depth, catches)| *depth == self.depth && catches.iter().any(|c| c.catches(fault)))
1279    }
1280
1281    /// Whether a `try` running in a call further up catches `fault`.
1282    fn catches_above(&self, fault: Option<Fault>) -> bool {
1283        self.handlers
1284            .iter()
1285            .any(|(depth, catches)| *depth < self.depth && catches.iter().any(|c| c.catches(fault)))
1286    }
1287
1288    /// Inside a `try` of this call: a copy of `w`, for the catch it may
1289    /// fault to. Elsewhere nothing is copied.
1290    #[inline]
1291    fn snapshot(&self, w: &World) -> Option<World> {
1292        if self.handlers.is_empty() {
1293            return None;
1294        }
1295        self.handlers
1296            .iter()
1297            .any(|(depth, _)| *depth == self.depth)
1298            .then(|| w.clone())
1299    }
1300
1301    /// A world at `at` failed with `e`. A fault that a `try` of this call
1302    /// catches sends `saved`, the world as its statement began, on to the
1303    /// catch. One that a `try` further up catches, or any fault in partial
1304    /// mode, ends the world here: it's recorded, for the call's caller to
1305    /// catch, or as a failure. Either way the statement goes on with the
1306    /// other worlds. Anything else stops the run.
1307    fn fail(&mut self, e: RuntimeError, at: At, saved: Option<World>) -> Result<()> {
1308        if e.fault.is_none() {
1309            return Err(e);
1310        }
1311        if let Some(w) = saved.filter(|_| self.catches_here(e.fault)) {
1312            self.faulted.push((w, e));
1313            return Ok(());
1314        }
1315        if !self.partial() && !self.catches_above(e.fault) {
1316            return Err(e);
1317        }
1318        let fun = &self.prog.functions[at.f as usize];
1319        let e = match fun.kind {
1320            FnKind::Named => e.with_note(format!("in a call to `{}`", fun.name)),
1321            FnKind::Lambda => e.with_note("inside a lambda"),
1322            FnKind::Simulate => e.with_note("inside a `simulate` block"),
1323            FnKind::Main => e,
1324        };
1325        let stmt = at.stmt as usize;
1326        let before_evidence = self.evidence.at[stmt] || self.evidence.after[stmt];
1327        let run = self.sampler.is_some().then_some(at.run);
1328        self.failures.record(e, at.weight, run, before_evidence, at.densities);
1329        Ok(())
1330    }
1331
1332    /// Stop if a statement produced more worlds than allowed. (When sampling,
1333    /// worlds never multiply: there's one per run in the batch.)
1334    fn check_worlds(&self, n: usize, span: Span) -> Result<()> {
1335        if n > self.config.max_worlds && self.sampler.is_none() {
1336            return Err(RuntimeError::limit(
1337                span,
1338                format!("more than {} worlds are too many to follow", self.config.max_worlds),
1339            )
1340            .with_help("simplify the model, raise the limit, or sample it with `@mode sample(runs: 10_000)`"));
1341        }
1342        Ok(())
1343    }
1344
1345    fn spend(&mut self, n: u64, span: Span) -> Result<()> {
1346        self.budget.work(n).map_err(|e| e.at(span))?;
1347        if let Some(cancel) = &self.config.cancel {
1348            if cancel.load(Ordering::Relaxed) {
1349                return Err(RuntimeError::limit(span, "the run was cancelled"));
1350            }
1351        }
1352        Ok(())
1353    }
1354
1355    // ── Statements ───────────────────────────────────────────────────────
1356
1357    fn exec_block(&mut self, f: FnId, block: &'p Block, worlds: Vec<World>) -> Result<Flow> {
1358        let mut flow = Flow::next(worlds);
1359        for stmt in &block.stmts {
1360            if flow.next.is_empty() {
1361                break;
1362            }
1363            let input = std::mem::take(&mut flow.next);
1364            let out = self.exec_stmt(f, stmt, input)?;
1365            flow.next = out.next;
1366            clear(&mut flow.next, &self.live.dies[stmt.id as usize]);
1367            flow.broke.extend(out.broke);
1368            flow.continued.extend(out.continued);
1369            flow.returned.extend(out.returned);
1370            flow.faulted.extend(out.faulted);
1371        }
1372        Ok(flow)
1373    }
1374
1375    /// Weighs an enumerated world by the density `exp(ln)` of what an
1376    /// observation at `span` saw.
1377    fn observe_density(&mut self, f: FnId, ln: f64, w: &mut World, span: Span) -> Result<()> {
1378        self.check_density(f, span)?;
1379        if ln == f64::INFINITY {
1380            return Err(RuntimeError::new(span, "the density at the observed value is infinite"));
1381        }
1382        if ln == f64::NEG_INFINITY {
1383            self.lost += w.weight;
1384            w.weight = Weight::ZERO;
1385            self.last_ruling_out = Some(span);
1386        } else {
1387            w.weight = w.weight * Weight::from_ln(ln);
1388            w.densities += 1;
1389        }
1390        Ok(())
1391    }
1392
1393    /// An enumerated world's weight can include a density only where the
1394    /// world counts it (see `World::densities`).
1395    /// A `simulate` block's own observations can be densities too: its
1396    /// result is normalized, as long as its worlds include as many.
1397    #[inline(never)]
1398    fn check_density(&mut self, f: FnId, span: Span) -> Result<()> {
1399        if self.depth == 0 {
1400            self.density_at.get_or_insert(span);
1401            return Ok(());
1402        }
1403        if self.prog.functions[f as usize].kind == FnKind::Simulate {
1404            self.local_density_at.get_or_insert(span);
1405            return Ok(());
1406        }
1407        Err(OpError::unsupported(
1408            "observing a value from a continuous distribution inside a function isn't supported when enumerating yet",
1409        )
1410        .help("observe it in the main program or directly inside `simulate`, or sample the model with `@mode sample(runs: 10_000)`")
1411        .at(span))
1412    }
1413
1414    /// A world returning from a `simulate` block: its weight must include as
1415    /// many densities as the others'.
1416    #[inline(never)]
1417    fn simulate_returns(&mut self, w: &World, span: Span) -> Result<()> {
1418        match self.returned_densities {
1419            Some(n) if n != w.densities => Err(OpError::unsupported(
1420                "a `simulate` whose worlds observe different numbers of continuous values isn't supported when enumerating",
1421            )
1422            .help("a density is per unit of what it observes, so worlds with more of them can't be weighed against the others; observe the same values in every world")
1423            .at(self.local_density_at.unwrap_or(span))),
1424            _ => {
1425                self.returned_densities = Some(w.densities);
1426                Ok(())
1427            }
1428        }
1429    }
1430
1431    /// The parts of a world where an event holds, or doesn't: each with its
1432    /// probability given what the world knew, and the restrictions there.
1433    /// An event of several independent draws has a part for each box.
1434    fn event_parts(
1435        &mut self,
1436        event: &analytic::Event,
1437        yes: bool,
1438        constraints: &analytic::Constraints,
1439        span: Span,
1440    ) -> Result<Vec<(f64, analytic::Constraints)>> {
1441        if !event.restrictable() {
1442            return Err(analytic::joint_condition("a condition on several normal draws together").at(span));
1443        }
1444        let boxes = event.side(yes).map_err(|e| e.at(span))?;
1445        let n = boxes.len().max(1);
1446        self.budget
1447            .collection(((constraints.len() + event.draws.len()) * n) as u128)
1448            .map_err(|e| e.at(span))?;
1449        self.budget
1450            .work(((constraints.len() + event.size()) * n) as u64)
1451            .map_err(|e| e.at(span))?;
1452        let mut parts = Vec::with_capacity(boxes.len());
1453        for part in &boxes {
1454            let mut c = constraints.clone();
1455            let p = event.restrict(part, &mut c);
1456            parts.push((p, c));
1457        }
1458        Ok(parts)
1459    }
1460
1461    /// `if` on an event: each part of the world where it holds goes to
1462    /// `yes`, and each where it doesn't to `no`, restricted to that part.
1463    #[inline(never)]
1464    fn branch_on_event(
1465        &mut self,
1466        event: &analytic::Event,
1467        w: World,
1468        yes: &mut Vec<World>,
1469        no: &mut Vec<World>,
1470        span: Span,
1471    ) -> Result<()> {
1472        let holds = self.event_parts(event, true, &w.constraints, span)?;
1473        let fails = self.event_parts(event, false, &w.constraints, span)?;
1474        let parts: Vec<_> = (holds.into_iter().map(|(p, c)| (true, p, c)))
1475            .chain(fails.into_iter().map(|(p, c)| (false, p, c)))
1476            .filter(|(_, p, _)| *p > 0.0)
1477            .collect();
1478        let last = parts.len().saturating_sub(1);
1479        let mut w = Some(w);
1480        for (i, (side, p, c)) in parts.into_iter().enumerate() {
1481            // The last part takes the world itself.
1482            let mut nw = if i == last { w.take() } else { w.clone() }.expect("the world");
1483            nw.constraints = c;
1484            let out = if side { &mut *yes } else { &mut *no };
1485            out.push(nw.scaled(p));
1486        }
1487        Ok(())
1488    }
1489
1490    /// `observe` of an event, or of one side of it: each part of the world
1491    /// where it holds goes on, restricted to that part.
1492    #[inline(never)]
1493    fn observe_event(
1494        &mut self,
1495        event: &analytic::Event,
1496        yes: bool,
1497        w: World,
1498        out: &mut Vec<World>,
1499        span: Span,
1500    ) -> Result<()> {
1501        let parts = self.event_parts(event, yes, &w.constraints, span)?;
1502        let p: f64 = parts.iter().map(|(p, _)| p).sum();
1503        if p < 1.0 && self.sampler.is_none() {
1504            self.lost += w.weight.scale(1.0 - p);
1505        }
1506        if p <= 0.0 {
1507            self.last_ruling_out = Some(span);
1508        }
1509        let parts: Vec<_> = parts.into_iter().filter(|(p, _)| *p > 0.0).collect();
1510        let last = parts.len().saturating_sub(1);
1511        let mut w = Some(w);
1512        for (i, (p, c)) in parts.into_iter().enumerate() {
1513            let mut nw = if i == last { w.take() } else { w.clone() }.expect("the world");
1514            nw.constraints = c;
1515            out.push(nw.scaled(p));
1516        }
1517        Ok(())
1518    }
1519
1520    fn exec_stmt(&mut self, f: FnId, stmt: &'p Stmt, worlds: Vec<World>) -> Result<Flow> {
1521        // Outside every `try`, no world can fault to a catch.
1522        if self.handlers.is_empty() {
1523            return self.exec_stmt_kind(f, stmt, worlds);
1524        }
1525        // The worlds that fault in this statement, on their way to a catch,
1526        // leave with its flow.
1527        let outer = std::mem::take(&mut self.faulted);
1528        let flow = self.exec_stmt_kind(f, stmt, worlds);
1529        let faulted = std::mem::replace(&mut self.faulted, outer);
1530        let mut flow = flow?;
1531        flow.faulted.extend(faulted);
1532        Ok(flow)
1533    }
1534
1535    fn exec_stmt_kind(&mut self, f: FnId, stmt: &'p Stmt, worlds: Vec<World>) -> Result<Flow> {
1536        let span = stmt.span;
1537        if matches!(
1538            stmt.kind,
1539            StmtKind::Draw { .. } | StmtKind::Take { .. } | StmtKind::Observe { .. } | StmtKind::Chance { .. }
1540        ) {
1541            self.check_callback_effect(span)?;
1542        }
1543        let n = worlds.len();
1544        self.stats.world_steps += n as u64;
1545        self.stats.peak_worlds = self.stats.peak_worlds.max(n);
1546        self.spend(n as u64, span)?;
1547        self.check_worlds(n, span)?;
1548        let mut worlds = worlds;
1549        if self.delaying() {
1550            // Delayed variables this statement reads are drawn first.
1551            let conj = self.conj;
1552            let first = &conj.draws_first[stmt.id as usize];
1553            if !first.is_empty() {
1554                for w in &mut worlds {
1555                    for &slot in first {
1556                        self.draw_delayed(w, slot);
1557                    }
1558                }
1559            }
1560        }
1561        match &stmt.kind {
1562            StmtKind::Set { place, value } | StmtKind::Update { place, value, .. } => {
1563                let update = match &stmt.kind {
1564                    StmtKind::Update { check, .. } => Some(check.as_ref()),
1565                    _ => None,
1566                };
1567                let mut worlds = worlds;
1568                let mut failed = Vec::new();
1569                let mut discrete = false;
1570                for (i, w) in worlds.iter_mut().enumerate() {
1571                    let mut saved = self.snapshot(w);
1572                    if let Some(check) = update {
1573                        let updated = self.update(f, place, value, check, w, span);
1574                        each!(self, At::new(f, stmt, w), saved saved, updated, { failed.push(i) });
1575                        continue;
1576                    }
1577                    let v = match self.record_place_type(place, w) {
1578                        Some(ty) => self.eval_expected(f, value, w, &ty),
1579                        None => self.eval(f, value, w),
1580                    };
1581                    let v = each!(self, At::new(f, stmt, w), saved saved, v, { failed.push(i) });
1582                    if let Value::Analytic(_) = v {
1583                        discrete |= is_discrete(&v);
1584                    }
1585                    each!(self, At::new(f, stmt, w), saved saved, self.assign(f, place, v, w, span), {
1586                        failed.push(i)
1587                    });
1588                }
1589                if discrete {
1590                    self.split_discrete(f, stmt, place, &mut worlds, &mut failed)?;
1591                }
1592                drop_failed(&mut worlds, &failed);
1593                // When the live values say less than before, as when a
1594                // call's result is used up (`f = f - $5`), worlds may now be
1595                // the same: merging spares the statements after from
1596                // following each of them (a call next would multiply them).
1597                if self.live.narrows[stmt.id as usize] && self.merging() {
1598                    return Ok(Flow::next(self.merge(worlds, stmt.id)));
1599                }
1600                Ok(Flow::next(worlds))
1601            }
1602            StmtKind::Draw { place, dist } => {
1603                // A conjugate prior's draw may be delayed (section 14).
1604                let delay = if self.delaying() {
1605                    self.conj.delays[stmt.id as usize]
1606                } else {
1607                    None
1608                };
1609                let mut out = Vec::with_capacity(worlds.len());
1610                for w in worlds {
1611                    let at = At::new(f, stmt, &w);
1612                    let mut saved = self.snapshot(&w);
1613                    if self.sampler.is_some() {
1614                        if let Some(counts) = each!(self, at, saved saved, self.direct_counts(f, dist, &w)) {
1615                            let k = counts.sample(self.sampler.as_mut().expect("sampling"));
1616                            let mut w = w;
1617                            each!(self, at, saved saved, self.assign(f, place, Value::Int(k.into()), &mut w, dist.span));
1618                            out.push(w);
1619                            continue;
1620                        }
1621                    }
1622                    if self.sampler.is_none()
1623                        && each!(self, at, saved saved, self.draw_dependent(f, place, dist, &w, &mut out))
1624                    {
1625                        self.check_worlds(out.len(), span)?;
1626                        continue;
1627                    }
1628                    let d = each!(self, at, saved saved, self.eval(f, dist, &w));
1629                    if let (Some(variable), Value::Continuous(family)) = (delay, &d) {
1630                        if conjugate::is_prior(family) {
1631                            let mut w = w;
1632                            w.slots[place.slot as usize] = Value::Delayed(Arc::new(Delayed {
1633                                family: **family,
1634                                variable,
1635                            }));
1636                            self.updates(variable).delayed += 1;
1637                            out.push(w);
1638                            continue;
1639                        }
1640                    }
1641                    let mark = out.len();
1642                    each!(self, at, saved saved, self.split_by(f, place, d, w, dist.span, &mut out), {
1643                        out.truncate(mark)
1644                    });
1645                    self.check_worlds(out.len(), span)?;
1646                }
1647                Ok(Flow::next(self.merge(out, stmt.id)))
1648            }
1649            StmtKind::Take { place, bag } => {
1650                let mut out = Vec::new();
1651                for w in worlds {
1652                    let at = At::new(f, stmt, &w);
1653                    let current = each!(self, at, world w, self.read_place(f, bag, &w, span));
1654                    let Value::Bag(cards) = &current else {
1655                        return Err(RuntimeError::new(
1656                            span,
1657                            format!("`take` needs a bag, found {}", ops::article(&current.kind())),
1658                        )
1659                        .with_help("make one with `bag([card: count, …])`"));
1660                    };
1661                    let total: u128 = cards.values().map(|n| *n as u128).sum();
1662                    if total == 0 {
1663                        let empty = RuntimeError::new(span, "can't take a card from an empty bag")
1664                            .as_fault(Fault::EmptyCollection);
1665                        each!(self, at, world w, Err::<(), _>(empty));
1666                    }
1667                    let chosen = match &mut self.sampler {
1668                        Some(rng) => rng.choose(cards.values().map(|n| *n as f64)),
1669                        None => None,
1670                    };
1671                    for (i, (card, count)) in cards.iter().enumerate() {
1672                        if self.sampler.is_some() && chosen != Some(i) {
1673                            continue;
1674                        }
1675                        let rest = cards.without_nth(i);
1676                        let mut nw = if self.sampler.is_some() {
1677                            w.clone()
1678                        } else {
1679                            w.clone().scaled(*count as f64 / total as f64)
1680                        };
1681                        self.assign(f, bag, Value::multiset(rest), &mut nw, span)?;
1682                        self.assign(f, place, card.clone(), &mut nw, span)?;
1683                        out.push(nw);
1684                    }
1685                    self.check_worlds(out.len(), span)?;
1686                }
1687                Ok(Flow::next(self.merge(out, stmt.id)))
1688            }
1689            StmtKind::Call { dest, callee, args } => {
1690                let mut out = Vec::with_capacity(worlds.len());
1691                for w in worlds {
1692                    let at = At::new(f, stmt, &w);
1693                    let mut key = Vec::with_capacity(args.len() + 4);
1694                    let evaluated: Result<()> = args.iter().try_for_each(|a| {
1695                        key.push(self.eval(f, a, &w)?);
1696                        Ok(())
1697                    });
1698                    each!(self, at, world w, evaluated);
1699                    let func = match callee {
1700                        Callee::Fn { func, capture_args } => {
1701                            self.capture(&mut key, f, *func, capture_args, &w, span, false)?;
1702                            *func
1703                        }
1704                        Callee::Value(e, named) => match each!(self, at, world w, self.eval(f, e, &w)) {
1705                            Value::Builtin(b) => {
1706                                let value =
1707                                    each!(self, at, world w, self.builtin_values(b, &key, named, w.weight, span));
1708                                let mut saved = self.snapshot(&w);
1709                                let mut nw = w;
1710                                each!(self, at, saved saved, self.assign(f, dest, value, &mut nw, span));
1711                                out.push(nw);
1712                                continue;
1713                            }
1714                            Value::Closure(c) => {
1715                                if !named.is_empty() {
1716                                    return Err(RuntimeError::new(
1717                                        span,
1718                                        "only builtin minimum/maximum accept named defaults",
1719                                    ));
1720                                }
1721                                self.check_arity(&c, args.len(), span)?;
1722                                key.extend(c.captured.iter().cloned());
1723                                c.func
1724                            }
1725                            other => {
1726                                return Err(RuntimeError::new(
1727                                    e.span,
1728                                    format!("can't call {}", ops::article(&other.kind())),
1729                                ));
1730                            }
1731                        },
1732                    };
1733                    let result = each!(self, at, world w, self.call(func, key, span));
1734                    self.unresolved += w.weight * result.unresolved;
1735                    self.lost += w.weight * result.lost;
1736                    // What faulted in the call faulted in this world: a catch
1737                    // here takes it with the world as it was, and otherwise
1738                    // the world fails, for a `try` further up or as a
1739                    // failure, and may still meet evidence after the call.
1740                    let run = self.sampler.is_some().then_some(w.run);
1741                    let after = self.evidence.after[stmt.id as usize];
1742                    for g in &result.failures.groups {
1743                        if self.catches_here(g.error.fault) {
1744                            let mut caught = w.clone();
1745                            caught.weight = caught.weight * g.weight;
1746                            self.faulted.push((caught, g.error.clone()));
1747                        } else if self.partial() || self.catches_above(g.error.fault) {
1748                            self.failures.absorb_group(g, w.weight, run, after, w.densities);
1749                        } else {
1750                            return Err(g.error.clone());
1751                        }
1752                    }
1753                    add_pending(&mut self.pending, &result.pending, w.weight);
1754                    let last = result.outcomes.len().saturating_sub(1);
1755                    let mut w = Some(w);
1756                    for (i, (v, p, restrictions)) in result.outcomes.iter().enumerate() {
1757                        let mut nw = if i == last {
1758                            w.take().unwrap()
1759                        } else {
1760                            w.clone().unwrap()
1761                        };
1762                        nw.weight = nw.weight * *p;
1763                        // Only copy the constraints (shared with the other
1764                        // outcomes' worlds) when the call added some.
1765                        if !restrictions.is_empty() {
1766                            Arc::make_mut(&mut nw.constraints)
1767                                .extend(restrictions.iter().map(|(id, d)| (*id, d.clone())));
1768                        }
1769                        self.assign(f, dest, v.clone(), &mut nw, span)?;
1770                        out.push(nw);
1771                    }
1772                    self.check_worlds(out.len(), span)?;
1773                }
1774                Ok(Flow::next(self.merge(out, stmt.id)))
1775            }
1776            StmtKind::If { cond, then, otherwise } => {
1777                let (mut yes, mut no) = (Vec::new(), Vec::new());
1778                for w in worlds {
1779                    let condition = each!(self, At::new(f, stmt, &w), world w, self.eval(f, cond, &w));
1780                    if let Value::Event(event) = condition {
1781                        self.check_callback_effect(cond.span)?;
1782                        self.branch_on_event(&event, w, &mut yes, &mut no, cond.span)?;
1783                        continue;
1784                    }
1785                    if let Value::Analytic(x) = &condition {
1786                        if self.branch_on_draw(x, &w, &mut yes, &mut no, cond.span)? {
1787                            continue;
1788                        }
1789                    }
1790                    let c = ops::condition(&condition).map_err(|e| e.at(cond.span))?;
1791                    if (c.yes > 0.0 && c.no > 0.0) || c.missing > 0.0 {
1792                        self.check_callback_effect(cond.span)?;
1793                    }
1794                    if let Some(rng) = &mut self.sampler {
1795                        // One branch, chosen with its probability.
1796                        match rng.choose([c.yes, c.no].into_iter()) {
1797                            Some(0) => yes.push(w),
1798                            Some(_) => no.push(w),
1799                            None => self.unresolved += w.weight,
1800                        }
1801                        continue;
1802                    }
1803                    let (p, q) = (c.yes, c.no);
1804                    if c.missing > 0.0 {
1805                        self.unresolved += w.weight.scale(c.missing);
1806                    }
1807                    if q <= 0.0 {
1808                        yes.push(w.scaled(p));
1809                    } else if p <= 0.0 {
1810                        no.push(w.scaled(q));
1811                    } else {
1812                        no.push(w.clone().scaled(q));
1813                        yes.push(w.scaled(p));
1814                    }
1815                }
1816                self.check_worlds(yes.len() + no.len(), span)?;
1817                let mut flow = Flow::default();
1818                if !yes.is_empty() {
1819                    flow.join(self.exec_block(f, then, yes)?);
1820                }
1821                if !no.is_empty() {
1822                    flow.join(self.exec_block(f, otherwise, no)?);
1823                }
1824                flow.next = self.merge(flow.next, stmt.id);
1825                Ok(flow)
1826            }
1827            StmtKind::Chance {
1828                arms,
1829                otherwise,
1830                exhaustive,
1831            } => {
1832                let mut buckets: Vec<Vec<World>> = vec![Vec::new(); arms.len()];
1833                let mut rest = Vec::new();
1834                for w in worlds {
1835                    let at = At::new(f, stmt, &w);
1836                    // Chance weights are explicit probabilities.
1837                    let chances: Result<Vec<f64>> = arms
1838                        .iter()
1839                        .map(|(weight, _)| {
1840                            let value = self.eval(f, weight, &w)?;
1841                            ops::to_prob(&value).map_err(|e| e.at(weight.span))
1842                        })
1843                        .collect();
1844                    let chances = each!(self, at, world w, chances);
1845                    let sum: f64 = chances.iter().sum();
1846                    if sum > 1.0 + 1e-9 {
1847                        let over = RuntimeError::new(
1848                            span,
1849                            format!(
1850                                "the chances in this `chance` add up to {}, more than 100%",
1851                                fmt_prob(sum)
1852                            ),
1853                        );
1854                        each!(self, at, world w, Err::<(), _>(over.as_fault(Fault::DomainError)));
1855                    }
1856                    let remainder = (1.0 - sum).max(0.0);
1857                    if remainder > 1e-9 && *exhaustive && otherwise.is_none() {
1858                        let short = RuntimeError::new(
1859                            span,
1860                            format!("the chances add up to {} and there's no `else` arm", fmt_prob(sum)),
1861                        )
1862                        .with_help("a `chance` used as a value needs its chances to add up to 100%, or an `else`");
1863                        each!(self, at, world w, Err::<(), _>(short.as_fault(Fault::DomainError)));
1864                    }
1865                    if let Some(rng) = &mut self.sampler {
1866                        // One arm, chosen with its probability.
1867                        let options = chances.iter().copied().chain(std::iter::once(remainder));
1868                        match rng.choose(options) {
1869                            Some(i) if i < arms.len() => buckets[i].push(w),
1870                            Some(_) => rest.push(w),
1871                            None => self.unresolved += w.weight,
1872                        }
1873                        continue;
1874                    }
1875                    for (bucket, p) in buckets.iter_mut().zip(&chances) {
1876                        if *p > 0.0 {
1877                            bucket.push(w.clone().scaled(*p));
1878                        }
1879                    }
1880                    if remainder > 1e-12 {
1881                        rest.push(w.scaled(remainder));
1882                    }
1883                }
1884                self.check_worlds(buckets.iter().map(Vec::len).sum::<usize>() + rest.len(), span)?;
1885                let mut flow = Flow::default();
1886                for ((_, body), bucket) in arms.iter().zip(buckets) {
1887                    if !bucket.is_empty() {
1888                        flow.join(self.exec_block(f, body, bucket)?);
1889                    }
1890                }
1891                if !rest.is_empty() {
1892                    match otherwise {
1893                        Some(body) => flow.join(self.exec_block(f, body, rest)?),
1894                        None if *exhaustive => {}
1895                        None => flow.next.extend(rest),
1896                    }
1897                }
1898                flow.next = self.merge(flow.next, stmt.id);
1899                Ok(flow)
1900            }
1901            StmtKind::Loop { body, bounded } => self.exec_loop(f, stmt, body, *bounded, worlds),
1902            StmtKind::Break => Ok(Flow {
1903                broke: worlds,
1904                ..Flow::default()
1905            }),
1906            StmtKind::Continue => Ok(Flow {
1907                continued: worlds,
1908                ..Flow::default()
1909            }),
1910            StmtKind::Return(value) => self.exec_return(f, stmt, value, worlds),
1911            StmtKind::Observe { value, from } => {
1912                self.observed = true;
1913                // When sampling, an update of a delayed draw; when
1914                // enumerating, of an analytic one.
1915                let exact = if self.delaying() || self.sampler.is_none() {
1916                    probl_sema::conjugate::update(value, from.as_ref())
1917                } else {
1918                    None
1919                };
1920                let mut out = Vec::with_capacity(worlds.len());
1921                for mut w in worlds {
1922                    let here = At::new(f, stmt, &w);
1923                    let mut saved = self.snapshot(&w);
1924                    if let Some(u) = &exact {
1925                        let at = from.as_ref().map_or(span, |d| d.span);
1926                        let ln = if self.sampler.is_some() {
1927                            self.observe_exactly(f, u, &mut w, at)
1928                        } else {
1929                            self.observe_posterior(f, u, &mut w, at)
1930                        };
1931                        if let Some(ln) = each!(self, here, saved saved, ln) {
1932                            w.weight = w.weight * Weight::from_ln(ln);
1933                            if w.weight.is_zero() {
1934                                self.last_ruling_out = Some(span);
1935                            } else {
1936                                out.push(w);
1937                            }
1938                            continue;
1939                        }
1940                    }
1941                    let (factor, missing, ruled_out) = match from {
1942                        None => {
1943                            let v = each!(self, here, saved saved, self.eval(f, value, &w));
1944                            if let Value::Event(event) = v {
1945                                self.observe_event(&event, true, w, &mut out, value.span)?;
1946                                continue;
1947                            } else {
1948                                let b = ops::fact(&v, "observe").map_err(|e| e.at(value.span))?;
1949                                if b { (1.0, 0.0, 0.0) } else { (0.0, 0.0, 1.0) }
1950                            }
1951                        }
1952                        Some(d) => {
1953                            let v = each!(self, here, saved saved, self.eval(f, value, &w));
1954                            match each!(self, here, saved saved, self.direct_counts(f, d, &w)) {
1955                                Some(counts) if self.sampler.is_some() => {
1956                                    let x = match v {
1957                                        Value::Bool(_) => None,
1958                                        _ => v.as_f64(),
1959                                    };
1960                                    (x.map_or(0.0, |x| counts.pmf(x)), 0.0, 0.0)
1961                                }
1962                                _ => {
1963                                    if self.sampler.is_none() {
1964                                        let mut exact = self.observe_gaussian(f, &v, d, &mut w, span);
1965                                        if matches!(exact, Ok(None)) {
1966                                            exact = self.observe_numerically(f, &v, d, &mut w, span);
1967                                        }
1968                                        if let Some(ln) = each!(self, here, saved saved, exact) {
1969                                            w.weight = w.weight * Weight::from_ln(ln);
1970                                            if w.weight.is_zero() {
1971                                                self.last_ruling_out = Some(span);
1972                                            } else {
1973                                                out.push(w);
1974                                            }
1975                                            continue;
1976                                        }
1977                                    }
1978                                    let dist = self.eval(f, d, &w).map_err(|e| self.integration_hint(d, e));
1979                                    let dist = each!(self, here, saved saved, dist);
1980                                    if let (Value::Bool(b), Value::Event(event)) = (&v, &dist) {
1981                                        self.observe_event(event, *b, w, &mut out, span)?;
1982                                        continue;
1983                                    } else {
1984                                        if analytic::contains(&v) || analytic::contains(&dist) {
1985                                            return Err(analytic::unsupported("this likelihood observation").at(span));
1986                                        }
1987                                        self.densities |= is_density(&dist);
1988                                        match each!(
1989                                            self,
1990                                            here,
1991                                            saved saved,
1992                                            likelihood(&dist, &v, self.sampler.is_some()).map_err(|e| e.at(span))
1993                                        ) {
1994                                            Weighs::Probability(p, missing, other) => (p, missing, other),
1995                                            Weighs::LogDensity(ln) => {
1996                                                self.observe_density(f, ln, &mut w, span)?;
1997                                                if !w.weight.is_zero() {
1998                                                    out.push(w);
1999                                                }
2000                                                continue;
2001                                            }
2002                                        }
2003                                    }
2004                                }
2005                            }
2006                        }
2007                    };
2008                    if missing > 0.0 && self.sampler.is_none() {
2009                        self.unresolved += w.weight.scale(missing);
2010                    }
2011                    if ruled_out > 0.0 && self.sampler.is_none() {
2012                        self.lost += w.weight.scale(ruled_out);
2013                    }
2014                    if factor > 0.0 {
2015                        w.weight = w.weight.scale(factor);
2016                        out.push(w);
2017                    } else {
2018                        self.last_ruling_out = Some(span);
2019                    }
2020                }
2021                Ok(Flow::next(out))
2022            }
2023            StmtKind::Report { site, value, key } => {
2024                let mut failed = Vec::new();
2025                for (i, w) in worlds.iter().enumerate() {
2026                    let v = each!(self, At::new(f, stmt, w), world w, self.eval(f, value, w), { failed.push(i) });
2027                    // A report of a count law lists it.
2028                    let v = match v {
2029                        Value::Counts(_) => ops::listed(&v, &mut self.budget)
2030                            .map_err(|e| e.at(value.span))?
2031                            .into_owned(),
2032                        v => v,
2033                    };
2034                    let k = match key {
2035                        Some(k) => {
2036                            let k_value =
2037                                each!(self, At::new(f, stmt, w), world w, self.eval(f, k, w), { failed.push(i) });
2038                            if analytic::contains(&k_value) && !matches!(k_value, Value::Event(_)) {
2039                                return Err(analytic::unsupported("grouping by a continuous outcome").at(k.span));
2040                            }
2041                            if k_value.is_uncertain() {
2042                                return Err(RuntimeError::new(
2043                                    k.span,
2044                                    "a `by` key must be a plain value, not a distribution",
2045                                )
2046                                .with_help("draw a value first with `~`"));
2047                            }
2048                            k_value
2049                        }
2050                        None => Value::Unit,
2051                    };
2052                    // A distribution of outcomes, from `max(x, d6)` say, reports
2053                    // as their mixture, and a record or list of them field by
2054                    // field.
2055                    let (v, run) = if self.sampler.is_some() {
2056                        (self.sample_continuous(v), Some(w.run))
2057                    } else {
2058                        (v, None)
2059                    };
2060                    if let Value::Event(event) = &k {
2061                        // An event of a draw is a key with two values: each
2062                        // group has its part of the world.
2063                        let span = key.as_ref().map_or(span, |k| k.span);
2064                        for yes in [true, false] {
2065                            for (p, constraints) in self.event_parts(event, yes, &w.constraints, span)? {
2066                                if p <= 0.0 {
2067                                    continue;
2068                                }
2069                                let v = analytic::resolve(&v, &constraints, &mut self.budget)
2070                                    .map_err(|e| e.at(value.span))?;
2071                                let k = Value::Bool(yes);
2072                                self.same_units(w, value.span)?;
2073                                let sink = &self.sinks[*site as usize];
2074                                sink.validate_analytic(&k, &v).map_err(|e| e.at(value.span))?;
2075                                sink.validate_structure(&k, &v).map_err(|e| e.at(value.span))?;
2076                                self.sinks[*site as usize].add(k, &v, w.weight.scale(p), run);
2077                            }
2078                        }
2079                        continue;
2080                    }
2081                    if run.is_none() {
2082                        self.same_units(w, value.span)?;
2083                        self.sinks[*site as usize]
2084                            .validate_structure(&k, &v)
2085                            .map_err(|e| e.at(value.span))?;
2086                    }
2087                    self.sinks[*site as usize]
2088                        .validate_analytic(&k, &v)
2089                        .map_err(|e| e.at(value.span))?;
2090                    self.sinks[*site as usize].add(k, &v, w.weight, run);
2091                }
2092                let mut worlds = worlds;
2093                drop_failed(&mut worlds, &failed);
2094                Ok(Flow::next(worlds))
2095            }
2096            StmtKind::Fail { message } => Err(RuntimeError::new(span, message.clone())),
2097            StmtKind::Try { body, catches } => {
2098                self.handlers.push((self.depth, catches.as_slice()));
2099                let flow = self.exec_block(f, body, worlds);
2100                self.handlers.pop();
2101                let mut flow = flow?;
2102                // Each faulted world goes on in the first catch that takes
2103                // its fault; the others are for a `try` around this one.
2104                let mut caught: Vec<Vec<World>> = vec![Vec::new(); catches.len()];
2105                for (w, e) in std::mem::take(&mut flow.faulted) {
2106                    match catches.iter().position(|c| c.catches(e.fault)) {
2107                        Some(i) => caught[i].push(w),
2108                        None => flow.faulted.push((w, e)),
2109                    }
2110                }
2111                for (c, worlds) in catches.iter().zip(caught) {
2112                    if !worlds.is_empty() {
2113                        flow.join(self.exec_block(f, &c.body, worlds)?);
2114                    }
2115                }
2116                flow.next = self.merge(flow.next, stmt.id);
2117                Ok(flow)
2118            }
2119            StmtKind::Check { slot, ty } => {
2120                for w in &mut worlds {
2121                    // A delayed variable is drawn only if some of its values
2122                    // could fail the check.
2123                    if let Value::Delayed(_) = &w.slots[*slot as usize] {
2124                        let always = *ty == TypeSpec::Float;
2125                        if always {
2126                            continue;
2127                        }
2128                        self.draw_delayed(w, *slot);
2129                    }
2130                    self.check_slot(f, *slot, ty, w, span)?;
2131                }
2132                Ok(Flow::next(worlds))
2133            }
2134        }
2135    }
2136
2137    /// `return`: each world's value, with the restrictions its caller can
2138    /// see. (Out of `exec_stmt_kind`, whose frame every statement pays for.)
2139    #[inline(never)]
2140    fn exec_return(&mut self, f: FnId, stmt: &'p Stmt, value: &'p Expr, worlds: Vec<World>) -> Result<Flow> {
2141        let mut flow = Flow::default();
2142        let simulate = self.prog.functions[f as usize].kind == FnKind::Simulate;
2143        for w in worlds {
2144            let v = each!(self, At::new(f, stmt, &w), world w, self.eval(f, value, &w));
2145            if simulate {
2146                self.simulate_returns(&w, stmt.span)?;
2147            }
2148            let mut constraints = w.constraints;
2149            if constraints.keys().any(|id| !w.inherited.contains(id)) {
2150                Arc::make_mut(&mut constraints).retain(|id, _| w.inherited.contains(id));
2151            }
2152            flow.returned.push((v, w.weight, constraints));
2153        }
2154        Ok(flow)
2155    }
2156
2157    fn exec_loop(
2158        &mut self,
2159        f: FnId,
2160        stmt: &'p Stmt,
2161        body: &'p Block,
2162        bounded: bool,
2163        worlds: Vec<World>,
2164    ) -> Result<Flow> {
2165        let live = self.live;
2166        let span = stmt.span;
2167        let entered = total_weight(&worlds);
2168        let cutoff = entered.scale(self.config.epsilon);
2169        let mut inside = worlds;
2170        let mut out = Flow::default();
2171        let mut iterations: u64 = 0;
2172        // The states met at the loop's start in earlier iterations, while it
2173        // may be solved. Once one comes back, the loop cycles, and it's
2174        // solved as a Markov chain from there (section 10).
2175        let mut met = self.may_solve(stmt, bounded).then(FxHashSet::<u64>::default);
2176        while !inside.is_empty() {
2177            let cycles = match &mut met {
2178                Some(met) => {
2179                    // A loop's rewards don't make its states new.
2180                    let rewards = self.rewards.get(&stmt.id).map_or(&[][..], Vec::as_slice);
2181                    let mut head = live_slots(&live.loop_head[stmt.id as usize], inside[0].slots.len());
2182                    head.retain(|&i| !rewards.contains(&(i as SlotId)));
2183                    let hashes: Vec<u64> = inside.iter().map(|w| state_hash(w, &head)).collect();
2184                    let again = hashes.iter().any(|h| met.contains(h));
2185                    met.extend(hashes);
2186                    again
2187                }
2188                None => false,
2189            };
2190            if cycles {
2191                met = None;
2192                if let Some(solved) = self.solve_loop(f, stmt, body, &inside, cutoff, &mut iterations)? {
2193                    out.join(solved);
2194                    break;
2195                }
2196            }
2197            let mass = total_weight(&inside);
2198            if !bounded && mass < cutoff && self.sampler.is_none() {
2199                self.unresolved += mass;
2200                break;
2201            }
2202            iterations += 1;
2203            if iterations > self.config.max_iterations {
2204                return Err(RuntimeError::limit(
2205                    span,
2206                    format!("this loop ran {} times without finishing", self.config.max_iterations),
2207                )
2208                .with_note(format!(
2209                    "worlds still inside the loop weigh {} of what entered it",
2210                    fmt_prob(mass.ratio(entered))
2211                ))
2212                .with_help("check that every world can leave the loop"));
2213            }
2214            let mut flow = self.exec_block(f, body, inside)?;
2215            clear_dead(&mut flow.broke, &live.after[stmt.id as usize]);
2216            clear_dead(&mut flow.continued, &live.loop_head[stmt.id as usize]);
2217            out.next.extend(flow.broke);
2218            out.returned.extend(flow.returned);
2219            out.faulted.extend(flow.faulted);
2220            let mut again = flow.next;
2221            again.extend(flow.continued);
2222            inside = merge(again, &live.loop_head[stmt.id as usize], self.merging());
2223        }
2224        out.next = self.merge(out.next, stmt.id);
2225        Ok(out)
2226    }
2227
2228    /// Whether a loop may be solved as a Markov chain: an unbounded loop
2229    /// whose body doesn't report or print, when enumerating, unless its
2230    /// chain was already found too large.
2231    fn may_solve(&self, stmt: &Stmt, bounded: bool) -> bool {
2232        !bounded
2233            && self.config.solving
2234            && self.sampler.is_none()
2235            && self.solvable[stmt.id as usize]
2236            && !self.too_large.contains(&stmt.id)
2237    }
2238
2239    /// Solve a loop as an absorbing Markov chain, from the worlds `inside`
2240    /// at its start (section 10). Its states are the worlds at its start,
2241    /// told apart by the variables read later. Running the body once from
2242    /// each state reachable from `inside` gives the chance of going on to
2243    /// each state, and the worlds that leave. The expected number of visits
2244    /// to each state then says how much of each leaves. `None` if the chain
2245    /// is too large, and the loop should be unrolled instead.
2246    fn solve_loop(
2247        &mut self,
2248        f: FnId,
2249        stmt: &'p Stmt,
2250        body: &'p Block,
2251        inside: &[World],
2252        cutoff: Weight,
2253        iterations: &mut u64,
2254    ) -> Result<Option<Flow>> {
2255        // The chain's transitions are probabilities: a density isn't one.
2256        if inside
2257            .iter()
2258            .any(|w| w.densities > 0 || !w.constraints.is_empty() || w.slots.iter().any(analytic::contains))
2259        {
2260            self.too_large.insert(stmt.id);
2261            return Ok(None);
2262        }
2263        let head_set = &self.live.loop_head[stmt.id as usize];
2264        let after = &self.live.after[stmt.id as usize];
2265        let n = inside[0].slots.len();
2266        let rewards: Vec<usize> = self
2267            .rewards
2268            .get(&stmt.id)
2269            .map_or(&[][..], Vec::as_slice)
2270            .iter()
2271            .map(|&r| r as usize)
2272            .collect();
2273        let mut head = live_slots(head_set, n);
2274        head.retain(|i| !rewards.contains(i));
2275        let dead: Vec<SlotId> = head_set.iter_missing(n).collect();
2276        if rewards.len() > MAX_REWARDS {
2277            self.too_large.insert(stmt.id);
2278            return Ok(None);
2279        }
2280        // A world's rewards, where they're plain numbers: whole numbers or
2281        // floats, by their bits, and which are floats.
2282        let reward_of = |slots: &[Value]| -> Option<Raw> {
2283            let (mut bits, mut floats) = ([0u64; MAX_REWARDS], 0u8);
2284            for (i, &r) in rewards.iter().enumerate() {
2285                bits[i] = match &slots[r] {
2286                    Value::Int(n) => n.to_i64()? as u64,
2287                    Value::Float(x) => {
2288                        floats |= 1 << i;
2289                        x.to_bits()
2290                    }
2291                    _ => return None,
2292                };
2293            }
2294            Some((bits, floats))
2295        };
2296        // Each state: its world, with dead slots cleared and rewards at 0,
2297        // found by its other live slots.
2298        let mut states: Vec<Vec<Value>> = Vec::new();
2299        let mut index: FxHashMap<Vec<Value>, usize> = FxHashMap::default();
2300        let mut intern = |mut slots: Vec<Value>, states: &mut Vec<Vec<Value>>| -> usize {
2301            for &d in &dead {
2302                slots[d as usize] = Value::Dead;
2303            }
2304            for &r in &rewards {
2305                slots[r] = Value::Int(0.into());
2306            }
2307            let key: Vec<Value> = head.iter().map(|&i| slots[i].clone()).collect();
2308            *index.entry(key).or_insert_with(|| {
2309                states.push(slots);
2310                states.len() - 1
2311            })
2312        };
2313        let total = total_weight(inside);
2314        let mut start: Vec<f64> = Vec::new();
2315        // With rewards: each state and reward the worlds start with.
2316        let mut starts: Vec<((usize, Raw), f64)> = Vec::new();
2317        for w in inside {
2318            let reward = if rewards.is_empty() {
2319                ([0; MAX_REWARDS], 0)
2320            } else {
2321                match reward_of(&w.slots) {
2322                    Some(r) => r,
2323                    None => return Ok(None),
2324                }
2325            };
2326            let i = intern(w.slots.clone(), &mut states);
2327            if start.len() <= i {
2328                start.resize(i + 1, 0.0);
2329            }
2330            start[i] += w.weight.ratio(total);
2331            starts.push(((i, reward), w.weight.ratio(total)));
2332        }
2333        // With rewards: each state's steps, to each state with how much the
2334        // rewards grow.
2335        let mut moves: Vec<Vec<(u32, Raw, f64)>> = Vec::new();
2336        let mut chain = Chain::default();
2337        let mut exits: Vec<Vec<World>> = Vec::new();
2338        let mut returns: Vec<Vec<Returned>> = Vec::new();
2339        let mut unresolved: Vec<Weight> = Vec::new();
2340        let mut failed: Vec<Failures> = Vec::new();
2341        let mut faulted: Vec<Vec<(World, RuntimeError)>> = Vec::new();
2342        // A body that uses a call's result so far can't be solved as a chain:
2343        // that result changes from round to round.
2344        let saved_pending = std::mem::take(&mut self.pending);
2345        let mut k = 0;
2346        while k < states.len() {
2347            if !self.pending.is_empty() {
2348                self.pending = saved_pending;
2349                return Ok(None);
2350            }
2351            if states.len() > self.config.max_chain_states {
2352                self.pending = saved_pending;
2353                self.too_large.insert(stmt.id);
2354                return Ok(None);
2355            }
2356            let world = World {
2357                slots: states[k].clone(),
2358                constraints: Default::default(),
2359                inherited: Default::default(),
2360                weight: Weight::ONE,
2361                run: 0,
2362                densities: 0,
2363            };
2364            let saved_unresolved = std::mem::replace(&mut self.unresolved, Weight::ZERO);
2365            let saved_lost = std::mem::replace(&mut self.lost, Weight::ZERO);
2366            let saved_failures = std::mem::take(&mut self.failures);
2367            let flow = self.exec_block(f, body, vec![world]);
2368            let left = std::mem::replace(&mut self.unresolved, saved_unresolved);
2369            let lost = std::mem::replace(&mut self.lost, saved_lost);
2370            let failures = std::mem::replace(&mut self.failures, saved_failures);
2371            let mut flow = flow?;
2372            if flow
2373                .next
2374                .iter()
2375                .chain(&flow.continued)
2376                .chain(&flow.broke)
2377                .chain(flow.faulted.iter().map(|(w, _)| w))
2378                .any(|w| w.densities > 0 || !w.constraints.is_empty() || w.slots.iter().any(analytic::contains))
2379                || flow
2380                    .returned
2381                    .iter()
2382                    .any(|(v, _, c)| !c.is_empty() || analytic::contains(v))
2383            {
2384                self.pending = saved_pending;
2385                self.too_large.insert(stmt.id);
2386                return Ok(None);
2387            }
2388            clear_dead(&mut flow.broke, after);
2389            // Leaving: by `break` or `return`, unresolved, ruled out,
2390            // failed, or faulted to a catch outside the loop.
2391            let mut leave = left + lost + failures.weight;
2392            let faults = std::mem::take(&mut flow.faulted);
2393            for (w, _) in &faults {
2394                leave += w.weight;
2395            }
2396            for w in &flow.broke {
2397                leave += w.weight;
2398            }
2399            for (_, w, _) in &flow.returned {
2400                leave += *w;
2401            }
2402            let mut next: FxHashMap<usize, f64> = FxHashMap::default();
2403            let mut steps: Vec<(u32, Raw, f64)> = Vec::new();
2404            for w in flow.next.into_iter().chain(flow.continued) {
2405                let p = w.weight.to_f64();
2406                if p == 0.0 {
2407                    // Too small for the chain's arithmetic.
2408                    self.too_large.insert(stmt.id);
2409                    self.pending = saved_pending;
2410                    return Ok(None);
2411                }
2412                let grows = if rewards.is_empty() {
2413                    ([0; MAX_REWARDS], 0)
2414                } else {
2415                    match reward_of(&w.slots) {
2416                        Some(r) => r,
2417                        None => {
2418                            self.too_large.insert(stmt.id);
2419                            self.pending = saved_pending;
2420                            return Ok(None);
2421                        }
2422                    }
2423                };
2424                let j = intern(w.slots, &mut states);
2425                *next.entry(j).or_insert(0.0) += p;
2426                if !rewards.is_empty() {
2427                    steps.push((j as u32, grows, p));
2428                }
2429            }
2430            moves.push(steps);
2431            let mut next: Vec<(usize, f64)> = next.into_iter().collect();
2432            next.sort_by_key(|&(j, _)| j);
2433            chain.next.push(next);
2434            chain.leave.push(leave.to_f64());
2435            exits.push(flow.broke);
2436            returns.push(flow.returned);
2437            unresolved.push(left);
2438            failed.push(failures);
2439            faulted.push(faults);
2440            k += 1;
2441        }
2442        let waiting = !self.pending.is_empty();
2443        self.pending = saved_pending;
2444        if waiting {
2445            return Ok(None);
2446        }
2447        start.resize(states.len(), 0.0);
2448        // Rewards are iterated, which needs no elimination: only states that
2449        // can't be left are found first.
2450        if !rewards.is_empty() {
2451            if let Some(s) = chain.stuck() {
2452                return Err(self.never_leaves(f, stmt, &states[s], &head));
2453            }
2454            drop(index);
2455            drop(states);
2456            let chain = Rewarded {
2457                rewards: &rewards,
2458                moves: &moves,
2459                exits,
2460                returns,
2461                unresolved,
2462                failed,
2463                faulted,
2464            };
2465            let solved = self.iterate_rewards(stmt, chain, starts, total, cutoff, iterations)?;
2466            if solved.is_none() {
2467                self.too_large.insert(stmt.id);
2468            }
2469            return Ok(solved);
2470        }
2471        let mut steps = MAX_ELIMINATION;
2472        let solution = chain.visits(&start, &mut steps);
2473        self.spend(MAX_ELIMINATION - steps, stmt.span)?;
2474        let visits = match solution {
2475            Solution::Visits(v) => v,
2476            Solution::TooBig => {
2477                self.too_large.insert(stmt.id);
2478                return Ok(None);
2479            }
2480            Solution::Stuck(s) => return Err(self.never_leaves(f, stmt, &states[s], &head)),
2481        };
2482        self.stats.solved_loops += 1;
2483        self.stats.chain_states += states.len() as u64;
2484        let mut flow = Flow::default();
2485        let mut left = Weight::ZERO;
2486        for (k, v) in visits.into_iter().enumerate() {
2487            if v == 0.0 {
2488                continue;
2489            }
2490            let times = total.scale(v);
2491            for mut w in std::mem::take(&mut exits[k]) {
2492                w.weight = w.weight * times;
2493                flow.next.push(w);
2494            }
2495            for (value, w, c) in std::mem::take(&mut returns[k]) {
2496                flow.returned.push((value, w * times, c));
2497            }
2498            left += unresolved[k] * times;
2499            self.failures.absorb(&failed[k], times, None, false);
2500            for (mut w, e) in std::mem::take(&mut faulted[k]) {
2501                w.weight = w.weight * times;
2502                flow.faulted.push((w, e));
2503            }
2504        }
2505        self.unresolved += left;
2506        Ok(Some(flow))
2507    }
2508
2509    /// A loop solved as a chain whose steps carry rewards: the law of each
2510    /// state and reward, step after step, from the worlds that start it, as
2511    /// unrolling would follow them, but without running the body again.
2512    /// Like unrolling, it stops once what's still inside weighs less than
2513    /// `cutoff`, which is then unresolved, or after the most iterations.
2514    /// `None` if a reward doesn't fit the numbers it's followed in.
2515    fn iterate_rewards(
2516        &mut self,
2517        stmt: &Stmt,
2518        chain: Rewarded,
2519        starts: Vec<((usize, Raw), f64)>,
2520        total: Weight,
2521        cutoff: Weight,
2522        iterations: &mut u64,
2523    ) -> Result<Option<Flow>> {
2524        let span = stmt.span;
2525        let count = chain.rewards.len();
2526        // Each reward is followed as a whole number while every value it
2527        // takes is an int, and as a float once one is a float.
2528        let floats = starts
2529            .iter()
2530            .map(|((_, (_, f)), _)| *f)
2531            .chain(chain.moves.iter().flatten().map(|(_, (_, f), _)| *f))
2532            .fold(0u8, |a, f| a | f);
2533        let float: [bool; MAX_REWARDS] = std::array::from_fn(|i| floats & (1 << i) != 0);
2534        let compact = |(bits, floats): &Raw| -> Option<Reward> {
2535            let mut out = *bits;
2536            for i in 0..count {
2537                if float[i] && floats & (1 << i) == 0 {
2538                    let n = bits[i] as i64;
2539                    if n.unsigned_abs() > 1 << 53 {
2540                        return None;
2541                    }
2542                    out[i] = (n as f64).to_bits();
2543                }
2544            }
2545            Some(out)
2546        };
2547        let plus = |a: &Reward, b: &Reward| -> Option<Reward> {
2548            let mut out = [0u64; MAX_REWARDS];
2549            for i in 0..count {
2550                out[i] = if float[i] {
2551                    (f64::from_bits(a[i]) + f64::from_bits(b[i])).to_bits()
2552                } else {
2553                    (a[i] as i64).checked_add(b[i] as i64)? as u64
2554                };
2555            }
2556            Some(out)
2557        };
2558        let value = |r: &Reward, i: usize| -> Value {
2559            if float[i] {
2560                Value::Float(f64::from_bits(r[i]))
2561            } else {
2562                Value::Int((r[i] as i64).into())
2563            }
2564        };
2565        let mut moves: Vec<Vec<(u32, Reward, f64)>> = Vec::with_capacity(chain.moves.len());
2566        for steps in chain.moves {
2567            let mut out = Vec::with_capacity(steps.len());
2568            for (j, d, p) in steps {
2569                let Some(d) = compact(d) else { return Ok(None) };
2570                out.push((*j, d, *p));
2571            }
2572            moves.push(out);
2573        }
2574        let mut current: FxHashMap<(u32, Reward), f64> = FxHashMap::default();
2575        for ((s, r), p) in &starts {
2576            let Some(r) = compact(r) else { return Ok(None) };
2577            *current.entry((*s as u32, r)).or_insert(0.0) += p;
2578        }
2579        // How much weight ran the body from each state; leaves by each of
2580        // the distinct worlds that leave, with each reward, as unrolling's
2581        // merging would add it up; and faulted, by state and reward.
2582        let mut visits = vec![0.0; moves.len()];
2583        let mut templates: Vec<&World> = Vec::new();
2584        let mut template_index: FxHashMap<&[Value], u32> = FxHashMap::default();
2585        let exits: Vec<Vec<(u32, f64)>> = chain
2586            .exits
2587            .iter()
2588            .map(|worlds| {
2589                worlds
2590                    .iter()
2591                    .map(|w| {
2592                        let t = *template_index.entry(&w.slots[..]).or_insert_with(|| {
2593                            templates.push(w);
2594                            (templates.len() - 1) as u32
2595                        });
2596                        (t, w.weight.to_f64())
2597                    })
2598                    .collect()
2599            })
2600            .collect();
2601        drop(template_index);
2602        let mut leaving: FxHashMap<(u32, Reward), f64> = FxHashMap::default();
2603        let mut ran: FxHashMap<(u32, Reward), f64> = FxHashMap::default();
2604        let mut next: FxHashMap<(u32, Reward), f64> = FxHashMap::default();
2605        // When every step adds the same, and every world starts with the
2606        // same reward, as a counter of rounds does, all the worlds have the
2607        // same reward at each step: the weights are a vector by state.
2608        let mut increments = moves.iter().flatten().map(|(_, d, _)| d);
2609        let same = increments.next().copied().filter(|d| increments.all(|e| e == d));
2610        let mut rewards = current.keys().map(|(_, r)| r);
2611        let start = rewards.next().copied().filter(|r| rewards.all(|q| q == r));
2612        if let (Some(d), Some(mut reward)) = (same, start) {
2613            let n = moves.len();
2614            let mut weights = vec![0.0; n];
2615            for (&(s, _), &m) in &current {
2616                weights[s as usize] += m;
2617            }
2618            let mut after = vec![0.0; n];
2619            loop {
2620                let mass: f64 = weights.iter().sum();
2621                if mass == 0.0 {
2622                    break;
2623                }
2624                if total.scale(mass) < cutoff {
2625                    self.unresolved += total.scale(mass);
2626                    break;
2627                }
2628                *iterations += 1;
2629                if *iterations > self.config.max_iterations {
2630                    return Err(RuntimeError::limit(
2631                        span,
2632                        format!("this loop ran {} times without finishing", self.config.max_iterations),
2633                    )
2634                    .with_note(format!(
2635                        "worlds still inside the loop weigh {} of what entered it",
2636                        fmt_prob(mass)
2637                    ))
2638                    .with_help("check that every world can leave the loop"));
2639                }
2640                self.spend(n as u64, span)?;
2641                for (s, &m) in weights.iter().enumerate() {
2642                    if m == 0.0 {
2643                        continue;
2644                    }
2645                    visits[s] += m;
2646                    for &(j, _, p) in &moves[s] {
2647                        after[j as usize] += m * p;
2648                    }
2649                    for &(t, p) in &exits[s] {
2650                        *leaving.entry((t, reward)).or_insert(0.0) += m * p;
2651                    }
2652                    if !chain.faulted[s].is_empty() {
2653                        *ran.entry((s as u32, reward)).or_insert(0.0) += m;
2654                    }
2655                }
2656                std::mem::swap(&mut weights, &mut after);
2657                after.iter_mut().for_each(|w| *w = 0.0);
2658                let Some(r) = plus(&reward, &d) else { return Ok(None) };
2659                reward = r;
2660            }
2661            current.clear();
2662        }
2663        while !current.is_empty() {
2664            let mass: f64 = current.values().sum();
2665            if total.scale(mass) < cutoff {
2666                self.unresolved += total.scale(mass);
2667                break;
2668            }
2669            *iterations += 1;
2670            if *iterations > self.config.max_iterations {
2671                return Err(RuntimeError::limit(
2672                    span,
2673                    format!("this loop ran {} times without finishing", self.config.max_iterations),
2674                )
2675                .with_note(format!(
2676                    "worlds still inside the loop weigh {} of what entered it",
2677                    fmt_prob(mass)
2678                ))
2679                .with_help("check that every world can leave the loop"));
2680            }
2681            self.spend(current.len() as u64, span)?;
2682            for (&(s, r), &m) in &current {
2683                visits[s as usize] += m;
2684                for (j, d, p) in &moves[s as usize] {
2685                    let Some(r) = plus(&r, d) else { return Ok(None) };
2686                    *next.entry((*j, r)).or_insert(0.0) += m * p;
2687                }
2688                for &(t, p) in &exits[s as usize] {
2689                    *leaving.entry((t, r)).or_insert(0.0) += m * p;
2690                }
2691                if !chain.faulted[s as usize].is_empty() {
2692                    *ran.entry((s, r)).or_insert(0.0) += m;
2693                }
2694            }
2695            self.check_worlds(next.len(), span)?;
2696            std::mem::swap(&mut current, &mut next);
2697            next.clear();
2698        }
2699        self.stats.solved_loops += 1;
2700        self.stats.chain_states += moves.len() as u64;
2701        // A world leaving from a state with a reward: what it gained in
2702        // its last step, added to the reward.
2703        let add = |a: Value, b: &Value, budget: &mut crate::dist::Budget| -> Result<Value> {
2704            ops::binary(probl_syntax::ast::BinOp::Add, &a, b, budget).map_err(|e| e.at(span))
2705        };
2706        let mut flow = Flow::default();
2707        let rewarded = |w: &World, r: &Reward, budget: &mut crate::dist::Budget| -> Result<World> {
2708            let mut w = w.clone();
2709            for (i, &slot) in chain.rewards.iter().enumerate() {
2710                if !matches!(w.slots[slot], Value::Dead) {
2711                    w.slots[slot] = add(value(r, i), &w.slots[slot], budget)?;
2712                }
2713            }
2714            Ok(w)
2715        };
2716        let mut leaving: Vec<((u32, Reward), f64)> = leaving.into_iter().collect();
2717        leaving.sort_by_key(|a| a.0);
2718        for ((t, r), m) in leaving {
2719            let mut w = rewarded(templates[t as usize], &r, &mut self.budget)?;
2720            w.weight = total.scale(m);
2721            flow.next.push(w);
2722        }
2723        let mut ran: Vec<((u32, Reward), f64)> = ran.into_iter().collect();
2724        ran.sort_by_key(|a| a.0);
2725        for ((s, r), m) in ran {
2726            let times = total.scale(m);
2727            for (w, e) in &chain.faulted[s as usize] {
2728                let mut w = rewarded(w, &r, &mut self.budget)?;
2729                w.weight = w.weight * times;
2730                flow.faulted.push((w, e.clone()));
2731            }
2732        }
2733        let mut left = Weight::ZERO;
2734        for (s, v) in visits.into_iter().enumerate() {
2735            if v == 0.0 {
2736                continue;
2737            }
2738            let times = total.scale(v);
2739            for (value, w, c) in &chain.returns[s] {
2740                flow.returned.push((value.clone(), *w * times, c.clone()));
2741            }
2742            left += chain.unresolved[s] * times;
2743            self.failures.absorb(&chain.failed[s], times, None, false);
2744        }
2745        self.unresolved += left;
2746        Ok(Some(flow))
2747    }
2748
2749    /// The error for a loop that some worlds can never leave, with one of
2750    /// them.
2751    fn never_leaves(&self, f: FnId, stmt: &Stmt, slots: &[Value], head: &[usize]) -> RuntimeError {
2752        let names = &self.prog.functions[f as usize].slots;
2753        let parts: Vec<String> = head
2754            .iter()
2755            .filter(|&&i| names[i].name != TEMP)
2756            .map(|&i| format!("`{}` is {:?}", names[i].name, slots[i]))
2757            .collect();
2758        let err = RuntimeError::new(stmt.span, "some worlds can never leave this loop")
2759            .with_help("check that every world can leave the loop");
2760        if parts.is_empty() {
2761            err
2762        } else {
2763            err.with_note(format!("for example, the worlds where {}", parts.join(" and ")))
2764        }
2765    }
2766
2767    /// Store each possible value of `d` into `place`, one world per outcome.
2768    fn split_by(&mut self, f: FnId, place: &Place, d: Value, w: World, span: Span, out: &mut Vec<World>) -> Result<()> {
2769        match d {
2770            Value::Dist(_) | Value::Continuous(_) | Value::Counts(_) | Value::Joint(_) if self.sampler.is_some() => {
2771                let v = self.sample(&d);
2772                let mut w = w;
2773                self.assign(f, place, v, &mut w, span)?;
2774                out.push(w);
2775            }
2776            Value::Continuous(family) => {
2777                self.next_latent += 1;
2778                let v = Analytic::new(self.next_latent, *family)
2779                    .value()
2780                    .map_err(|e| e.at(span))?;
2781                let mut w = w;
2782                self.assign(f, place, v, &mut w, span)?;
2783                out.push(w);
2784            }
2785            Value::Dist(dist) => {
2786                self.unresolved += w.weight.scale(dist.missing);
2787                let last = dist.outcomes.len().saturating_sub(1);
2788                let mut w = Some(w);
2789                for (i, (v, p)) in dist.outcomes.iter().enumerate() {
2790                    let base = if i == last {
2791                        w.take().unwrap()
2792                    } else {
2793                        w.clone().unwrap()
2794                    };
2795                    let mut nw = base.scaled(*p);
2796                    if matches!(v, Value::Continuous(_) | Value::Joint(_)) {
2797                        self.split_by(f, place, v.clone(), nw, span, out)?;
2798                    } else {
2799                        self.assign(f, place, v.clone(), &mut nw, span)?;
2800                        out.push(nw);
2801                    }
2802                }
2803            }
2804            Value::Prob(p) => {
2805                return self.split_by(f, place, Dist::bernoulli(p).into_value(), w, span, out);
2806            }
2807            Value::Joint(dist) => self.draw_joint(f, place, &dist, w, span, out)?,
2808            // Enumerating needs its outcomes.
2809            Value::Counts(_) => {
2810                let d = ops::listed(&d, &mut self.budget).map_err(|e| e.at(span))?.into_owned();
2811                return self.split_by(f, place, d, w, span, out);
2812            }
2813            other => {
2814                let mut w = w;
2815                self.assign(f, place, other, &mut w, span)?;
2816                out.push(w);
2817            }
2818        }
2819        Ok(())
2820    }
2821
2822    /// A draw from a distribution whose outcomes own their draws: a world
2823    /// for each outcome, with its draws fresh there.
2824    #[inline(never)]
2825    fn draw_joint(
2826        &mut self,
2827        f: FnId,
2828        place: &Place,
2829        dist: &Dist,
2830        w: World,
2831        span: Span,
2832        out: &mut Vec<World>,
2833    ) -> Result<()> {
2834        self.unresolved += w.weight.scale(dist.missing);
2835        for (v, p) in &dist.outcomes {
2836            let v = crate::joint::fresh(v, &mut self.next_latent);
2837            let mut nw = w.clone().scaled(*p);
2838            self.assign(f, place, v, &mut nw, span)?;
2839            out.push(nw);
2840        }
2841        Ok(())
2842    }
2843
2844    fn check_arity(&self, c: &Closure, given: usize, span: Span) -> Result<()> {
2845        let expected = self.prog.functions[c.func as usize].n_params as usize;
2846        if expected != given {
2847            return Err(RuntimeError::new(
2848                span,
2849                format!("this function takes {expected} argument(s), but {given} were given"),
2850            ));
2851        }
2852        Ok(())
2853    }
2854
2855    /// Convert values at declared type boundaries, recursively through
2856    /// containers. Numeric probability conversions are checked, never clamped.
2857    fn coerce(&mut self, v: Value, ty: &TypeSpec, span: Span) -> Result<Value> {
2858        self.coerce_at(v, ty, span, 0)
2859    }
2860
2861    fn coerce_at(&mut self, v: Value, ty: &TypeSpec, span: Span, depth: usize) -> Result<Value> {
2862        self.budget.work(1).map_err(|e| e.at(span))?;
2863        if depth > 64 {
2864            return Err(RuntimeError::new(
2865                span,
2866                "type conversion nesting exceeds the limit of 64",
2867            ));
2868        }
2869        Ok(match (v, ty) {
2870            (Value::Analytic(_) | Value::Event(_), TypeSpec::Prob) => {
2871                return Err(analytic::unsupported("converting this outcome to prob").at(span));
2872            }
2873            (v @ Value::Float(_), TypeSpec::Int) => Value::Int(
2874                ops::integer(&v, "int conversion", &mut self.budget)
2875                    .map_err(|e| e.at(span))?
2876                    .into_owned(),
2877            ),
2878            // An outcome like `floor(x)` is an int; another is never rounded
2879            // into one.
2880            (Value::Analytic(a), TypeSpec::Int) if a.int => Value::Analytic(a),
2881            (Value::Analytic(_), TypeSpec::Int) => {
2882                return Err(analytic::unsupported("converting this outcome to int").at(span));
2883            }
2884            (Value::Counts(c), TypeSpec::Dist(t)) if **t == TypeSpec::Int => Value::Counts(c),
2885            (Value::Prob(p), TypeSpec::Float) => Value::Float(p),
2886            // A declared type is a contract: failing it isn't a fault of the
2887            // world's values, even in partial mode.
2888            (v @ (Value::Int(_) | Value::Float(_)), TypeSpec::Prob) => {
2889                ops::make_prob(&v).map_err(|e| OpError { fault: None, ..e }.at(span))?
2890            }
2891            (Value::List(xs), TypeSpec::List(t)) => {
2892                self.budget.collection(xs.len() as u128).map_err(|e| e.at(span))?;
2893                let mut out = Vec::with_capacity(xs.len());
2894                for x in xs.iter() {
2895                    out.push(self.coerce_at(x.clone(), t, span, depth + 1)?);
2896                }
2897                Value::list(out)
2898            }
2899            (Value::Map(xs), TypeSpec::Map(kt, vt)) => {
2900                let mut out = BTreeMap::new();
2901                for (k, v) in xs.iter() {
2902                    let key = self.coerce_at(k.clone(), kt, span, depth + 1)?;
2903                    let value = self.coerce_at(v.clone(), vt, span, depth + 1)?;
2904                    if out.insert(key, value).is_some() {
2905                        return Err(RuntimeError::new(span, "map key collision during type conversion"));
2906                    }
2907                }
2908                Value::map(out)
2909            }
2910            (Value::Bag(xs), TypeSpec::Bag(t)) => {
2911                let mut out = BTreeMap::<Value, u64>::new();
2912                for (x, n) in xs.iter() {
2913                    let key = self.coerce_at(x.clone(), t, span, depth + 1)?;
2914                    let count = out.entry(key).or_default();
2915                    *count = count
2916                        .checked_add(*n)
2917                        .ok_or_else(|| RuntimeError::new(span, "bag count overflow during type conversion"))?;
2918                }
2919                Value::multiset(crate::value::Multiset::new(out))
2920            }
2921            (Value::Dist(d), TypeSpec::Dist(t)) => {
2922                let mut out = Vec::with_capacity(d.outcomes.len());
2923                for (x, p) in &d.outcomes {
2924                    out.push((self.coerce_at(x.clone(), t, span, depth + 1)?, *p));
2925                }
2926                ops::combine(out, d.missing, &mut self.budget).map_err(|e| e.at(span))?
2927            }
2928            (Value::Joint(d), TypeSpec::Dist(t)) => {
2929                let mut out = Vec::with_capacity(d.outcomes.len());
2930                for (x, p) in &d.outcomes {
2931                    out.push((self.coerce_at(x.clone(), t, span, depth + 1)?, *p));
2932                }
2933                crate::joint::of(out, d.missing, &mut self.budget).map_err(|e| e.at(span))?
2934            }
2935            (Value::Record(r), TypeSpec::Record(t)) => {
2936                let types: Vec<_> = self.prog.records[*t as usize]
2937                    .fields
2938                    .iter()
2939                    .map(|f| (f.name.clone(), f.ty.clone()))
2940                    .collect();
2941                let mut fields = Vec::with_capacity(r.fields.len());
2942                for (name, value) in &r.fields {
2943                    let value = match types.iter().find(|(n, _)| n == &**name) {
2944                        Some((_, t)) => self.coerce_at(value.clone(), t, span, depth + 1)?,
2945                        None => value.clone(),
2946                    };
2947                    fields.push((name.clone(), value));
2948                }
2949                ops::make_record(r.ty.clone(), fields)
2950            }
2951            (Value::Record(r), TypeSpec::AnonRecord(types)) => {
2952                let mut fields = Vec::with_capacity(r.fields.len());
2953                for (name, value) in &r.fields {
2954                    let value = match types.iter().find(|(n, _)| n == &**name) {
2955                        Some((_, t)) => self.coerce_at(value.clone(), t, span, depth + 1)?,
2956                        None => value.clone(),
2957                    };
2958                    fields.push((name.clone(), value));
2959                }
2960                ops::make_record(r.ty.clone(), fields)
2961            }
2962            (v, _) => v,
2963        })
2964    }
2965
2966    /// Does `v` have the declared type?
2967    fn conforms(&self, v: &Value, ty: &TypeSpec) -> bool {
2968        match (ty, v) {
2969            (TypeSpec::Int, Value::Int(_)) => true,
2970            (TypeSpec::Float, Value::Float(_) | Value::Int(_) | Value::Analytic(_)) => true,
2971            (TypeSpec::Complex, Value::Complex(_)) => true,
2972            (TypeSpec::Prob, Value::Prob(_)) => true,
2973            (TypeSpec::Bool, Value::Bool(_) | Value::Event(_)) => true,
2974            (TypeSpec::Str, Value::Str(_)) => true,
2975            (TypeSpec::Date, Value::Date(_)) => true,
2976            (TypeSpec::Unit, Value::Unit) => true,
2977            (TypeSpec::Function, Value::Closure(_) | Value::Builtin(_)) => true,
2978            (TypeSpec::List(t), Value::List(items)) => items.iter().all(|x| self.conforms(x, t)),
2979            (TypeSpec::List(t), Value::Range(..)) => **t == TypeSpec::Int,
2980            (TypeSpec::Map(k, t), Value::Map(m)) => m.iter().all(|(a, b)| self.conforms(a, k) && self.conforms(b, t)),
2981            (TypeSpec::Bag(t), Value::Bag(b)) => b.keys().all(|x| self.conforms(x, t)),
2982            (TypeSpec::Dist(t), Value::Dist(d) | Value::Joint(d)) => d.outcomes.iter().all(|(x, _)| match x {
2983                Value::Continuous(_) => **t == TypeSpec::Float,
2984                _ => self.conforms(x, t),
2985            }),
2986            (TypeSpec::Dist(t), Value::Continuous(_)) => **t == TypeSpec::Float,
2987            (TypeSpec::Dist(t), Value::Counts(_)) => **t == TypeSpec::Int,
2988            (TypeSpec::Record(r), Value::Record(rec)) => {
2989                let declared = &self.prog.records[*r as usize];
2990                rec.ty.as_deref() == Some(declared.name.as_str())
2991                    && rec.fields.len() == declared.fields.len()
2992                    && declared
2993                        .fields
2994                        .iter()
2995                        .all(|f| rec.get(&f.name).is_some_and(|x| self.conforms(x, &f.ty)))
2996            }
2997            (TypeSpec::Enum(e), Value::Enum(x)) => x.ty == *e,
2998            (TypeSpec::AnonRecord(fields), Value::Record(rec)) => {
2999                rec.fields.len() == fields.len()
3000                    && rec
3001                        .fields
3002                        .iter()
3003                        .zip(fields)
3004                        .all(|((n, x), (m, t))| &**n == m.as_str() && self.conforms(x, t))
3005            }
3006            _ => false,
3007        }
3008    }
3009
3010    // ── Calls ────────────────────────────────────────────────────────────
3011
3012    /// Run a function as a sub-simulation and return the unnormalized
3013    /// distribution of its result. Results are memoized unless the function
3014    /// prints: its behaviour depends only on its arguments and the outside
3015    /// values it reads, which are all in `key`.
3016    ///
3017    /// When enumerating, a call can come back to itself: the same function,
3018    /// with the same arguments, while it's still running. Its result is then
3019    /// solved by iteration (section 6). The call that comes back gets the
3020    /// result so far, which starts with nothing resolved, and the call it
3021    /// came back to runs again with each new result, until the weight still
3022    /// waiting on it is below ε (`run_call`).
3023    fn call(&mut self, func: FnId, key: Vec<Value>, span: Span) -> Result<Arc<CallResult>> {
3024        self.stats.calls += 1;
3025        let prog = self.prog;
3026        let fun = &prog.functions[func as usize];
3027        // When sampling, every call makes its own choices (section 14).
3028        let sampling = self.sampler.is_some();
3029        // A cached result does not prove a callback's effects were permitted:
3030        // a draw or observation can collapse to one outcome of weight one.
3031        let memoizable = self.config.memoizing && !sampling && !fun.effects.prints && self.callback.is_none();
3032        let call_key = (func, key);
3033        if memoizable {
3034            if let Some(r) = self.memo.get(&call_key) {
3035                self.stats.memo_hits += 1;
3036                self.observed |= r.observed;
3037                return Ok(r.clone());
3038            }
3039        }
3040        if !sampling {
3041            if let Some(&depth) = self.active.get(&call_key) {
3042                if fun.effects.prints {
3043                    return Err(RuntimeError::new(
3044                        span,
3045                        format!("`{}` comes back to itself, and prints", describe_call(fun, &call_key.1)),
3046                    )
3047                    .with_note("a call that comes back to itself runs several times, until its result settles, and would print each time")
3048                    .with_help("remove the `print`, or print the result where the call is made"));
3049                }
3050                let so_far = match self.approx.get(&call_key) {
3051                    Some(a) if self.frames_at.get(a.head) == Some(&a.frame) => a.result.clone(),
3052                    _ => {
3053                        let result = Arc::new(CallResult::waiting(depth));
3054                        let approx = Approx {
3055                            result: result.clone(),
3056                            head: depth,
3057                            frame: self.frames_at[depth],
3058                            round: self.rounds_at[depth],
3059                        };
3060                        self.approx.insert(call_key, approx);
3061                        result
3062                    }
3063                };
3064                return Ok(so_far);
3065            }
3066            // Worked out already, in this round of the call it waits on.
3067            if let Some(a) = self.approx.get(&call_key) {
3068                if self.frames_at.get(a.head) == Some(&a.frame) && self.rounds_at.get(a.head) == Some(&a.round) {
3069                    self.stats.memo_hits += 1;
3070                    return Ok(a.result.clone());
3071                }
3072            }
3073        }
3074        if self.depth >= self.config.max_call_depth {
3075            return Err(RuntimeError::limit(
3076                span,
3077                format!("calls are nested more than {} deep", self.config.max_call_depth),
3078            )
3079            .with_help("check for recursion that doesn't stop, or write it as a loop"));
3080        }
3081        let mut slots = vec![Value::Dead; fun.n_slots()];
3082        let n = fun.n_params as usize;
3083        for (i, v) in call_key.1[..n].iter().enumerate() {
3084            slots[i] = v.clone();
3085        }
3086        for (cap, v) in fun.captures.iter().zip(&call_key.1[n..]) {
3087            slots[cap.slot as usize] = v.clone();
3088        }
3089        let depth = self.depth + 1;
3090        if !sampling {
3091            self.active.insert(call_key.clone(), depth);
3092            if self.frames_at.len() <= depth {
3093                self.frames_at.resize(depth + 1, 0);
3094                self.rounds_at.resize(depth + 1, 0);
3095            }
3096            self.last_number += 1;
3097            self.frames_at[depth] = self.last_number;
3098        }
3099        let result = self.run_call(func, &call_key, slots, depth, span);
3100        if !sampling {
3101            self.active.remove(&call_key);
3102            // Results that waited on this frame no longer stand for anything.
3103            self.frames_at[depth] = 0;
3104        }
3105        let result = result?;
3106        // A result still waiting on a call being solved isn't final.
3107        if memoizable
3108            && result.pending.is_empty()
3109            && result
3110                .outcomes
3111                .iter()
3112                .all(|(v, _, c)| c.is_empty() && !analytic::contains(v))
3113        {
3114            if self.memo.len() >= self.config.max_cached_calls {
3115                self.memo.clear();
3116            }
3117            self.memo.insert(call_key, result.clone());
3118        }
3119        Ok(result)
3120    }
3121
3122    /// Run a call's body from `slots`, at `depth` in the stack, and return
3123    /// its result (section 6).
3124    ///
3125    /// If the call came back to itself, and waits on no call further down
3126    /// the stack, it's solved here: round after round, each using the result
3127    /// of the round before, until the weight still waiting on it is below ε.
3128    /// That weight stays unresolved. If it waits on a call further down,
3129    /// it's part of that call's solving instead: what waits on it waits on
3130    /// that call, and its result stands for it until that call's next round.
3131    fn run_call(
3132        &mut self,
3133        func: FnId,
3134        call_key: &CallKey,
3135        slots: Vec<Value>,
3136        depth: usize,
3137        span: Span,
3138    ) -> Result<Arc<CallResult>> {
3139        let fun = &self.prog.functions[func as usize];
3140        let sampling = self.sampler.is_some();
3141        let mut inherited = std::collections::BTreeSet::new();
3142        for v in &slots {
3143            analytic::collect_ids(v, &mut inherited);
3144        }
3145        let inherited = Arc::new(inherited);
3146        let mut rounds: u64 = 0;
3147        let mut before = f64::INFINITY;
3148        loop {
3149            if !sampling {
3150                self.last_number += 1;
3151                self.rounds_at[depth] = self.last_number;
3152            }
3153            let saved_unresolved = std::mem::replace(&mut self.unresolved, Weight::ZERO);
3154            let saved_lost = std::mem::replace(&mut self.lost, Weight::ZERO);
3155            let saved_observed = std::mem::replace(&mut self.observed, false);
3156            let saved_pending = std::mem::take(&mut self.pending);
3157            let saved_failures = std::mem::take(&mut self.failures);
3158            self.depth += 1;
3159            let flow = self.exec_block(
3160                func,
3161                &fun.body,
3162                vec![World {
3163                    slots: slots.clone(),
3164                    constraints: Default::default(),
3165                    inherited: inherited.clone(),
3166                    weight: Weight::ONE,
3167                    run: 0,
3168                    densities: 0,
3169                }],
3170            );
3171            self.depth -= 1;
3172            let unresolved = std::mem::replace(&mut self.unresolved, saved_unresolved);
3173            let lost = std::mem::replace(&mut self.lost, saved_lost);
3174            let observed = std::mem::replace(&mut self.observed, saved_observed);
3175            let pending = std::mem::replace(&mut self.pending, saved_pending);
3176            let failures = std::mem::replace(&mut self.failures, saved_failures);
3177            self.observed |= observed;
3178            let flow = flow.map_err(|e| match fun.kind {
3179                FnKind::Named => e.with_note(format!("in a call to `{}`", fun.name)),
3180                FnKind::Simulate => e.with_note("inside a `simulate` block"),
3181                FnKind::Lambda => e.with_note("inside a lambda"),
3182                FnKind::Main => e,
3183            })?;
3184            uncaught(&flow)?;
3185            let mut result = CallResult {
3186                outcomes: merge_values(flow.returned),
3187                unresolved,
3188                lost,
3189                observed,
3190                pending,
3191                failures,
3192            };
3193            if let Some(head) = result.pending.iter().map(|&(d, _)| d).filter(|&d| d < depth).min() {
3194                // Part of the solving of the call at `head`.
3195                let waiting = Weight::sum(result.pending.iter().map(|&(_, w)| w));
3196                result.pending = vec![(head, waiting)];
3197                let result = Arc::new(result);
3198                let approx = Approx {
3199                    result: result.clone(),
3200                    head,
3201                    frame: self.frames_at[head],
3202                    round: self.rounds_at[head],
3203                };
3204                self.approx.insert(call_key.clone(), approx);
3205                return Ok(result);
3206            }
3207            let waiting = result.take_pending(depth);
3208            if waiting.is_zero() {
3209                return Ok(Arc::new(result));
3210            }
3211            // It came back to itself: solved here.
3212            if rounds == 0 {
3213                self.stats.solved_calls += 1;
3214            }
3215            let w = waiting.to_f64();
3216            if w <= self.config.epsilon {
3217                self.settle(depth);
3218                return Ok(Arc::new(result));
3219            }
3220            rounds += 1;
3221            self.stats.call_rounds += 1;
3222            if w >= before {
3223                return Err(RuntimeError::new(
3224                    span,
3225                    format!(
3226                        "`{}` never returns for some of its worlds",
3227                        describe_call(fun, &call_key.1)
3228                    ),
3229                )
3230                .with_note(format!("{} of its weight keeps coming back to it", fmt_prob(w)))
3231                .with_help("check that the recursion can end"));
3232            }
3233            if rounds >= self.config.max_iterations {
3234                return Err(RuntimeError::limit(
3235                    span,
3236                    format!(
3237                        "`{}` came back to itself {} times without settling",
3238                        describe_call(fun, &call_key.1),
3239                        self.config.max_iterations
3240                    ),
3241                )
3242                .with_note(format!("{} of its weight is still waiting on it", fmt_prob(w)))
3243                .with_help("check that the recursion can end, or raise `@epsilon`"));
3244            }
3245            before = w;
3246            result.pending.push((depth, waiting));
3247            let approx = Approx {
3248                result: Arc::new(result),
3249                head: depth,
3250                frame: self.frames_at[depth],
3251                round: self.rounds_at[depth],
3252            };
3253            self.approx.insert(call_key.clone(), approx);
3254        }
3255    }
3256
3257    /// The call at `depth` is solved: the results that waited on it, from
3258    /// its last round, are final too if the weight still waiting in them is
3259    /// below ε. The others are dropped.
3260    fn settle(&mut self, depth: usize) {
3261        let frame = self.frames_at[depth];
3262        let epsilon = self.config.epsilon;
3263        let settled: Vec<(CallKey, CallResult)> = self
3264            .approx
3265            .iter()
3266            .filter(|(_, a)| a.head == depth && a.frame == frame)
3267            .filter(|(_, a)| a.result.pending.iter().all(|&(_, w)| w.to_f64() <= epsilon))
3268            .map(|(key, a)| {
3269                let mut result = (*a.result).clone();
3270                // What's still waiting stays unresolved.
3271                result.pending.clear();
3272                (key.clone(), result)
3273            })
3274            .collect();
3275        self.approx.retain(|_, a| a.head != depth);
3276        if !self.config.memoizing {
3277            return;
3278        }
3279        for (key, result) in settled {
3280            if !self.active.contains_key(&key)
3281                && !self.prog.functions[key.0 as usize].effects.prints
3282                && result
3283                    .outcomes
3284                    .iter()
3285                    .all(|(v, _, c)| c.is_empty() && !analytic::contains(v))
3286            {
3287                self.memo.insert(key, Arc::new(result));
3288            }
3289        }
3290    }
3291
3292    /// `env`: the host's environment snapshot, as a `map[str, str]`. It's
3293    /// the same value every time, so it's made once.
3294    fn env(&mut self, span: Span) -> Result<Value> {
3295        if let Some(env) = &self.env {
3296            return Ok(env.clone());
3297        }
3298        if self.config.env.len() > self.budget.max_collection {
3299            return Err(RuntimeError::limit(
3300                span,
3301                "the environment has more variables than the collection size limit",
3302            ));
3303        }
3304        let mut entries = BTreeMap::new();
3305        for (name, text) in self.config.env.iter() {
3306            if name.len().max(text.len()) > self.budget.max_string_bytes {
3307                return Err(RuntimeError::limit(
3308                    span,
3309                    format!("the environment variable `{name}` is longer than the string size limit"),
3310                ));
3311            }
3312            entries.insert(Value::str(name), Value::str(text));
3313        }
3314        let env = Value::map(entries);
3315        self.env = Some(env.clone());
3316        Ok(env)
3317    }
3318
3319    /// `simulate { … }`: run a block as a separate model and return its
3320    /// normalized distribution (docs/semantics.md, section 8).
3321    fn simulate(&mut self, func: FnId, key: Vec<Value>, span: Span) -> Result<Value> {
3322        if key.iter().any(analytic::contains) {
3323            return Err(analytic::unsupported("capturing an analytic draw inside `simulate`").at(span));
3324        }
3325        if !self.simulated.is_empty() {
3326            if let Some(v) = self.simulated.get(&(func, key.clone())) {
3327                self.stats.memo_hits += 1;
3328                return Ok(v.clone());
3329            }
3330        }
3331        let call_key = self.config.memoizing.then(|| (func, key.clone()));
3332        let saved_observed = self.observed;
3333        let saved_densities = self.densities;
3334        let saved_local = (self.local_density_at.take(), self.returned_densities.take());
3335        // Enumerated, even when sampling (section 14).
3336        let sampler = self.sampler.take();
3337        let callback = self.callback.take();
3338        self.nested += 1;
3339        let result = self.call(func, key, span);
3340        self.nested -= 1;
3341        self.sampler = sampler;
3342        self.callback = callback;
3343        // Observations inside `simulate` condition its result only.
3344        self.observed = saved_observed;
3345        self.densities = saved_densities;
3346        let local_density_at = std::mem::replace(&mut self.local_density_at, saved_local.0);
3347        self.returned_densities = saved_local.1;
3348        let result = result?;
3349        if let Some(span) = local_density_at.filter(|_| !result.unresolved.is_zero()) {
3350            return Err(OpError::unsupported(
3351                "observing a value from a continuous distribution while some probability is unresolved isn't supported when enumerating",
3352            )
3353            .help("the unresolved worlds could have any density at the value, so nothing would bound the evidence; sample the model with `@mode sample(runs: 10_000)`")
3354            .at(span));
3355        }
3356        if !result.pending.is_empty() {
3357            return Err(OpError::unsupported(
3358                "a `simulate` block that comes back to a call still running isn't supported yet",
3359            )
3360            .at(span));
3361        }
3362        // It's built whole or not at all: a fault on one of its paths, not
3363        // caught inside it, fails the world that runs it.
3364        if let Some(g) = result.failures.groups.first() {
3365            return Err(g.error.clone());
3366        }
3367        let resolved = Weight::sum(result.outcomes.iter().map(|(_, w, _)| *w));
3368        if resolved.is_zero() {
3369            let message = if result.unresolved.is_zero() {
3370                "every world in this `simulate` was ruled out by `observe`"
3371            } else {
3372                "this `simulate` left all of its weight unresolved"
3373            };
3374            return Err(RuntimeError::new(span, message));
3375        }
3376        let denom = resolved + result.unresolved;
3377        let pairs: Vec<(Value, f64)> = result
3378            .outcomes
3379            .iter()
3380            .map(|(v, w, _)| (v.clone(), w.ratio(denom)))
3381            .collect();
3382        let missing = result.unresolved.ratio(denom);
3383        // Outcomes of its own draws make a joint distribution.
3384        if pairs.iter().any(|(v, _)| analytic::contains(v)) {
3385            let v = crate::joint::of(pairs, missing, &mut self.budget).map_err(|e| e.at(span))?;
3386            if let Some(call_key) = call_key {
3387                if self.simulated.len() >= self.config.max_cached_calls {
3388                    self.simulated.clear();
3389                }
3390                self.simulated.insert(call_key, v.clone());
3391            }
3392            return Ok(v);
3393        }
3394        ops::combine(pairs, missing, &mut self.budget).map_err(|e| e.at(span))
3395    }
3396
3397    /// Call a closure that must not split worlds (used by `map`, `filter`, …).
3398    fn call_pure(&mut self, closure: &Value, args: Vec<Value>, what: &'static str, span: Span) -> Result<Value> {
3399        if let Value::Builtin(b) = closure {
3400            return self.builtin_values(*b, &args, &[], Weight::ONE, span);
3401        }
3402        let Value::Closure(c) = closure else {
3403            return Err(RuntimeError::new(
3404                span,
3405                format!(
3406                    "`{what}` needs a function, like `x -> x * 2`, found {}",
3407                    ops::article(&closure.kind())
3408                ),
3409            ));
3410        };
3411        self.check_arity(c, args.len(), span)?;
3412        let mut key = args;
3413        key.extend(c.captured.iter().cloned());
3414        let previous = self.callback.replace(what);
3415        let result = self.call(c.func, key, span);
3416        self.callback = previous;
3417        let result = result?;
3418        if !result.pending.is_empty() {
3419            return Err(OpError::unsupported(format!(
3420                "a function given to `{what}` that comes back to a call still running isn't supported yet"
3421            ))
3422            .at(span));
3423        }
3424        // A callback that faults fails the whole operation, in the world
3425        // that runs it.
3426        if let Some(g) = result.failures.groups.first() {
3427            return Err(g.error.clone());
3428        }
3429        match result.outcomes.as_slice() {
3430            [(v, p, c)] if c.is_empty() && (p.to_f64() - 1.0).abs() < 1e-12 && result.unresolved.is_zero() => {
3431                Ok(v.clone())
3432            }
3433            _ => Err(RuntimeError::new(
3434                span,
3435                format!("the function given to `{what}` can't branch on chances, draw values or observe"),
3436            )),
3437        }
3438    }
3439
3440    // ── Places ───────────────────────────────────────────────────────────
3441
3442    fn check_callback_effect(&self, span: Span) -> Result<()> {
3443        match self.callback {
3444            Some(what) => Err(RuntimeError::new(
3445                span,
3446                format!("the function given to `{what}` can't branch on chances, draw values or observe"),
3447            )
3448            .with_help("use a loop for probabilistic traversal, or `simulate` for a local distribution")),
3449            None => Ok(()),
3450        }
3451    }
3452
3453    fn record_place_type(&self, place: &Place, w: &World) -> Option<TypeSpec> {
3454        if place.path.is_empty() {
3455            return None;
3456        }
3457        let Value::Record(r) = &w.slots[place.slot as usize] else {
3458            return None;
3459        };
3460        let id = self
3461            .prog
3462            .records
3463            .iter()
3464            .position(|t| Some(t.name.as_str()) == r.ty.as_deref())?;
3465        let mut ty = TypeSpec::Record(id as u32);
3466        for element in &place.path {
3467            ty = match (element, ty) {
3468                (PathElem::Field(name), TypeSpec::Record(id)) => self.prog.records[id as usize]
3469                    .fields
3470                    .iter()
3471                    .find(|f| f.name == *name)?
3472                    .ty
3473                    .clone(),
3474                (PathElem::Field(name), TypeSpec::AnonRecord(fields)) => fields.into_iter().find(|(n, _)| n == name)?.1,
3475                (PathElem::Index(_), TypeSpec::List(t) | TypeSpec::Map(_, t)) => *t,
3476                _ => return None,
3477            };
3478        }
3479        Some(ty)
3480    }
3481
3482    fn assign(&mut self, f: FnId, place: &Place, v: Value, w: &mut World, span: Span) -> Result<()> {
3483        if place.path.is_empty() {
3484            w.slots[place.slot as usize] = v;
3485            return Ok(());
3486        }
3487        let v = match self.record_place_type(place, w) {
3488            Some(ty) => {
3489                let v = self.coerce(v, &ty, span)?;
3490                if !self.conforms(&v, &ty) {
3491                    let found = self.described(&v);
3492                    return Err(RuntimeError::new(
3493                        span,
3494                        format!("assigned field should be {}, found {found}", ty.describe(self.prog)),
3495                    ));
3496                }
3497                v
3498            }
3499            None => v,
3500        };
3501        let mut keys = Vec::with_capacity(place.path.len());
3502        for elem in &place.path {
3503            keys.push(match elem {
3504                PathElem::Field(name) => PathKey::Field(name.as_str()),
3505                PathElem::Index(e) => {
3506                    let key = self.eval(f, e, w)?;
3507                    if analytic::contains(&key) {
3508                        return Err(analytic::unsupported("an analytic assignment index").at(e.span));
3509                    }
3510                    PathKey::Index(key)
3511                }
3512            });
3513        }
3514        let target = &mut w.slots[place.slot as usize];
3515        if matches!(target, Value::Dead) {
3516            return Err(self.no_value(f, place.slot, span));
3517        }
3518        update(target, &keys, v, &mut self.budget).map_err(|e| e.at(span))
3519    }
3520
3521    /// Convert a variable to its declared type, or fail.
3522    fn check_slot(&mut self, f: FnId, slot: SlotId, ty: &TypeSpec, w: &mut World, span: Span) -> Result<()> {
3523        w.slots[slot as usize] = self.coerce(w.slots[slot as usize].clone(), ty, span)?;
3524        if self.conforms(&w.slots[slot as usize], ty) {
3525            return Ok(());
3526        }
3527        let name = &self.prog.functions[f as usize].slots[slot as usize].name;
3528        let what = if name == TEMP {
3529            "the result".to_string()
3530        } else {
3531            format!("`{name}`")
3532        };
3533        let found = self.described(&w.slots[slot as usize]);
3534        Err(RuntimeError::new(
3535            span,
3536            format!(
3537                "{what} should be {}, but it's {}",
3538                ops::article(&ty.describe(self.prog)),
3539                ops::article(&found)
3540            ),
3541        ))
3542    }
3543
3544    /// A value that fails a declared type, described as `typeof` would, or
3545    /// by its kind if that's over a limit.
3546    fn described(&mut self, v: &Value) -> String {
3547        match crate::type_name::of(v, &self.prog.enums, &mut self.budget) {
3548            Ok(Value::Str(s)) => s.to_string(),
3549            _ => v.kind(),
3550        }
3551    }
3552
3553    /// `place = value`, where `value` applies a method like `push` to
3554    /// `place`, then the check of a variable with a declared type. Out of
3555    /// line: every statement runs the statement loop.
3556    #[inline(never)]
3557    fn update(
3558        &mut self,
3559        f: FnId,
3560        place: &Place,
3561        value: &Expr,
3562        check: Option<&UpdateCheck>,
3563        w: &mut World,
3564        span: Span,
3565    ) -> Result<()> {
3566        let ExprKind::Builtin { func, args, named } = &value.kind else {
3567            return Err(RuntimeError::new(span, "internal error: an update that isn't a method"));
3568        };
3569        let mut values = Vec::with_capacity(args.len() + named.len());
3570        for a in args.iter().chain(named.iter().map(|(_, a)| a)) {
3571            values.push(self.eval(f, a, w)?);
3572        }
3573        let field = self.record_place_type(place, w);
3574        let types = match check {
3575            Some(c) => c.args.as_deref().map(Some),
3576            None => Some(None),
3577        };
3578        if let (None, Some(types), true) = (&field, types, named.is_empty()) {
3579            if self.in_place(f, place, *func, &mut values, types, w, span)? {
3580                return Ok(());
3581            }
3582        }
3583        let names: Vec<String> = named.iter().map(|(n, _)| n.clone()).collect();
3584        let mut v = self.builtin_values(*func, &values, &names, w.weight, span)?;
3585        if let Some(ty) = field {
3586            v = self.coerce(v, &ty, value.span)?;
3587        }
3588        self.assign(f, place, v, w, span)?;
3589        if let Some(c) = check {
3590            self.check_slot(f, place.slot, &c.ty, w, span)?;
3591        }
3592        Ok(())
3593    }
3594
3595    /// Change the list or map at `place` as `func` would, given the values
3596    /// of its arguments, the collection first, without copying it unless
3597    /// something else holds it. What it adds is checked against `types`, as
3598    /// a check of the whole variable would. False, with nothing changed,
3599    /// where it can't, as for a distribution, which the method lifts over.
3600    #[allow(clippy::too_many_arguments)]
3601    fn in_place(
3602        &mut self,
3603        f: FnId,
3604        place: &Place,
3605        func: Builtin,
3606        values: &mut [Value],
3607        types: Option<&[Option<TypeSpec>]>,
3608        w: &mut World,
3609        span: Span,
3610    ) -> Result<bool> {
3611        let plain = |v: &Value| {
3612            !v.is_uncertain()
3613                && !matches!(
3614                    v,
3615                    Value::Analytic(_) | Value::Event(_) | Value::Delayed(_) | Value::Dead
3616                )
3617        };
3618        let len = match (func, &values[0]) {
3619            (Builtin::Push | Builtin::Insert | Builtin::Remove | Builtin::DropLast, Value::List(items)) => items.len(),
3620            (Builtin::Insert | Builtin::Remove, Value::Map(m)) => m.len(),
3621            _ => return Ok(false),
3622        };
3623        if !values[1..].iter().all(plain) {
3624            return Ok(false);
3625        }
3626        let at = |e: OpError| e.at(span);
3627        // Whatever can fail does before anything changes.
3628        let list = matches!(values[0], Value::List(_));
3629        let index = match func {
3630            Builtin::Push | Builtin::Insert => {
3631                self.budget.collection(len as u128 + 1).map_err(at)?;
3632                match (func, list) {
3633                    (Builtin::Insert, true) => {
3634                        Some(ops::as_index(&values[1], len as u128 + 1, &mut self.budget).map_err(at)?)
3635                    }
3636                    _ => None,
3637                }
3638            }
3639            Builtin::Remove if list => Some(ops::as_index(&values[1], len as u128, &mut self.budget).map_err(at)?),
3640            _ => None,
3641        };
3642        if let Some(types) = types {
3643            for (i, ty) in types.iter().enumerate() {
3644                let Some(ty) = ty else { continue };
3645                let given = &values[i + 1];
3646                let v = self.coerce(given.clone(), ty, span)?;
3647                // The assignment's check of the whole variable reports it.
3648                if !self.conforms(&v, ty) {
3649                    return Ok(false);
3650                }
3651                // A key that converts could meet one already there, which the
3652                // check of the whole map reports.
3653                if !list && i == 0 && v != *given {
3654                    return Ok(false);
3655                }
3656                values[i + 1] = v;
3657            }
3658        }
3659        let mut keys = Vec::with_capacity(place.path.len());
3660        for elem in &place.path {
3661            keys.push(match elem {
3662                PathElem::Field(name) => PathKey::Field(name.as_str()),
3663                PathElem::Index(e) => {
3664                    let key = self.eval(f, e, w)?;
3665                    if analytic::contains(&key) {
3666                        return Ok(false);
3667                    }
3668                    PathKey::Index(key)
3669                }
3670            });
3671        }
3672        // The value read for the method is another reference to the
3673        // collection, which would make changing it copy it.
3674        values[0] = Value::Unit;
3675        let target = place_mut(&mut w.slots[place.slot as usize], &keys, &mut self.budget).map_err(at)?;
3676        let budget = &mut self.budget;
3677        match (func, target) {
3678            (Builtin::Push, Value::List(items)) => unshared(items, |items| items.len(), budget)
3679                .map_err(at)?
3680                .push(std::mem::replace(&mut values[1], Value::Unit)),
3681            (Builtin::Insert, Value::List(items)) => unshared(items, |items| items.len(), budget).map_err(at)?.insert(
3682                index.unwrap_or_default() as usize,
3683                std::mem::replace(&mut values[2], Value::Unit),
3684            ),
3685            (Builtin::Remove, Value::List(items)) => {
3686                unshared(items, |items| items.len(), budget)
3687                    .map_err(at)?
3688                    .remove(index.unwrap_or_default() as usize);
3689            }
3690            (Builtin::DropLast, Value::List(items)) => {
3691                unshared(items, |items| items.len(), budget).map_err(at)?.pop();
3692            }
3693            (Builtin::Insert, Value::Map(m)) => {
3694                let value = std::mem::replace(&mut values[2], Value::Unit);
3695                unshared(m, |m| m.len(), budget)
3696                    .map_err(at)?
3697                    .insert(std::mem::replace(&mut values[1], Value::Unit), value);
3698            }
3699            (Builtin::Remove, Value::Map(m)) => {
3700                unshared(m, |m| m.len(), budget).map_err(at)?.remove(&values[1]);
3701            }
3702            _ => {
3703                return Err(RuntimeError::new(
3704                    span,
3705                    "internal error: an update found another collection",
3706                ));
3707            }
3708        }
3709        Ok(true)
3710    }
3711
3712    fn read_place(&mut self, f: FnId, place: &Place, w: &World, span: Span) -> Result<Value> {
3713        let mut v = self.slot(f, place.slot, w, span)?;
3714        for elem in &place.path {
3715            v = match elem {
3716                PathElem::Field(name) => ops::field(&v, name, &mut self.budget),
3717                PathElem::Index(e) => {
3718                    let i = self.eval(f, e, w)?;
3719                    ops::index(&v, &i, &mut self.budget)
3720                }
3721            }
3722            .map_err(|e| e.at(span))?;
3723        }
3724        Ok(v)
3725    }
3726
3727    /// Add the values `callee` captures to `out`, from the slots of `f` that
3728    /// hold them. A global it holds only for the function values it may call
3729    /// can have no value yet, as can any, with `lenient`: that's an error
3730    /// only where the global is read.
3731    #[allow(clippy::too_many_arguments)]
3732    fn capture(
3733        &mut self,
3734        out: &mut Vec<Value>,
3735        f: FnId,
3736        callee: FnId,
3737        args: &[SlotId],
3738        w: &World,
3739        span: Span,
3740        lenient: bool,
3741    ) -> Result<()> {
3742        let prog = self.prog;
3743        let captures = &prog.functions[callee as usize].captures;
3744        for (&s, capture) in args.iter().zip(captures) {
3745            out.push(match &w.slots[s as usize] {
3746                Value::Dead if lenient || capture.for_values => Value::Dead,
3747                _ => self.slot(f, s, w, span)?,
3748            });
3749        }
3750        Ok(())
3751    }
3752
3753    fn slot(&mut self, f: FnId, s: SlotId, w: &World, span: Span) -> Result<Value> {
3754        match &w.slots[s as usize] {
3755            Value::Dead => Err(self.no_value(f, s, span)),
3756            Value::Delayed(_) => {
3757                let name = &self.prog.functions[f as usize].slots[s as usize].name;
3758                Err(
3759                    OpError::internal("internal error: a variable was read before it was drawn")
3760                        .at(span)
3761                        .with_note(format!(
3762                            "`{name}`'s draw was delayed for an exact update (docs/semantics.md, section 14)"
3763                        )),
3764                )
3765            }
3766            v => analytic::resolve(v, &w.constraints, &mut self.budget).map_err(|e| e.at(span)),
3767        }
3768    }
3769
3770    fn no_value(&self, f: FnId, s: SlotId, span: Span) -> RuntimeError {
3771        let name = &self.prog.functions[f as usize].slots[s as usize].name;
3772        RuntimeError::new(span, format!("`{name}` has no value here")).with_help(
3773            "it's used before it's given a value; if a function reads it, call the function after the variable is set",
3774        )
3775    }
3776
3777    // ── Expressions ──────────────────────────────────────────────────────
3778
3779    fn eval_expected(&mut self, f: FnId, e: &Expr, w: &World, ty: &TypeSpec) -> Result<Value> {
3780        match (&e.kind, ty) {
3781            (ExprKind::List(xs), TypeSpec::List(t)) => {
3782                self.budget.collection(xs.len() as u128).map_err(|err| err.at(e.span))?;
3783                let mut out = Vec::with_capacity(xs.len());
3784                for x in xs {
3785                    out.push(self.eval_expected(f, x, w, t)?);
3786                }
3787                Ok(Value::list(out))
3788            }
3789            (ExprKind::Map(xs), TypeSpec::Map(kt, vt)) => {
3790                let mut out = BTreeMap::new();
3791                for (k, v) in xs {
3792                    let key = self.eval_expected(f, k, w, kt)?;
3793                    if analytic::contains(&key) {
3794                        return Err(analytic::unsupported("analytic map keys").at(k.span));
3795                    }
3796                    if key.is_uncertain() {
3797                        return Err(RuntimeError::new(k.span, "map keys can't be distributions"));
3798                    }
3799                    out.insert(key, self.eval_expected(f, v, w, vt)?);
3800                }
3801                Ok(Value::map(out))
3802            }
3803            (ExprKind::Record { ty: None, fields }, TypeSpec::AnonRecord(types)) => {
3804                let mut out = Vec::with_capacity(fields.len());
3805                for (n, e) in fields {
3806                    let v = match types.iter().find(|(name, _)| name == n) {
3807                        Some((_, ty)) => self.eval_expected(f, e, w, ty)?,
3808                        None => self.eval(f, e, w)?,
3809                    };
3810                    out.push((Arc::from(n.as_str()), v));
3811                }
3812                Ok(ops::make_record(None, out))
3813            }
3814            _ => {
3815                let v = self.eval(f, e, w)?;
3816                self.coerce(v, ty, e.span)
3817            }
3818        }
3819    }
3820
3821    /// A shared field type supplies conversion context even when the record is
3822    /// represented by a finite distribution. Heterogeneous records have no
3823    /// shared context, and their individual declarations are checked below.
3824    fn record_field_type(&self, value: &Value, name: &str) -> Option<TypeSpec> {
3825        match value {
3826            Value::Record(r) => self
3827                .prog
3828                .records
3829                .iter()
3830                .find(|t| Some(t.name.as_str()) == r.ty.as_deref())?
3831                .fields
3832                .iter()
3833                .find(|f| f.name == name)
3834                .map(|f| f.ty.clone()),
3835            Value::Dist(d) => {
3836                let first = self.record_field_type(&d.outcomes.first()?.0, name)?;
3837                d.outcomes
3838                    .iter()
3839                    .all(|(v, _)| self.record_field_type(v, name).as_ref() == Some(&first))
3840                    .then_some(first)
3841            }
3842            _ => None,
3843        }
3844    }
3845
3846    fn check_updated_records(&mut self, value: Value, span: Span) -> Result<Value> {
3847        match value {
3848            Value::Record(ref r) => {
3849                let Some(id) = self
3850                    .prog
3851                    .records
3852                    .iter()
3853                    .position(|t| Some(t.name.as_str()) == r.ty.as_deref())
3854                else {
3855                    return Ok(value);
3856                };
3857                let ty = TypeSpec::Record(id as u32);
3858                let value = self.coerce(value, &ty, span)?;
3859                if !self.conforms(&value, &ty) {
3860                    return Err(RuntimeError::new(
3861                        span,
3862                        format!("updated record doesn't conform to {}", ty.describe(self.prog)),
3863                    ));
3864                }
3865                Ok(value)
3866            }
3867            Value::Dist(d) => {
3868                let mut outcomes = Vec::with_capacity(d.outcomes.len());
3869                for (v, p) in &d.outcomes {
3870                    outcomes.push((self.check_updated_records(v.clone(), span)?, *p));
3871                }
3872                ops::combine(outcomes, d.missing, &mut self.budget).map_err(|e| e.at(span))
3873            }
3874            v => Ok(v),
3875        }
3876    }
3877
3878    fn eval(&mut self, f: FnId, e: &Expr, w: &World) -> Result<Value> {
3879        let span = e.span;
3880        let at = |err: OpError| err.at(span);
3881        match &e.kind {
3882            ExprKind::Lit(l) => self.literal(l).map_err(at),
3883            ExprKind::Slot(s) => self.slot(f, *s, w, span),
3884            ExprKind::Unary(op, x) => {
3885                let v = self.eval(f, x, w)?;
3886                ops::unary(*op, &v, &mut self.budget).map_err(at)
3887            }
3888            ExprKind::Binary(op @ (BinOp::And | BinOp::Or), a, b) => {
3889                let and = *op == BinOp::And;
3890                let word = if and { "and" } else { "or" };
3891                let va = self.eval(f, a, w)?;
3892                let ta = ops::truth(&va, word).map_err(|err| err.at(a.span))?;
3893                if let Truth::Fact(x) = ta {
3894                    if x != and {
3895                        // `false and …` is false; `true or …` is true.
3896                        return Ok(Value::Bool(x));
3897                    }
3898                }
3899                let vb = self.eval(f, b, w)?;
3900                let tb = ops::truth(&vb, word).map_err(|err| err.at(b.span))?;
3901                ops::logic(and, ta, tb, &mut self.budget).map_err(at)
3902            }
3903            ExprKind::Binary(op, a, b) => {
3904                let va = self.eval(f, a, w)?;
3905                let vb = self.eval(f, b, w)?;
3906                ops::binary(*op, &va, &vb, &mut self.budget).map_err(at)
3907            }
3908            ExprKind::List(items) => {
3909                self.budget.collection(items.len() as u128).map_err(at)?;
3910                let mut values = Vec::with_capacity(items.len());
3911                for item in items {
3912                    values.push(self.eval(f, item, w)?);
3913                }
3914                Ok(Value::list(values))
3915            }
3916            ExprKind::Map(entries) => {
3917                let mut map = BTreeMap::new();
3918                for (k, v) in entries {
3919                    let key = self.eval(f, k, w)?;
3920                    if analytic::contains(&key) {
3921                        return Err(analytic::unsupported("analytic map keys").at(k.span));
3922                    }
3923                    if key.is_uncertain() {
3924                        return Err(RuntimeError::new(k.span, "map keys can't be distributions"));
3925                    }
3926                    let value = self.eval(f, v, w)?;
3927                    map.insert(key, value);
3928                }
3929                Ok(Value::map(map))
3930            }
3931            ExprKind::Record { ty, fields } => {
3932                let mut values = Vec::with_capacity(fields.len());
3933                for (name, v) in fields {
3934                    let field_ty = ty
3935                        .and_then(|t| self.prog.records[t as usize].fields.iter().find(|d| d.name == *name))
3936                        .map(|d| d.ty.clone());
3937                    let value = match field_ty {
3938                        Some(t) => {
3939                            let value = self.eval(f, v, w)?;
3940                            let value = self.coerce(value, &t, v.span)?;
3941                            if !self.conforms(&value, &t) {
3942                                let found = self.described(&value);
3943                                return Err(RuntimeError::new(
3944                                    v.span,
3945                                    format!("field `{name}` should be {}, found {found}", t.describe(self.prog)),
3946                                ));
3947                            }
3948                            value
3949                        }
3950                        None => self.eval(f, v, w)?,
3951                    };
3952                    values.push((Arc::from(name.as_str()), value));
3953                }
3954                Ok(ops::make_record(
3955                    ty.map(|t| self.record_names[t as usize].clone()),
3956                    values,
3957                ))
3958            }
3959            ExprKind::Field(x, name) => {
3960                let v = self.eval(f, x, w)?;
3961                ops::field(&v, name, &mut self.budget).map_err(at)
3962            }
3963            ExprKind::Index(x, i) => {
3964                let v = self.eval(f, x, w)?;
3965                let i = self.eval(f, i, w)?;
3966                ops::index(&v, &i, &mut self.budget).map_err(at)
3967            }
3968            ExprKind::With(x, fields) => {
3969                let base = self.eval(f, x, w)?;
3970                let mut updates = Vec::with_capacity(fields.len());
3971                for (name, v) in fields {
3972                    let field_ty = self.record_field_type(&base, name);
3973                    let value = match field_ty {
3974                        Some(ty) => {
3975                            let value = self.eval_expected(f, v, w, &ty)?;
3976                            if !self.conforms(&value, &ty) {
3977                                let found = self.described(&value);
3978                                return Err(RuntimeError::new(
3979                                    v.span,
3980                                    format!("field `{name}` should be {}, found {found}", ty.describe(self.prog)),
3981                                ));
3982                            }
3983                            value
3984                        }
3985                        None => self.eval(f, v, w)?,
3986                    };
3987                    updates.push((Arc::from(name.as_str()), value));
3988                }
3989                let value = ops::lift1(&base, &mut self.budget, |b, _| ops::with_fields(b, &updates)).map_err(at)?;
3990                self.check_updated_records(value, span)
3991            }
3992            ExprKind::Builtin { func, args, named } => self.builtin(f, *func, args, named, w, span),
3993            ExprKind::Closure { func, capture_args } => {
3994                // A named function reads its globals when it's called, so a
3995                // reference to it doesn't need them yet.
3996                let named = self.prog.functions[*func as usize].kind == FnKind::Named;
3997                let mut captured = Vec::with_capacity(capture_args.len());
3998                self.capture(&mut captured, f, *func, capture_args, w, span, named)?;
3999                Ok(Value::Closure(Arc::new(Closure {
4000                    func: *func,
4001                    captured,
4002                    named,
4003                })))
4004            }
4005            ExprKind::Bind { value, globals } => {
4006                let v = self.eval(f, value, w)?;
4007                let Value::Closure(c) = &v else {
4008                    return Ok(v);
4009                };
4010                if !c.named {
4011                    return Ok(v);
4012                }
4013                // A named function reads its globals now, as a call by name
4014                // does.
4015                let prog = self.prog;
4016                let captures = &prog.functions[c.func as usize].captures;
4017                let mut captured = Vec::with_capacity(captures.len());
4018                for capture in captures {
4019                    let CaptureSource::Global(g) = capture.source else {
4020                        return Err(RuntimeError::new(
4021                            span,
4022                            "internal error: a named function captures a local",
4023                        ));
4024                    };
4025                    let Ok(i) = globals.binary_search_by_key(&g, |&(g, _)| g) else {
4026                        return Err(RuntimeError::new(
4027                            span,
4028                            "internal error: a function value's globals aren't available here",
4029                        ));
4030                    };
4031                    captured.push(self.slot(f, globals[i].1, w, span)?);
4032                }
4033                Ok(Value::Closure(Arc::new(Closure {
4034                    func: c.func,
4035                    captured,
4036                    named: true,
4037                })))
4038            }
4039            ExprKind::Simulate { func, capture_args } => {
4040                let mut key = Vec::with_capacity(capture_args.len());
4041                self.capture(&mut key, f, *func, capture_args, w, span, false)?;
4042                self.simulate(*func, key, span)
4043            }
4044            ExprKind::Interp(parts) => {
4045                let mut text = String::new();
4046                for part in parts {
4047                    match part {
4048                        InterpPart::Lit(s) => crate::text::push(&mut text, s, &mut self.budget).map_err(at)?,
4049                        InterpPart::Expr(e) => {
4050                            let value = self.eval(f, e, w)?;
4051                            crate::text::push_value(&mut text, &value, &mut self.budget).map_err(at)?;
4052                        }
4053                    }
4054                }
4055                Ok(Value::str(&text))
4056            }
4057            ExprKind::Input(i) => Ok(self.inputs[*i as usize].clone()),
4058        }
4059    }
4060
4061    fn literal(&mut self, l: &Lit) -> OpResult<Value> {
4062        Ok(match l {
4063            Lit::Builtin(b) => Value::Builtin(*b),
4064            Lit::Unit => Value::Unit,
4065            Lit::Bool(b) => Value::Bool(*b),
4066            Lit::Int(i) => {
4067                self.budget.integer_bits(i.bits())?;
4068                self.budget.work(i.bits().div_ceil(64).max(1))?;
4069                Value::Int(i.clone())
4070            }
4071            Lit::Float(x) | Lit::FloatConstant(x) => Value::Float(*x),
4072            Lit::Prob(p) => Value::Prob(*p),
4073            Lit::Str(s) => crate::text::value(s, &mut self.budget)?,
4074            Lit::Dice { count, sides } => {
4075                if let Some(d) = self.dice.get(&(*count, *sides)) {
4076                    return Ok(d.clone());
4077                }
4078                let d = Dist::dice(*count, *sides, &mut self.budget)?.into_value();
4079                self.dice.insert((*count, *sides), d.clone());
4080                d
4081            }
4082            Lit::Enum { ty, variant } => self.enum_values[*ty as usize][*variant as usize].clone(),
4083        })
4084    }
4085
4086    fn builtin(
4087        &mut self,
4088        f: FnId,
4089        b: Builtin,
4090        args: &[Expr],
4091        named: &[(String, Expr)],
4092        w: &World,
4093        span: Span,
4094    ) -> Result<Value> {
4095        if b == Builtin::Typeof {
4096            // A direct inspection knows a delayed variable's outcome type
4097            // without drawing it and losing subsequent conjugate updates.
4098            let v = match args[0].kind {
4099                ExprKind::Slot(s) if matches!(w.slots[s as usize], Value::Delayed(_)) => w.slots[s as usize].clone(),
4100                _ => self.eval(f, &args[0], w)?,
4101            };
4102            return crate::type_name::of(&v, &self.prog.enums, &mut self.budget).map_err(|e| e.at(span));
4103        }
4104        let mut values = Vec::with_capacity(args.len());
4105        for a in args {
4106            values.push(self.eval(f, a, w)?);
4107        }
4108        let mut names = Vec::with_capacity(named.len());
4109        for (name, value) in named {
4110            names.push(name.clone());
4111            values.push(self.eval(f, value, w)?);
4112        }
4113        self.builtin_values(b, &values, &names, w.weight, span)
4114    }
4115
4116    fn check_function_arity(&self, value: &Value, given: usize, span: Span) -> Result<()> {
4117        match value {
4118            Value::Closure(c) => self.check_arity(c, given, span),
4119            Value::Builtin(b) => {
4120                let (min, max) = b.arity();
4121                if given < min || given > max || (*b == Builtin::Date && given == 2) {
4122                    return Err(RuntimeError::new(
4123                        span,
4124                        format!(
4125                            "`{}` takes {}, but {given} were given",
4126                            b.name(),
4127                            if *b == Builtin::Date {
4128                                "one ISO string or three integers".into()
4129                            } else if min == max {
4130                                format!("{min} argument(s)")
4131                            } else if max == usize::MAX {
4132                                format!("at least {min} arguments")
4133                            } else {
4134                                format!("{min} to {max} arguments")
4135                            }
4136                        ),
4137                    ));
4138                }
4139                Ok(())
4140            }
4141            _ => Err(RuntimeError::new(
4142                span,
4143                "expected a comparator function, like `(a, b) -> a - b`",
4144            )),
4145        }
4146    }
4147
4148    /// Shared dispatch for direct calls and first-class builtin values.
4149    fn builtin_values(
4150        &mut self,
4151        b: Builtin,
4152        values: &[Value],
4153        named: &[String],
4154        weight: Weight,
4155        span: Span,
4156    ) -> Result<Value> {
4157        let extrema = matches!(b, Builtin::Minimum | Builtin::Maximum);
4158        if !named.is_empty() && (!extrema || named != ["default"]) {
4159            return Err(RuntimeError::new(
4160                span,
4161                format!(
4162                    "`{}` does not accept these named arguments; only minimum/maximum support `default`",
4163                    b.name()
4164                ),
4165            ));
4166        }
4167        let positional = values.len() - named.len();
4168        if !named.is_empty() && positional == 0 {
4169            return Err(RuntimeError::new(
4170                span,
4171                "a collection argument is required before `default`",
4172            ));
4173        }
4174        if !named.is_empty() && positional >= 3 {
4175            return Err(RuntimeError::new(
4176                span,
4177                "the default was supplied both positionally and by name",
4178            ));
4179        }
4180        self.check_function_arity(&Value::Builtin(b), values.len(), span)?;
4181        let (values, default) = if extrema && (!named.is_empty() || values.len() == 3) {
4182            (&values[..values.len() - 1], values.last())
4183        } else {
4184            (values, None)
4185        };
4186        let at = |err: OpError| err.at(span);
4187        builtins::check_query_input(b, values).map_err(at)?;
4188        // Asking the built-in first spares a search of every collection
4189        // argument, on every call, for those that move values unseen.
4190        if !handles_analytic(b, values) && values.iter().any(analytic::contains) {
4191            // Internal built-ins stand for the syntax that uses them.
4192            let name = match b {
4193                Builtin::ScoreLaw => "score",
4194                Builtin::RepeatCount => "repeat",
4195                Builtin::Typeof => "typeof",
4196                b => b.name(),
4197            };
4198            return Err(analytic::unsupported(&format!("`{name}` on this outcome")).at(span));
4199        }
4200        match b {
4201            Builtin::Typeof => crate::type_name::of(&values[0], &self.prog.enums, &mut self.budget).map_err(at),
4202            Builtin::RunDate => self.config.today.map(Value::Date).ok_or_else(|| {
4203                RuntimeError::new(span, "the host did not supply an execution date for `today`")
4204                    .with_help("set Options.today, or supply today in the WASM request")
4205            }),
4206            Builtin::RunEnv => self.env(span),
4207            Builtin::Print => {
4208                let mut text = String::new();
4209                for (i, value) in values.iter().enumerate() {
4210                    if i != 0 {
4211                        crate::text::push(&mut text, " ", &mut self.budget).map_err(at)?;
4212                    }
4213                    crate::text::push_value(&mut text, value, &mut self.budget).map_err(at)?;
4214                }
4215                let line = if weight == Weight::ONE {
4216                    text
4217                } else {
4218                    format!("[{}] {text}", fmt_weight(weight))
4219                };
4220                self.printed += line.len() + 1;
4221                if self.printed > self.config.max_output {
4222                    return Err(too_much_output(span));
4223                }
4224                match &mut self.lines {
4225                    Some(lines) => lines.push((span, line)),
4226                    None => (self.print)(&line),
4227                }
4228                Ok(Value::Unit)
4229            }
4230            Builtin::Map | Builtin::Filter | Builtin::Reduce => self.higher_order(b, values, span),
4231            Builtin::Count if values.len() == 2 => self.higher_order(b, values, span),
4232            Builtin::Minimum | Builtin::Maximum => {
4233                if values.len() == 2 {
4234                    self.collection_extreme(b, values, default, span)
4235                } else {
4236                    builtins::population_extreme(&values[0], b == Builtin::Maximum, default, &mut self.budget)
4237                        .map_err(at)
4238                }
4239            }
4240            Builtin::Sort | Builtin::SortDesc if values.len() == 2 => self.higher_order(b, values, span),
4241            Builtin::Highest | Builtin::Lowest if values.len() == 3 => self.higher_order(b, values, span),
4242            Builtin::Min | Builtin::Max => builtins::call_plain(b, values, &mut self.budget).map_err(at),
4243            Builtin::Roll => self.roll(values).map_err(at),
4244            Builtin::Take => Err(RuntimeError::new(
4245                span,
4246                "use `deck.take()` to draw and remove an item from a mutable bag",
4247            )),
4248            _ if b.lifting() == Lifting::Raw => builtins::call_raw(b, values, &mut self.budget).map_err(at),
4249            _ if values.iter().any(|v| matches!(v, Value::Continuous(_))) => {
4250                let kind = values
4251                    .iter()
4252                    .find(|v| matches!(v, Value::Continuous(_)))
4253                    .unwrap()
4254                    .kind();
4255                Err(
4256                    RuntimeError::new(span, format!("`{}` needs a value, not a {kind}", b.name()))
4257                        .with_help("draw a value first, like `let x ~ normal(0, 1)`"),
4258                )
4259            }
4260            _ => ops::lift_n(values, &mut self.budget, &|a, budget| {
4261                builtins::call_plain(b, a, budget)
4262            })
4263            .map_err(at),
4264        }
4265    }
4266
4267    fn roll(&mut self, values: &[Value]) -> OpResult<Value> {
4268        if values[0].is_uncertain() {
4269            return Err(OpError::new("the number of dice to roll must be a plain number")
4270                .help("draw it first, like `let n ~ d4`"));
4271        }
4272        let count = ops::integer(&values[0], "roll count", &mut self.budget)?;
4273        let count = count
4274            .to_u64()
4275            .filter(|n| *n <= 1000)
4276            .ok_or_else(|| OpError::new("roll needs between 0 and 1000 dice"))? as u32;
4277        let die = match &values[1] {
4278            v @ (Value::Int(_) | Value::Float(_)) => {
4279                let sides = ops::integer(v, "roll sides", &mut self.budget)?;
4280                let sides = sides
4281                    .to_u64()
4282                    .filter(|n| *n >= 1 && *n <= u32::MAX as u64)
4283                    .ok_or_else(|| OpError::new("roll needs between 1 and 4294967295 sides"))?;
4284                Dist::dice(1, sides as u32, &mut self.budget)?
4285            }
4286            Value::Dist(d) => (**d).clone(),
4287            v => {
4288                return Err(OpError::new(format!(
4289                    "roll needs a die, like `d6`, found {}",
4290                    ops::article(&v.kind())
4291                )));
4292            }
4293        };
4294        let key = (count, die.clone().into_value());
4295        if let Some(pool) = self.pools.get(&key) {
4296            return Ok(pool.clone());
4297        }
4298        let pool = Dist::pool(count, &die, &mut self.budget)?.into_value();
4299        self.pools.insert(key, pool.clone());
4300        Ok(pool)
4301    }
4302
4303    fn compare_callback(
4304        &mut self,
4305        comparator: &Value,
4306        a: &Value,
4307        other: &Value,
4308        b: Builtin,
4309        span: Span,
4310    ) -> Result<std::cmp::Ordering> {
4311        self.budget.work(1).map_err(|e| e.at(span))?;
4312        match self.call_pure(comparator, vec![a.clone(), other.clone()], b.name(), span)? {
4313            Value::Int(n) => Ok(n.cmp(&probl_number::Integer::ZERO)),
4314            Value::Float(x) if x.is_finite() => Ok(x.partial_cmp(&0.0).expect("finite comparator result")),
4315            other => Err(RuntimeError::new(
4316                span,
4317                format!(
4318                    "the comparator given to `{}` must return a finite int or float (negative, zero, or positive), found {}",
4319                    b.name(),
4320                    ops::article(&other.kind())
4321                ),
4322            )),
4323        }
4324    }
4325
4326    fn collection_extreme(
4327        &mut self,
4328        b: Builtin,
4329        values: &[Value],
4330        default: Option<&Value>,
4331        span: Span,
4332    ) -> Result<Value> {
4333        if matches!(values[0], Value::Dist(_) | Value::Continuous(_)) {
4334            return Err(RuntimeError::new(
4335                span,
4336                "a comparator is supported only for collection extrema, not distribution support bounds",
4337            ));
4338        }
4339        let comparator = &values[1];
4340        self.check_function_arity(comparator, 2, span)?;
4341        let items = builtins::items(&values[0], b.name(), &mut self.budget).map_err(|e| e.at(span))?;
4342        let mut iter = items.into_iter();
4343        let Some(mut best) = iter.next() else {
4344            return builtins::empty_extreme(b.name(), default).map_err(|e| e.at(span));
4345        };
4346        for x in iter {
4347            let order = self.compare_callback(comparator, &x, &best, b, span)?;
4348            if (b == Builtin::Maximum && order.is_gt()) || (b == Builtin::Minimum && order.is_lt()) {
4349                best = x;
4350            }
4351        }
4352        Ok(best)
4353    }
4354
4355    /// Collection operations which call a function, including custom sorting.
4356    fn higher_order(&mut self, b: Builtin, values: &[Value], span: Span) -> Result<Value> {
4357        if let Value::Dist(d) = &values[0] {
4358            let mut results = Vec::with_capacity(d.outcomes.len());
4359            for (coll, p) in &d.outcomes {
4360                let mut args = values.to_vec();
4361                args[0] = coll.clone();
4362                results.push((self.higher_order(b, &args, span)?, *p));
4363            }
4364            return ops::combine(results, d.missing, &mut self.budget).map_err(|e| e.at(span));
4365        }
4366        if matches!(b, Builtin::Highest | Builtin::Lowest) {
4367            if let Value::Dist(d) = &values[1] {
4368                let mut results = Vec::with_capacity(d.outcomes.len());
4369                for (count, p) in &d.outcomes {
4370                    let mut args = values.to_vec();
4371                    args[1] = count.clone();
4372                    results.push((self.higher_order(b, &args, span)?, *p));
4373                }
4374                return ops::combine(results, d.missing, &mut self.budget).map_err(|e| e.at(span));
4375            }
4376        }
4377        let items = builtins::items(&values[0], b.name(), &mut self.budget).map_err(|e| e.at(span))?;
4378        match b {
4379            Builtin::Sort | Builtin::SortDesc | Builtin::Highest | Builtin::Lowest => {
4380                let comparator = values.last().unwrap();
4381                // Validate even empty/singleton collections without invoking the callback.
4382                self.check_function_arity(comparator, 2, span)?;
4383                let count = if matches!(b, Builtin::Highest | Builtin::Lowest) {
4384                    Some(builtins::extreme_count(&values[1], &mut self.budget).map_err(|e| e.at(span))?)
4385                } else {
4386                    None
4387                };
4388                crate::ordering::reserve_sort(items.len(), &mut self.budget).map_err(|e| e.at(span))?;
4389                let mut items = items;
4390                crate::ordering::try_sort_by(&mut items, |a, other| {
4391                    let order = self.compare_callback(comparator, a, other, b, span)?;
4392                    Ok(if matches!(b, Builtin::SortDesc | Builtin::Highest) {
4393                        order.reverse()
4394                    } else {
4395                        order
4396                    })
4397                })?;
4398                if let Some(n) = count {
4399                    items.truncate(n);
4400                }
4401                Ok(Value::list(items))
4402            }
4403            Builtin::Map => {
4404                let mut out = Vec::with_capacity(items.len());
4405                for x in items {
4406                    out.push(self.call_pure(&values[1], vec![x], "map", span)?);
4407                }
4408                Ok(Value::list(out))
4409            }
4410            Builtin::Filter | Builtin::Count => {
4411                let mut kept = Vec::new();
4412                for x in items {
4413                    let test = self.call_pure(&values[1], vec![x.clone()], b.name(), span)?;
4414                    match test {
4415                        Value::Bool(true) => kept.push(x),
4416                        Value::Bool(false) => {}
4417                        other => {
4418                            return Err(RuntimeError::new(
4419                                span,
4420                                format!("the test given to `{}` must give a fact (true or false)", b.name()),
4421                            )
4422                            .with_help(format!("it gave {}", ops::article(&other.kind()))));
4423                        }
4424                    }
4425                }
4426                Ok(if b == Builtin::Count {
4427                    Value::Int((kept.len() as i64).into())
4428                } else {
4429                    Value::list(kept)
4430                })
4431            }
4432            Builtin::Reduce => {
4433                let callback = &values[1];
4434                if !matches!(callback, Value::Closure(_) | Value::Builtin(_)) {
4435                    return Err(
4436                        RuntimeError::new(span, "`reduce` needs a function as its second argument")
4437                            .with_help("write `reduce(xs, f)` or `reduce(xs, f, initial)`"),
4438                    );
4439                }
4440                self.check_function_arity(callback, 2, span)?;
4441                let mut items = items.into_iter();
4442                let mut acc = match values.get(2) {
4443                    Some(initial) => initial.clone(),
4444                    None => items.next().ok_or_else(|| {
4445                        RuntimeError::new(span, "`reduce` of an empty collection needs an initial value")
4446                            .with_help("supply an initial value: `reduce(xs, f, initial)`")
4447                    })?,
4448                };
4449                for x in items {
4450                    acc = self.call_pure(callback, vec![acc, x], "reduce", span)?;
4451                }
4452                Ok(acc)
4453            }
4454            _ => unreachable!(),
4455        }
4456    }
4457}
4458
4459/// A call as written: `f(3, true)`.
4460/// The most rewards a loop's chain follows.
4461const MAX_REWARDS: usize = 4;
4462
4463/// A loop's rewards, as whole numbers or floats by their bits.
4464type Reward = [u64; MAX_REWARDS];
4465
4466/// Rewards as a world has them: their bits, and which are floats.
4467type Raw = (Reward, u8);
4468
4469/// A loop's chain with rewards, as `solve_loop` found it: for each state,
4470/// its steps with how much the rewards grow, and what leaves the loop from
4471/// it, in one run of the body from the state with its rewards at 0.
4472struct Rewarded<'a> {
4473    rewards: &'a [usize],
4474    moves: &'a [Vec<(u32, Raw, f64)>],
4475    exits: Vec<Vec<World>>,
4476    returns: Vec<Vec<Returned>>,
4477    unresolved: Vec<Weight>,
4478    failed: Vec<Failures>,
4479    faulted: Vec<Vec<(World, RuntimeError)>>,
4480}
4481
4482fn describe_call(fun: &Function, key: &[Value]) -> String {
4483    let args: Vec<String> = key[..fun.n_params as usize].iter().map(|v| format!("{v:?}")).collect();
4484    format!("{}({})", fun.name, args.join(", "))
4485}
4486
4487/// The error when `print` at `span` goes over the output limit.
4488/// A world faults to a catch only inside a `try` of its own call, which takes
4489/// it: none can leave the call that way.
4490fn uncaught(flow: &Flow) -> Result<()> {
4491    match flow.faulted.first() {
4492        Some((_, e)) => Err(OpError::internal("a fault went past the `try` meant to catch it").at(e.span)),
4493        None => Ok(()),
4494    }
4495}
4496
4497/// Remove the worlds at the indices `failed`, which are in increasing order.
4498fn drop_failed(worlds: &mut Vec<World>, failed: &[usize]) {
4499    if failed.is_empty() {
4500        return;
4501    }
4502    let mut failed = failed.iter().copied().peekable();
4503    let mut i = 0;
4504    worlds.retain(|_| {
4505        let drop = failed.next_if_eq(&i).is_some();
4506        i += 1;
4507        !drop
4508    });
4509}
4510
4511pub fn too_much_output(span: Span) -> RuntimeError {
4512    RuntimeError::limit(span, "the program printed more than the output limit")
4513}
4514
4515/// A weight as a percentage, even when it's too small for an `f64`.
4516pub fn fmt_weight(w: Weight) -> String {
4517    let x = w.to_f64();
4518    if x == 0.0 && !w.is_zero() {
4519        let l = w.log10() + 2.0;
4520        return format!("{:.1}e{}%", libm::pow(10.0, l - l.floor()), l.floor());
4521    }
4522    fmt_prob(x)
4523}
4524
4525enum PathKey<'a> {
4526    Field(&'a str),
4527    Index(Value),
4528}
4529
4530fn update(target: &mut Value, keys: &[PathKey], v: Value, budget: &mut Budget) -> OpResult<()> {
4531    let Some((first, rest)) = keys.split_first() else {
4532        *target = v;
4533        return Ok(());
4534    };
4535    match (first, target) {
4536        (PathKey::Field(name), Value::Record(r)) => {
4537            let r = unshared(r, |r| r.fields.len(), budget)?;
4538            let slot = r
4539                .get_mut(name)
4540                .ok_or_else(|| OpError::new(format!("this record has no field `{name}`")))?;
4541            update(slot, rest, v, budget)
4542        }
4543        (PathKey::Index(i), Value::List(items)) => {
4544            let items = unshared(items, |items| items.len(), budget)?;
4545            let k = ops::as_index(i, items.len() as u128, budget)? as usize;
4546            update(&mut items[k], rest, v, budget)
4547        }
4548        (PathKey::Index(k), Value::Map(m)) => {
4549            let m = unshared(m, |m| m.len(), budget)?;
4550            if rest.is_empty() {
4551                m.insert(k.clone(), v);
4552                return Ok(());
4553            }
4554            let slot = m
4555                .get_mut(k)
4556                .ok_or_else(|| OpError::new(format!("the key {k:?} isn't in the map")))?;
4557            update(slot, rest, v, budget)
4558        }
4559        (PathKey::Field(name), other) => Err(OpError::new(format!(
4560            "can't set the field `{name}` of {}",
4561            ops::article(&other.kind())
4562        ))),
4563        (PathKey::Index(_), other) => Err(OpError::new(format!(
4564            "can't set an element of {}",
4565            ops::article(&other.kind())
4566        ))),
4567    }
4568}
4569
4570/// `Arc::make_mut`, counting a unit of work for each element of the copy
4571/// it makes when something else holds the value.
4572fn unshared<'a, T: Clone>(arc: &'a mut Arc<T>, size: impl Fn(&T) -> usize, budget: &mut Budget) -> OpResult<&'a mut T> {
4573    if Arc::get_mut(arc).is_none() {
4574        budget.work(size(arc) as u64)?;
4575    }
4576    Ok(Arc::make_mut(arc))
4577}
4578
4579/// The value at `keys` in `target`, to change in place: each level is
4580/// copied only if something else holds it.
4581fn place_mut<'v>(target: &'v mut Value, keys: &[PathKey], budget: &mut Budget) -> OpResult<&'v mut Value> {
4582    let Some((first, rest)) = keys.split_first() else {
4583        return Ok(target);
4584    };
4585    match (first, target) {
4586        (PathKey::Field(name), Value::Record(r)) => {
4587            let slot = unshared(r, |r| r.fields.len(), budget)?
4588                .get_mut(name)
4589                .ok_or_else(|| OpError::new(format!("this record has no field `{name}`")))?;
4590            place_mut(slot, rest, budget)
4591        }
4592        (PathKey::Index(i), Value::List(items)) => {
4593            let items = unshared(items, |items| items.len(), budget)?;
4594            let k = ops::as_index(i, items.len() as u128, budget)? as usize;
4595            place_mut(&mut items[k], rest, budget)
4596        }
4597        (PathKey::Index(k), Value::Map(m)) => {
4598            let slot = unshared(m, |m| m.len(), budget)?
4599                .get_mut(k)
4600                .ok_or_else(|| OpError::new(format!("the key {k:?} isn't in the map")))?;
4601            place_mut(slot, rest, budget)
4602        }
4603        (PathKey::Field(name), other) => Err(OpError::new(format!(
4604            "can't set the field `{name}` of {}",
4605            ops::article(&other.kind())
4606        ))),
4607        (PathKey::Index(_), other) => Err(OpError::new(format!(
4608            "can't set an element of {}",
4609            ops::article(&other.kind())
4610        ))),
4611    }
4612}
4613
4614/// Whether observing a value from `d` uses a density.
4615fn is_density(d: &Value) -> bool {
4616    match d {
4617        Value::Continuous(_) => true,
4618        Value::Dist(dist) => dist.outcomes.iter().any(|(x, _)| matches!(x, Value::Continuous(_))),
4619        _ => false,
4620    }
4621}
4622
4623/// Whether a value is a discrete outcome of a draw, like `floor(x)`. Out of
4624/// line: assignments run it on every analytic value they store.
4625#[cold]
4626#[inline(never)]
4627fn is_discrete(v: &Value) -> bool {
4628    matches!(v, Value::Analytic(a) if a.is_discrete())
4629}
4630
4631/// Whether a built-in deals with these analytic arguments deliberately.
4632/// The others would look inside them as if they were plain values.
4633fn handles_analytic(b: Builtin, values: &[Value]) -> bool {
4634    use Builtin as B;
4635    let plain = |v: &Value| !analytic::contains(v);
4636    let number = |v: &Value| matches!(v, Value::Analytic(_) | Value::Int(_) | Value::Float(_) | Value::Prob(_));
4637    let numbers = |v: &Value| matches!(v, Value::List(xs) if xs.iter().all(number));
4638    match b {
4639        B::Map | B::Filter | B::Reduce | B::Len | B::IterItems | B::Settled | B::BooleanLaw | B::IsListOfLen => true,
4640        // They move values without looking at them.
4641        B::Slice
4642        | B::Reverse
4643        | B::Keys
4644        | B::Values
4645        | B::Enumerate
4646        | B::Zip
4647        | B::Push
4648        | B::Pop
4649        | B::Last
4650        | B::DropLast => true,
4651        B::Count => values.len() == 1,
4652        // A key is compared with others, so it can't be an outcome; an index
4653        // is a number, which an outcome isn't.
4654        B::Get | B::Insert | B::Remove => matches!(values[0], Value::List(_)) || plain(&values[1]),
4655        // Affine in one draw, piecewise affine.
4656        B::Sum | B::Mean => numbers(&values[0]),
4657        B::Floor | B::Ceil | B::Trunc => matches!(values[0], Value::Analytic(_)),
4658        B::Sqrt | B::Exp | B::Ln | B::Exp2 | B::Log10 | B::Log2 | B::Log1p | B::Expm1 => {
4659            matches!(values[0], Value::Analytic(_))
4660        }
4661        B::Round => matches!(values[0], Value::Analytic(_)) && values[1..].iter().all(plain),
4662        B::Abs | B::Min | B::Max | B::Clamp => values
4663            .iter()
4664            .all(|v| number(v) || matches!(v, Value::Dist(_)) && plain(v)),
4665        B::Minimum | B::Maximum => values.len() == 1 && numbers(&values[0]),
4666        _ => false,
4667    }
4668}
4669
4670/// What observing `v` from `d` weighs.
4671enum Weighs {
4672    /// The probability of `v`, the probability that is missing from `d` (so
4673    /// the true value may be up to that much higher), and the probability
4674    /// of the other outcomes, which the observation rules out. When
4675    /// sampling, a continuous `d` gives its density here (section 13).
4676    Probability(f64, f64, f64),
4677    /// When enumerating, the logarithm of a continuous `d`'s density at `v`.
4678    LogDensity(f64),
4679}
4680
4681fn likelihood(d: &Value, v: &Value, sampling: bool) -> OpResult<Weighs> {
4682    let continuous = match d {
4683        Value::Continuous(f) => Some(vec![(**f, 1.0)]),
4684        Value::Dist(dist) if dist.outcomes.iter().any(|(x, _)| matches!(x, Value::Continuous(_))) => {
4685            let parts: Option<Vec<_>> = dist
4686                .outcomes
4687                .iter()
4688                .map(|(x, p)| match x {
4689                    Value::Continuous(f) => Some((**f, *p)),
4690                    _ => None,
4691                })
4692                .collect();
4693            let parts = parts.ok_or_else(|| {
4694                OpError::new("`observe … from` can't mix a density with the probabilities of single values")
4695            })?;
4696            Some(parts)
4697        }
4698        _ => None,
4699    };
4700    if let Some(parts) = continuous {
4701        if let Value::Dist(dist) = d {
4702            if dist.missing > 0.0 && !sampling {
4703                return Err(OpError::unsupported(
4704                    "observing a value from a mixture with unresolved probability isn't supported when enumerating",
4705                )
4706                .help("the unresolved part could have any density at the value; sample the model with `@mode sample(runs: 10_000)`"));
4707            }
4708        }
4709        let x = match v {
4710            Value::Bool(_) => None,
4711            _ => v.as_f64(),
4712        }
4713        .ok_or_else(|| {
4714            OpError::new(format!(
4715                "a continuous distribution can't produce {}",
4716                ops::article(&v.kind())
4717            ))
4718        })?;
4719        if !sampling {
4720            return Ok(Weighs::LogDensity(crate::stats::log_sum_exp(
4721                parts.iter().map(|(f, p)| libm::log(*p) + f.ln_pdf(x)),
4722            )));
4723        }
4724        return Ok(Weighs::Probability(
4725            parts.iter().map(|(f, p)| p * f.pdf(x)).sum(),
4726            0.0,
4727            0.0,
4728        ));
4729    }
4730    match d {
4731        // As observing compares values: `2.0` is the outcome 2.
4732        Value::Counts(law) => {
4733            let x = match v.as_complex() {
4734                Some(z) if z.im() == 0.0 => z.re(),
4735                _ => f64::NAN,
4736            };
4737            let p = if x.is_nan() { 0.0 } else { law.pmf(x) };
4738            Ok(Weighs::Probability(p, 0.0, 1.0 - p))
4739        }
4740        Value::Dist(dist) => {
4741            let (mut seen, mut other) = (0.0, 0.0);
4742            for (x, p) in &dist.outcomes {
4743                if ops::equals(x, v) {
4744                    seen += p;
4745                } else {
4746                    other += p;
4747                }
4748            }
4749            Ok(Weighs::Probability(seen, dist.missing, other))
4750        }
4751        Value::Prob(_) => Err(OpError::new("`observe … from` needs a distribution, not a probability")
4752            .help("to observe that a fact with probability p is true, write `observe true from bernoulli(p)`")),
4753        other => Ok(if ops::equals(other, v) {
4754            Weighs::Probability(1.0, 0.0, 0.0)
4755        } else {
4756            Weighs::Probability(0.0, 0.0, 1.0)
4757        }),
4758    }
4759}