Skip to main content

probl_engine/
lib.rs

1//! The Probl engine: values, distributions and the world-set interpreter.
2//! The rules it implements are in docs/semantics.md.
3
4pub mod analytic;
5pub mod builtins;
6pub mod chain;
7pub mod complex;
8pub mod conjugate;
9pub mod continuous;
10pub mod data;
11pub mod dates;
12pub mod dist;
13pub mod error;
14pub mod failure;
15pub mod interp;
16pub mod joint;
17pub mod likelihood;
18mod math;
19pub mod ops;
20mod ordering;
21pub mod quadrature;
22pub mod report;
23mod stats;
24pub mod symbolic;
25mod text;
26mod type_name;
27pub mod value;
28pub mod weight;
29pub mod world;
30
31pub use error::{ErrorKind, Fault, RuntimeError};
32pub use failure::{Failure, Failures};
33pub use interp::{Stats, Updates};
34pub use probl_sema::ir::FailureMode;
35pub use symbolic::{Diagrams, Symbolic};
36pub use weight::Weight;
37
38use continuous::Rng;
39use dist::Budget;
40use error::OpError;
41use probl_sema::Liveness;
42use probl_sema::conjugate::{Conjugacy, Variable};
43use probl_sema::ir::{Mode, Program};
44use probl_syntax::Span;
45use report::{Format, Sink};
46use std::any::Any;
47use std::collections::BTreeMap;
48use std::panic::AssertUnwindSafe;
49use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
50use std::sync::{Arc, Condvar, Mutex, PoisonError, mpsc};
51use value::Value;
52
53/// Upper bounds set by whoever runs a program. A program's `@max_worlds` and
54/// `@max_iterations` can lower them, never raise them.
55#[derive(Clone, Debug)]
56pub struct Limits {
57    /// Maximum bits per integer; cannot exceed the parser's hard ceiling.
58    pub max_integer_bits: u64,
59    /// Cumulative allowance for large integer results (including shared copies).
60    pub max_integer_bytes: u64,
61    /// Maximum UTF-8 bytes in one string.
62    pub max_string_bytes: usize,
63    /// Cumulative UTF-8 payload bytes produced by the runtime (not live memory).
64    pub max_string_alloc_bytes: u64,
65    /// Worlds one statement may produce.
66    pub max_worlds: usize,
67    /// Outcomes of one distribution, or of combining distributions.
68    pub max_outcomes: usize,
69    /// Elements of one collection built by the program.
70    pub max_collection: usize,
71    /// Units of work: world-steps plus outcomes computed.
72    pub max_work: u64,
73    /// Iterations of one loop.
74    pub max_iterations: u64,
75    /// Nested calls.
76    pub max_call_depth: usize,
77    /// Function results kept for reuse (the cache is cleared when full).
78    pub max_cached_calls: usize,
79    /// Bytes of `print` output.
80    pub max_output: usize,
81    /// Stack size of the engine's thread; it bounds how deep calls can nest,
82    /// together with `max_call_depth`.
83    pub stack_size: usize,
84    /// Threads that sample batches of runs at the same time. The output
85    /// doesn't depend on it.
86    pub max_threads: usize,
87    /// States of a loop solved as a Markov chain; a loop with more is
88    /// unrolled instead.
89    pub max_chain_states: usize,
90    /// Nodes of the decision diagrams of a program answered symbolically
91    /// (`Options::symbolic`).
92    pub max_nodes: usize,
93}
94
95impl Default for Limits {
96    fn default() -> Limits {
97        Limits {
98            max_integer_bits: probl_number::MAX_INTEGER_BITS,
99            max_integer_bytes: 256 * 1024 * 1024,
100            max_string_bytes: 16 * 1024 * 1024,
101            max_string_alloc_bytes: 256 * 1024 * 1024,
102            max_worlds: 10_000_000,
103            max_outcomes: 2_000_000,
104            max_collection: 10_000_000,
105            max_work: 20_000_000_000,
106            max_iterations: 10_000_000,
107            max_call_depth: 500,
108            max_cached_calls: 1_000_000,
109            max_output: 64 * 1024 * 1024,
110            stack_size: 64 * 1024 * 1024,
111            max_threads: std::thread::available_parallelism().map_or(1, |n| n.get()),
112            max_chain_states: 50_000,
113            max_nodes: 10_000_000,
114        }
115    }
116}
117
118/// How the engine runs. The defaults are what `probl run` uses, except that
119/// the CLI also supplies an execution-date snapshot.
120#[derive(Clone, Debug)]
121pub struct Options {
122    /// Execution-date snapshot, as days since 1970-01-01. The host captures
123    /// it once or supplies a replay date. The engine never reads a clock.
124    /// Required only when the program evaluates the `today` constant.
125    pub today: Option<i32>,
126    /// The environment snapshot that the `env` constant stands for: each
127    /// variable's text, by name. The host captures it once; the engine never
128    /// reads the process's environment, so it's empty unless the host gives
129    /// one.
130    pub env: Arc<BTreeMap<String, String>>,
131    /// Merge worlds that reach the same state (turning it off is only useful
132    /// for testing that merging doesn't change results).
133    pub merge: bool,
134    /// Reuse the result of a function called again with the same inputs.
135    pub memoize: bool,
136    /// Overrides the program's `@epsilon`.
137    pub epsilon: Option<f64>,
138    /// Also print the simplest fraction near each probability.
139    pub fractions: bool,
140    pub limits: Limits,
141    /// Set to stop a running program.
142    pub cancel: Option<Arc<AtomicBool>>,
143    /// Overrides the program's `@mode`.
144    pub mode: Option<Mode>,
145    /// The program's data, loaded with [`data::load`]. A program that reads
146    /// data can't run without it.
147    pub inputs: Option<Arc<data::Inputs>>,
148    /// When sampling, update conjugate priors exactly (docs/semantics.md,
149    /// section 14). Turning it off draws every variable from its prior, for
150    /// comparing: the estimates mean the same either way.
151    pub conjugate: bool,
152    /// When enumerating, solve loops that cycle as Markov chains
153    /// (docs/semantics.md, section 10). Turning it off unrolls them, which
154    /// is only useful for checking the engine.
155    pub solve: bool,
156    /// When sampling, told after each batch how many runs are done, and how
157    /// many there are.
158    pub progress: Option<Progress>,
159    /// Overrides the program's `@on_error`. Without either, a run is total
160    /// when it enumerates and partial when it samples (docs/semantics.md,
161    /// section 11).
162    pub on_error: Option<FailureMode>,
163    /// Turns on numerical integration, to this relative tolerance, as the
164    /// program's `@integrate` does (docs/semantics.md, section 13).
165    pub integrate: Option<f64>,
166    /// Experimental: when enumerating, answer programs with decision
167    /// diagrams instead of worlds (docs/design/symbolic-inference.md).
168    pub symbolic: Symbolic,
169}
170
171/// A function told how many runs are done, and how many there are.
172#[derive(Clone)]
173pub struct Progress(pub Arc<dyn Fn(u64, u64) + Send + Sync>);
174
175impl std::fmt::Debug for Progress {
176    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177        f.write_str("Progress")
178    }
179}
180
181impl Default for Options {
182    fn default() -> Options {
183        Options {
184            today: None,
185            env: Arc::default(),
186            merge: true,
187            memoize: true,
188            epsilon: None,
189            fractions: false,
190            limits: Limits::default(),
191            cancel: None,
192            mode: None,
193            inputs: None,
194            conjugate: true,
195            solve: true,
196            progress: None,
197            on_error: None,
198            integrate: None,
199            symbolic: Symbolic::Off,
200        }
201    }
202}
203
204#[derive(Clone, Debug)]
205pub struct Outcome {
206    /// The execution-date snapshot supplied by the host, for reproducibility.
207    pub today: Option<i32>,
208    /// The summary line and the reports, as printed by `probl run`.
209    pub output: String,
210    pub stats: Stats,
211    /// Weight left unresolved (upper bound on the weight of cut-off worlds).
212    pub unresolved: Weight,
213    /// The probability of the evidence, if the program observed anything.
214    pub evidence: Option<Weight>,
215    /// With `@integrate`, if results came from numerical integration: its
216    /// relative tolerance. Every result of the run, and the evidence, can
217    /// then be approximate to about that tolerance.
218    pub integration: Option<f64>,
219    /// What each `report` collected, in the order of `Program::reports`.
220    pub reports: Vec<Sink>,
221    /// What each report says, as numbers: what `output` prints them from.
222    pub results: Vec<report::ReportResult>,
223    /// How the reports were rendered into `output`, to render some of them
224    /// again with [`report::render`].
225    pub format: Format,
226    /// When sampling: how.
227    pub sample: Option<Sampled>,
228    /// The data the program read, if any.
229    pub data: Vec<data::SourceInfo>,
230    /// The variables whose draws could be delayed for exact updates when
231    /// sampling, and what happened to them.
232    pub updates: Vec<(Variable, Updates)>,
233    /// What a fault did to the other worlds: the program's `@on_error`, the
234    /// host's choice, or the mode's default.
235    pub on_error: FailureMode,
236    /// In partial mode: the worlds that failed. A run with failures is a
237    /// partial result, which the reports describe only in part.
238    pub failures: Failures,
239    /// The total weight of the worlds that finished (when sampling, the
240    /// runs').
241    pub finished: Weight,
242    /// Whether an observation used a density, which makes the evidence a
243    /// density too.
244    pub densities: bool,
245}
246
247/// How a program was sampled (docs/semantics.md, section 14).
248#[derive(Clone, Debug)]
249pub struct Sampled {
250    pub runs: u64,
251    pub seed: u64,
252    /// The effective sample size: (Σ w)² / Σ w².
253    pub effective: f64,
254    /// The runs' total weight, Σ w.
255    pub weight: Weight,
256    /// Σ w².
257    pub squares: Weight,
258    /// The standard error of the evidence estimate (`Outcome::evidence`),
259    /// relative to it.
260    pub evidence_se: f64,
261    /// Whether an observation used a density, which makes the evidence a
262    /// density too.
263    pub densities: bool,
264}
265
266/// Run a program. `print` receives the output of `print(…)` as it happens.
267/// Every problem, including a crash of the engine itself, comes back as an
268/// error.
269pub fn run(
270    program: &Program,
271    options: &Options,
272    print: &mut (dyn FnMut(&str) + Send),
273) -> Result<Outcome, RuntimeError> {
274    std::thread::scope(|scope| {
275        let handle = std::thread::Builder::new()
276            .name("probl-engine".into())
277            .stack_size(options.limits.stack_size)
278            .spawn_scoped(scope, || run_here(program, options, print));
279        let handle = match handle {
280            Ok(h) => h,
281            Err(e) => {
282                return Err(internal(format!("couldn't start the engine: {e}")));
283            }
284        };
285        handle
286            .join()
287            .unwrap_or_else(|panic| Err(internal(panic_detail(&*panic))))
288    })
289}
290
291/// Run a program on the calling thread, as [`run`] does on a thread of its
292/// own. The calling thread's stack must be large enough for the program's
293/// calls (`Limits::stack_size` doesn't apply), and a panic isn't caught.
294/// With one thread allowed (`Limits::max_threads`), sampling starts no
295/// thread either, so this works where threads aren't available, such as in
296/// WebAssembly.
297pub fn run_on_this_thread(
298    program: &Program,
299    options: &Options,
300    print: &mut (dyn FnMut(&str) + Send),
301) -> Result<Outcome, RuntimeError> {
302    run_here(program, options, print)
303}
304
305/// What a panic said.
306fn panic_detail(panic: &(dyn Any + Send)) -> String {
307    panic
308        .downcast_ref::<String>()
309        .cloned()
310        .or_else(|| panic.downcast_ref::<&str>().map(|s| s.to_string()))
311        .unwrap_or_else(|| "no details".to_string())
312}
313
314fn internal(detail: String) -> RuntimeError {
315    OpError::internal("internal error: the engine crashed")
316        .at(Span::default())
317        .with_note(detail)
318}
319
320fn run_here(
321    program: &Program,
322    options: &Options,
323    print: &mut (dyn FnMut(&str) + Send),
324) -> Result<Outcome, RuntimeError> {
325    if options.today.is_some_and(|d| !dates::valid(d)) {
326        return Err(RuntimeError::new(
327            Default::default(),
328            "execution date must be within 0001-01-01..9999-12-31",
329        ));
330    }
331    let settings = &program.settings;
332    let mode = options.mode.clone().unwrap_or_else(|| settings.mode.clone());
333    let sample = match mode {
334        Mode::Auto | Mode::Enumerate => None,
335        Mode::Sample { runs, seed } => Some((runs, seed)),
336        ref other => {
337            return Err(
338                OpError::unsupported(format!("`{}` mode isn't implemented yet", other.name()))
339                    .help("enumeration and `@mode sample(runs: 10_000)` are; particles and beam search come later")
340                    .at(settings.mode_span.unwrap_or_default()),
341            );
342        }
343    };
344    if let Some((runs, _)) = sample {
345        if runs == 0 || runs > u32::MAX as u64 {
346            return Err(OpError::new(format!(
347                "sample mode needs between 1 and {} runs",
348                report::thousands(u32::MAX as i64)
349            ))
350            .at(settings.mode_span.unwrap_or_default()));
351        }
352    }
353    // Enumeration stops at the first fault by default, and sampling ends
354    // only the run that failed (section 11).
355    let on_error = options.on_error.or(settings.on_error).unwrap_or(match sample {
356        Some(_) => FailureMode::Partial,
357        None => FailureMode::Total,
358    });
359    let limits = &options.limits;
360    let config = interp::Config {
361        today: options.today,
362        env: options.env.clone(),
363        epsilon: options.epsilon.unwrap_or(settings.epsilon),
364        merging: options.merge,
365        memoizing: options.memoize,
366        max_worlds: limits.max_worlds.min(settings.max_worlds),
367        max_iterations: limits.max_iterations.min(settings.max_iterations),
368        max_call_depth: limits.max_call_depth,
369        max_cached_calls: limits.max_cached_calls,
370        max_output: limits.max_output,
371        budget: Budget {
372            max_integer_bits: limits.max_integer_bits,
373            integer_bytes_left: Arc::new(AtomicU64::new(limits.max_integer_bytes)),
374            max_string_bytes: limits.max_string_bytes,
375            string_bytes_left: Arc::new(AtomicU64::new(limits.max_string_alloc_bytes)),
376            cancel: options.cancel.clone(),
377            max_outcomes: limits.max_outcomes,
378            max_collection: limits.max_collection,
379            work_left: limits.max_work,
380            shared: None,
381        },
382        cancel: options.cancel.clone(),
383        sample_seed: sample.map(|(_, seed)| seed),
384        conjugate: options.conjugate,
385        solving: options.solve,
386        max_chain_states: limits.max_chain_states,
387        partial: on_error == FailureMode::Partial,
388        integrate: options.integrate.or(settings.integrate),
389    };
390    let epsilon = config.epsilon;
391    let inputs = inputs(program, options)?;
392    if let Some((runs, seed)) = sample {
393        let live = probl_sema::analyze(program);
394        let conj = probl_sema::conjugate::analyze(program);
395        return sampled(program, &live, &conj, config, options, print, runs, seed, on_error);
396    }
397    let integrate = config.integrate;
398    if let Some(t) = integrate.filter(|t| !(*t > 0.0 && *t <= 1e-3)) {
399        return Err(RuntimeError::new(
400            Default::default(),
401            format!("integration needs a relative tolerance above 0 and at most 0.001, not {t}"),
402        ));
403    }
404    if options.symbolic != Symbolic::Off {
405        let budget = config.budget.clone();
406        match symbolic::run(program, budget, limits.max_nodes, options.cancel.clone()) {
407            Ok(answer) => {
408                let run = Enumerated {
409                    sinks: answer.sinks,
410                    stats: Stats {
411                        diagrams: Some(answer.diagrams),
412                        ..Stats::default()
413                    },
414                    finished: answer.finished,
415                    unresolved: Weight::ZERO,
416                    failures: Failures::default(),
417                    observed: answer.observed,
418                    densities: false,
419                    integrated: None,
420                    last_ruling_out: answer.last_ruling_out,
421                };
422                return enumerated(program, options, run, on_error, epsilon, integrate);
423            }
424            Err(symbolic::Fallback::Stopped(e)) => return Err(e),
425            Err(symbolic::Fallback::Unsupported(e) | symbolic::Fallback::TooLarge(e))
426                if options.symbolic == Symbolic::Only =>
427            {
428                return Err(e);
429            }
430            // Enumerate it instead.
431            Err(_) => {}
432        }
433    }
434    let live = probl_sema::analyze(program);
435    let conj = probl_sema::conjugate::analyze(program);
436    let mut engine = interp::Engine::new(program, &live, &conj, config, inputs, print);
437    let finished = engine.run_main()?;
438    let unresolved = engine.unresolved;
439    let mut failures = std::mem::take(&mut engine.failures);
440    failures.other_units = failures.mixed_densities
441        || matches!((failures.densities, engine.reported_densities), (Some(a), Some(b)) if a != b);
442    if let Some(span) = engine.density_at.filter(|_| !unresolved.is_zero()) {
443        return Err(OpError::unsupported(
444            "observing a value from a continuous distribution while some probability is unresolved isn't supported when enumerating",
445        )
446        .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)`")
447        .at(span));
448    }
449    let run = Enumerated {
450        sinks: std::mem::take(&mut engine.sinks),
451        stats: engine.stats.clone(),
452        finished,
453        unresolved,
454        failures,
455        observed: engine.observed,
456        densities: engine.densities,
457        integrated: engine.integrated,
458        last_ruling_out: engine.last_ruling_out,
459    };
460    enumerated(program, options, run, on_error, epsilon, integrate)
461}
462
463/// What enumerating a program found, to make its outcome from.
464struct Enumerated {
465    sinks: Vec<Sink>,
466    stats: Stats,
467    /// The weight of the worlds that finished.
468    finished: Weight,
469    unresolved: Weight,
470    failures: Failures,
471    /// Whether an `observe` was reached.
472    observed: bool,
473    /// Whether an observation used a density.
474    densities: bool,
475    /// If numerical integration computed the worlds' weights: its tolerance.
476    integrated: Option<f64>,
477    /// The last `observe` that ruled out a world.
478    last_ruling_out: Option<Span>,
479}
480
481/// The outcome of an enumerated run: its summary line, its reports' results
482/// and their text.
483fn enumerated(
484    program: &Program,
485    options: &Options,
486    run: Enumerated,
487    on_error: FailureMode,
488    epsilon: f64,
489    integrate: Option<f64>,
490) -> Result<Outcome, RuntimeError> {
491    let Enumerated {
492        sinks,
493        stats,
494        finished,
495        unresolved,
496        failures,
497        observed,
498        densities,
499        integrated,
500        last_ruling_out,
501    } = run;
502    if observed && finished.is_zero() && failures.is_empty() {
503        let span = last_ruling_out.unwrap_or_default();
504        return Err(if unresolved.is_zero() {
505            RuntimeError::new(
506                span,
507                "the evidence is impossible: every world was ruled out by `observe`",
508            )
509            .with_help("check the observations; there's no answer to condition on")
510        } else {
511            RuntimeError::new(span, "every world that fits the evidence was left unresolved")
512                .with_help("lower `@epsilon` so that loops run longer")
513        });
514    }
515
516    // Reach describes control flow; unresolved weight is shown separately.
517    // Failed worlds count in it when their weight can be compared with the
518    // finished worlds' (section 11).
519    let comparable = failures.comparable();
520    let total = if comparable {
521        finished + failures.weight
522    } else {
523        finished
524    };
525    let format = Format {
526        fractions: options.fractions,
527        weighted: program.main().effects.observes,
528        unresolved,
529        program_total: total,
530        run_squares: None,
531        reach_known: comparable,
532        integrate,
533    };
534    let plain = Format {
535        fractions: false,
536        ..format
537    };
538    let results = report::results(&program.reports, &sinks, format, unresolved);
539    // What a numerical integration computed, if anything: the worlds'
540    // weights, or a mean or spread without a formula.
541    let integration = integrated.or_else(|| {
542        let numeric = |g: &report::GroupResult| g.numeric.as_ref().is_some_and(|n| n.integrated);
543        results
544            .iter()
545            .flat_map(|r| &r.groups)
546            .any(|g| numeric(g) || g.fields.iter().flat_map(|f| &f.0).any(|(_, g)| numeric(g)))
547            .then_some(integrate)
548            .flatten()
549    });
550    let mut header = String::from("enumerated");
551    if let Some(tolerance) = integration {
552        header.push_str(&format!(" · integrated to {tolerance:e}"));
553    }
554    if !failures.is_empty() {
555        header.push_str(" · partial result");
556        if comparable && !total.is_zero() {
557            header.push_str(&format!(
558                " · {} failed",
559                report::pct(failures.weight.ratio(total), plain)
560            ));
561        }
562    }
563    // With failures before later evidence, only the finished worlds'
564    // contribution to the evidence is known.
565    let evidence = observed.then_some(total);
566    let of = if comparable { "" } else { " of the finished worlds" };
567    if let Some(z) = evidence {
568        let (lo, hi) = (total.to_f64(), (total + unresolved).to_f64());
569        if densities {
570            header.push_str(&format!(" · log evidence{of} {:.4}", z.ln()));
571        } else if hi - lo >= 0.00005 {
572            header.push_str(&format!(
573                " · evidence{of} {}–{}",
574                report::pct(lo, plain),
575                report::pct(hi, plain)
576            ));
577        } else if z.to_f64() < 0.0001 {
578            header.push_str(&format!(" · evidence{of} {}", scientific(z)));
579        } else {
580            header.push_str(&format!(" · evidence{of} {}", report::pct(z.to_f64(), plain)));
581        }
582    }
583    if !unresolved.is_zero() {
584        if unresolved.to_f64() <= epsilon {
585            header.push_str(&format!(" · unresolved < {epsilon:e}"));
586        } else {
587            header.push_str(&format!(" · unresolved {:.1e}", unresolved.to_f64()));
588        }
589    }
590    let body = report::render_results(&program.reports, &results, format);
591    let output = if body.is_empty() {
592        header
593    } else {
594        format!("{header}\n\n{body}")
595    };
596    Ok(Outcome {
597        today: options.today,
598        output,
599        stats,
600        unresolved,
601        evidence,
602        integration,
603        reports: sinks,
604        results,
605        format,
606        sample: None,
607        data: sources(options),
608        updates: Vec::new(),
609        on_error,
610        failures,
611        finished,
612        densities,
613    })
614}
615
616/// A weight too small for a percentage, in scientific notation: `2.9e-25`.
617fn scientific(w: Weight) -> String {
618    let l = w.log10();
619    let exponent = l.floor();
620    let mantissa = libm::pow(10.0, l - exponent);
621    // Rounding the mantissa can make it 10.0.
622    let (mantissa, exponent) = if format!("{mantissa:.2}") == "10.00" {
623        (1.0, exponent + 1.0)
624    } else {
625        (mantissa, exponent)
626    };
627    format!("{mantissa:.2}e{exponent}")
628}
629
630/// The sampled estimate of the evidence, for the summary line (section 14):
631/// a probability, or the logarithm of a density.
632fn evidence_estimate(z: Weight, relative_se: f64, densities: bool) -> String {
633    let known = relative_se.is_finite();
634    if densities {
635        let ln = z.log10() * std::f64::consts::LN_10;
636        if !known {
637            return format!("log evidence {ln:.2}");
638        }
639        let decimals = if relative_se > 0.0 {
640            (-libm::log10(relative_se).floor()).clamp(2.0, 6.0) as usize
641        } else {
642            2
643        };
644        return format!("log evidence {ln:.decimals$} ± {relative_se:.decimals$}");
645    }
646    let x = z.to_f64();
647    if x >= 1e-4 {
648        return match known {
649            true => format!("evidence {}", report::estimate(x, x * relative_se)),
650            false => format!("evidence {}", value::fmt_prob(x)),
651        };
652    }
653    if !known {
654        return format!("evidence {}", scientific(z));
655    }
656    let pct = relative_se * 100.0;
657    let decimals = if pct >= 10.0 {
658        0
659    } else if pct >= 1.0 {
660        1
661    } else {
662        2
663    };
664    format!("evidence {} (± {pct:.decimals$}%)", scientific(z))
665}
666
667/// The values of the program's inputs: loaded, and loaded for it.
668fn inputs<'a>(program: &Program, options: &'a Options) -> Result<&'a [Value], RuntimeError> {
669    let Some(first) = program.inputs.first() else {
670        return Ok(&[]);
671    };
672    match &options.inputs {
673        Some(inputs) if inputs.fit(program) => {
674            if inputs.max_string_bytes() > options.limits.max_string_bytes {
675                return Err(RuntimeError::limit(
676                    first.span,
677                    "an input string exceeds the runtime string size limit",
678                ));
679            }
680            if inputs.max_integer_bits() > options.limits.max_integer_bits {
681                return Err(RuntimeError::limit(
682                    first.span,
683                    "an input integer exceeds the runtime integer size limit",
684                ));
685            }
686            Ok(inputs.values())
687        }
688        Some(_) => Err(internal("the data was loaded for a different program".to_string())),
689        None => Err(
690            RuntimeError::new(first.span, "the program reads data, which wasn't loaded")
691                .with_help("load it with `probl_engine::data::load` before running the program"),
692        ),
693    }
694}
695
696fn sources(options: &Options) -> Vec<data::SourceInfo> {
697    options.inputs.as_ref().map_or_else(Vec::new, |i| i.sources().to_vec())
698}
699
700/// What the batches produced so far, added up in batch order.
701struct Combined {
702    sinks: Vec<Sink>,
703    totals: interp::SampleTotals,
704    observed: bool,
705    densities: bool,
706    unresolved: Weight,
707    last_ruling_out: Option<Span>,
708    stats: Stats,
709    failures: Failures,
710}
711
712impl Combined {
713    fn new(reports: usize) -> Combined {
714        Combined {
715            sinks: vec![Sink::default(); reports],
716            totals: interp::SampleTotals {
717                weight: Weight::ZERO,
718                squares: Weight::ZERO,
719            },
720            observed: false,
721            densities: false,
722            unresolved: Weight::ZERO,
723            last_ruling_out: None,
724            stats: Stats::default(),
725            failures: Failures::default(),
726        }
727    }
728
729    fn absorb(&mut self, batch: interp::Batch) {
730        for (sink, theirs) in self.sinks.iter_mut().zip(batch.sinks) {
731            sink.absorb(theirs);
732        }
733        self.totals.weight += batch.totals.weight;
734        self.totals.squares += batch.totals.squares;
735        self.observed |= batch.observed;
736        self.densities |= batch.densities;
737        self.unresolved += batch.unresolved;
738        self.last_ruling_out = batch.last_ruling_out.or(self.last_ruling_out);
739        self.stats.absorb(&batch.stats);
740        self.failures.append(batch.failures);
741    }
742}
743
744/// A batch's result as a worker sends it: its number, what it produced (or
745/// its error), and what it printed.
746type Sent = (u64, Result<interp::Batch, RuntimeError>, interp::Printed);
747
748/// Run the batches of a sampled program on up to `max_threads` threads, and
749/// combine them in batch order: what's printed, the first error, and every
750/// sum come out the same whatever the number of threads. The threads share
751/// one budget of work.
752#[allow(clippy::too_many_arguments)]
753fn run_batches(
754    program: &Program,
755    live: &Liveness,
756    conj: &Conjugacy,
757    config: &interp::Config,
758    options: &Options,
759    print: &mut (dyn FnMut(&str) + Send),
760    runs: u64,
761    seed: u64,
762) -> Result<Combined, RuntimeError> {
763    let batches = runs.div_ceil(interp::BATCH);
764    let threads = (options.limits.max_threads.max(1) as u64).min(batches);
765    let inputs = inputs(program, options)?;
766    // How far past the batch being combined the threads may go: it bounds
767    // the finished batches waiting in memory.
768    let ahead = 2 * threads;
769    let mut config = config.clone();
770    config.budget.shared = Some(Arc::new(AtomicU64::new(config.budget.work_left)));
771    config.budget.work_left = 0;
772    let config = &config;
773    if threads == 1 {
774        return run_batches_here(program, live, conj, config, inputs, options, print, runs, seed);
775    }
776
777    let next = AtomicU64::new(0);
778    // The lowest batch that failed: no batch after it needs to run.
779    let failed = AtomicU64::new(u64::MAX);
780    // Bytes printed so far, in batch order.
781    let printed = AtomicUsize::new(0);
782    // Batches combined so far, which the threads wait for when too far ahead.
783    let combined_upto = Mutex::new(0u64);
784    let caught_up = Condvar::new();
785    let stop = |error: RuntimeError| {
786        // With the lock held, so that no thread misses the wake-up.
787        let _upto = combined_upto.lock().unwrap_or_else(PoisonError::into_inner);
788        failed.store(0, Ordering::Relaxed);
789        caught_up.notify_all();
790        Err(error)
791    };
792    let (tx, rx) = mpsc::channel::<Sent>();
793    std::thread::scope(|scope| {
794        for t in 0..threads {
795            let tx = tx.clone();
796            let (next, failed, printed) = (&next, &failed, &printed);
797            let (combined_upto, caught_up) = (&combined_upto, &caught_up);
798            let worker = move || {
799                let mut ignore = |_: &str| {};
800                // The data, copied: threads sharing it would all touch its
801                // reference counts, and slow each other down.
802                let copied: Vec<Value> = inputs.iter().map(Value::unshared).collect();
803                let mut engine = interp::Engine::new(program, live, conj, config.clone(), &copied, &mut ignore);
804                loop {
805                    let index = next.fetch_add(1, Ordering::Relaxed);
806                    if index >= batches {
807                        break;
808                    }
809                    let mut upto = combined_upto.lock().unwrap_or_else(PoisonError::into_inner);
810                    while index >= *upto + ahead && index <= failed.load(Ordering::Relaxed) {
811                        upto = caught_up.wait(upto).unwrap_or_else(PoisonError::into_inner);
812                    }
813                    drop(upto);
814                    if index > failed.load(Ordering::Relaxed) {
815                        break;
816                    }
817                    let first = index * interp::BATCH;
818                    let n = (runs - first).min(interp::BATCH);
819                    let before = printed.load(Ordering::Relaxed);
820                    let (result, lines) = std::panic::catch_unwind(AssertUnwindSafe(|| {
821                        engine.run_batch(Rng::stream(seed, index), first, n, before)
822                    }))
823                    .unwrap_or_else(|panic| (Err(internal(panic_detail(&*panic))), Vec::new()));
824                    let stopped = result.is_err();
825                    if stopped {
826                        failed.fetch_min(index, Ordering::Relaxed);
827                    }
828                    if tx.send((index, result, lines)).is_err() || stopped {
829                        break;
830                    }
831                }
832            };
833            let spawned = std::thread::Builder::new()
834                .name("probl-sampler".into())
835                .stack_size(options.limits.stack_size)
836                .spawn_scoped(scope, worker);
837            if let Err(e) = spawned {
838                // The threads already started can do the work.
839                if t == 0 {
840                    return Err(internal(format!("couldn't start a sampling thread: {e}")));
841                }
842                break;
843            }
844        }
845        drop(tx);
846
847        let mut combined = Combined::new(program.reports.len());
848        let mut pending = BTreeMap::new();
849        let mut bytes = 0;
850        for index in 0..batches {
851            let (result, lines) = loop {
852                if let Some(sent) = pending.remove(&index) {
853                    break sent;
854                }
855                match rx.recv() {
856                    Ok((i, result, lines)) => {
857                        pending.insert(i, (result, lines));
858                    }
859                    Err(_) => return stop(internal("a sampling thread stopped before its batch was done".into())),
860                }
861            };
862            for (span, line) in lines {
863                bytes += line.len() + 1;
864                if bytes > options.limits.max_output {
865                    return stop(interp::too_much_output(span));
866                }
867                print(&line);
868            }
869            printed.store(bytes, Ordering::Relaxed);
870            match result {
871                Ok(batch) => combined.absorb(batch),
872                Err(e) => return stop(e),
873            }
874            if let Some(progress) = &options.progress {
875                (progress.0)(((index + 1) * interp::BATCH).min(runs), runs);
876            }
877            *combined_upto.lock().unwrap_or_else(PoisonError::into_inner) = index + 1;
878            caught_up.notify_all();
879        }
880        Ok(combined)
881    })
882}
883
884/// Run the batches one after another on the calling thread, starting no
885/// thread: the same output as on several threads.
886#[allow(clippy::too_many_arguments)]
887fn run_batches_here(
888    program: &Program,
889    live: &Liveness,
890    conj: &Conjugacy,
891    config: &interp::Config,
892    inputs: &[Value],
893    options: &Options,
894    print: &mut (dyn FnMut(&str) + Send),
895    runs: u64,
896    seed: u64,
897) -> Result<Combined, RuntimeError> {
898    let mut ignore = |_: &str| {};
899    let mut engine = interp::Engine::new(program, live, conj, config.clone(), inputs, &mut ignore);
900    let mut combined = Combined::new(program.reports.len());
901    let mut bytes = 0;
902    for index in 0..runs.div_ceil(interp::BATCH) {
903        let first = index * interp::BATCH;
904        let n = (runs - first).min(interp::BATCH);
905        let (result, lines) = engine.run_batch(Rng::stream(seed, index), first, n, bytes);
906        for (span, line) in lines {
907            bytes += line.len() + 1;
908            if bytes > options.limits.max_output {
909                return Err(interp::too_much_output(span));
910            }
911            print(&line);
912        }
913        combined.absorb(result?);
914        if let Some(progress) = &options.progress {
915            (progress.0)(first + n, runs);
916        }
917    }
918    Ok(combined)
919}
920
921/// Sample the program and print its estimates (docs/semantics.md, section 14).
922#[allow(clippy::too_many_arguments)]
923fn sampled(
924    program: &Program,
925    live: &Liveness,
926    conj: &Conjugacy,
927    config: interp::Config,
928    options: &Options,
929    print: &mut (dyn FnMut(&str) + Send),
930    runs: u64,
931    seed: u64,
932    on_error: FailureMode,
933) -> Result<Outcome, RuntimeError> {
934    let mut engine = run_batches(program, live, conj, &config, options, print, runs, seed)?;
935    let totals = engine.totals;
936    let failures = std::mem::take(&mut engine.failures);
937    if totals.weight.is_zero() && failures.is_empty() {
938        let span = engine.last_ruling_out.unwrap_or_default();
939        return Err(RuntimeError::new(span, "every run was ruled out by `observe`")
940            .with_note(format!("{} runs were tried", report::thousands(runs as i64)))
941            .with_help(
942                "the evidence may be impossible, or too unlikely for this many runs: enumerate the model, or add runs",
943            ));
944    }
945    let effective = (totals.weight * totals.weight).ratio(totals.squares);
946    let mut header = format!("sample · {} runs · seed {seed}", report::thousands(runs as i64));
947    if !failures.is_empty() {
948        let failed = match failures.runs {
949            1 => "1 run".to_string(),
950            n => format!("{} runs", report::thousands(n as i64)),
951        };
952        header.push_str(&format!(" · partial result · {failed} failed"));
953    }
954    // Failed runs count in the evidence and reach when their weight can be
955    // compared with the finished runs' (section 11).
956    let comparable = failures.comparable();
957    let (total, total_squares) = match comparable {
958        true => (totals.weight + failures.weight, totals.squares + failures.squares),
959        false => (totals.weight, totals.squares),
960    };
961    // The evidence: the average final weight, and its standard error
962    // relative to it (section 14).
963    let evidence = total.scale(1.0 / runs as f64);
964    let n = runs as f64;
965    let evidence_se = if runs > 1 {
966        let effective = (total * total).ratio(total_squares);
967        ((n / effective - 1.0).max(0.0) / (n - 1.0)).sqrt()
968    } else {
969        f64::NAN
970    };
971    if engine.observed {
972        let of = if comparable { "" } else { " of the finished runs" };
973        header.push_str(&format!(
974            " · {}{of}",
975            evidence_estimate(evidence, evidence_se, engine.densities)
976        ));
977        header.push_str(&format!(
978            " · effective sample size {}",
979            report::thousands(effective.round() as i64)
980        ));
981    }
982    let unresolved = engine.unresolved;
983    if !unresolved.is_zero() {
984        header.push_str(&format!(
985            " · unresolved {:.1e}",
986            unresolved.ratio(totals.weight + unresolved)
987        ));
988    }
989    let format = Format {
990        fractions: false,
991        unresolved: Weight::ZERO,
992        program_total: total,
993        run_squares: Some(total_squares),
994        weighted: program.main().effects.observes,
995        reach_known: comparable,
996        integrate: None,
997    };
998    // The report lines don't show unresolved weight when sampling (the
999    // summary line does), but the results still say what's incomplete.
1000    let results = report::results(&program.reports, &engine.sinks, format, unresolved);
1001    let body = report::render_results(&program.reports, &results, format);
1002    let output = if body.is_empty() {
1003        header
1004    } else {
1005        format!("{header}\n\n{body}")
1006    };
1007    Ok(Outcome {
1008        today: options.today,
1009        output,
1010        stats: engine.stats.clone(),
1011        unresolved,
1012        evidence: engine.observed.then_some(evidence),
1013        integration: None,
1014        densities: engine.densities,
1015        reports: std::mem::take(&mut engine.sinks),
1016        results,
1017        format,
1018        sample: Some(Sampled {
1019            runs,
1020            seed,
1021            effective,
1022            weight: totals.weight,
1023            squares: totals.squares,
1024            evidence_se,
1025            densities: engine.densities,
1026        }),
1027        data: sources(options),
1028        updates: conj
1029            .variables
1030            .iter()
1031            .enumerate()
1032            .map(|(i, v)| (v.clone(), engine.stats.updates.get(i).copied().unwrap_or_default()))
1033            .collect(),
1034        on_error,
1035        failures,
1036        finished: totals.weight,
1037    })
1038}