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