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