1use 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 pub peak_worlds: usize,
39 pub world_steps: u64,
41 pub calls: u64,
42 pub memo_hits: u64,
43 pub solved_loops: u64,
45 pub chain_states: u64,
46 pub solved_calls: u64,
49 pub call_rounds: u64,
50 pub updates: Vec<Updates>,
53 pub diagrams: Option<crate::symbolic::Diagrams>,
56}
57
58#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
61pub struct Updates {
62 pub delayed: u64,
64 pub exact: u64,
66 pub drawn: u64,
68}
69
70impl Stats {
71 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#[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 pub sample_seed: Option<u64>,
109 pub conjugate: bool,
112 pub solving: bool,
114 pub max_chain_states: usize,
116 pub partial: bool,
119 pub integrate: Option<f64>,
122}
123
124const MAX_ELIMINATION: u64 = 200_000_000;
127
128pub const BATCH: u64 = 1000;
132
133#[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
147pub type Printed = Vec<(Span, String)>;
149
150#[derive(Clone, Copy, Debug)]
152pub struct SampleTotals {
153 pub weight: Weight,
154 pub squares: Weight,
156}
157
158#[derive(Clone, Debug)]
160pub struct CallResult {
161 pub outcomes: Vec<Returned>,
162 pub unresolved: Weight,
164 pub lost: Weight,
166 pub observed: bool,
168 pub pending: Vec<(usize, Weight)>,
172 pub failures: Failures,
174}
175
176impl CallResult {
177 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 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
199struct Approx {
202 result: Arc<CallResult>,
203 head: usize,
207 frame: u64,
208 round: u64,
209}
210
211fn 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 conj: &'p Conjugacy,
228 config: Config,
229 budget: Budget,
230 memo: FxHashMap<CallKey, Arc<CallResult>>,
231 simulated: FxHashMap<CallKey, Value>,
235 active: FxHashMap<CallKey, usize>,
237 approx: FxHashMap<CallKey, Approx>,
240 frames_at: Vec<u64>,
243 rounds_at: Vec<u64>,
244 last_number: u64,
245 pending: Vec<(usize, Weight)>,
247 depth: usize,
248 pub unresolved: Weight,
250 lost: Weight,
253 solvable: Vec<bool>,
256 rewards: FxHashMap<StmtId, Vec<SlotId>>,
259 too_large: FxHashSet<StmtId>,
262 pub observed: bool,
264 pub densities: bool,
267 pub last_ruling_out: Option<Span>,
269 pub density_at: Option<Span>,
271 pub reported_densities: Option<u32>,
274 local_density_at: Option<Span>,
278 returned_densities: Option<u32>,
279 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 inputs: &'p [Value],
289 printed: usize,
291 lines: Option<Printed>,
294 pub stats: Stats,
295 sampler: Option<Rng>,
298 nested: usize,
300 next_latent: u64,
301 callback: Option<&'static str>,
304 pub failures: Failures,
306 handlers: Vec<(usize, &'p [Catch])>,
309 faulted: Vec<(World, RuntimeError)>,
312 evidence: EvidenceOrder,
315}
316
317#[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
339macro_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 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 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 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 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 #[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 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 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 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 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 #[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 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(®ion),
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 #[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 #[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 #[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 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 #[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 #[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 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 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 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 #[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 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 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 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 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 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 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 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 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 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 fn merging(&self) -> bool {
1258 self.config.merging && self.sampler.is_none()
1259 }
1260
1261 fn partial(&self) -> bool {
1265 self.config.partial && self.nested == 0 && self.callback.is_none()
1266 }
1267
1268 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 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 #[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 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 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 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 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 #[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 #[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 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 #[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 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 #[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 if self.handlers.is_empty() {
1517 return self.exec_stmt_kind(f, stmt, worlds);
1518 }
1519 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 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 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 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) = ¤t else {
1640 return Err(RuntimeError::new(
1641 span,
1642 format!("`take` needs a bag, found {}", ops::article(¤t.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 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 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 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 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 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 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 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 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 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 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 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 #[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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 ¤t {
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 ¤t {
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 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 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 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 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 #[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 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 (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 (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 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 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 let sampling = self.sampler.is_some();
3033 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 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 self.frames_at[depth] = 0;
3108 }
3109 let result = result?;
3110 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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
4188const MAX_REWARDS: usize = 4;
4191
4192type Reward = [u64; MAX_REWARDS];
4194
4195type Raw = (Reward, u8);
4197
4198struct 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
4216fn 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
4226fn 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
4244pub 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
4299fn 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#[cold]
4311#[inline(never)]
4312fn is_discrete(v: &Value) -> bool {
4313 matches!(v, Value::Analytic(a) if a.is_discrete())
4314}
4315
4316fn 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 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 B::Get | B::Insert | B::Remove => matches!(values[0], Value::List(_)) || plain(&values[1]),
4340 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
4355enum Weighs {
4357 Probability(f64, f64, f64),
4362 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 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}