Skip to main content

fugue/runtime/
async_handler.rs

1#![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/docs/runtime/async_handler.md"))]
2
3use crate::core::address::Address;
4use crate::core::distribution::Distribution;
5use crate::core::model::Model;
6use crate::runtime::handler::Handler;
7use crate::runtime::trace::Trace;
8
9/// Asynchronous counterpart of [`Handler`], for effects resolved by I/O.
10///
11/// Every method has the name, arguments and meaning of the [`Handler`] method
12/// it mirrors, but is an `async fn`, so a handler can resolve a site over the
13/// network (a tool call, a remote model, an LLM) and await the answer instead
14/// of blocking its thread. [`run_async`] calls the methods in the order
15/// [`run`](crate::runtime::handler::run) calls the synchronous ones and awaits
16/// each before the model continues. The `addr` and `dist` a method receives
17/// are borrowed from the model node, which `run_async` keeps alive until the
18/// method's future completes, so both stay usable after an `.await`.
19///
20/// Any synchronous handler is an `AsyncHandler` through [`FromSync`]. A
21/// complete custom handler is in the
22/// [module documentation](crate::runtime::async_handler).
23///
24/// # `Send`
25///
26/// The methods are declared with `async fn`, so the trait does not bound their
27/// futures by `Send`, and generic code over `H: AsyncHandler` cannot require
28/// that they are (stable Rust cannot yet name that bound; return-type
29/// notation, rust-lang/rust#109417, is unstable). For a concrete handler type
30/// the compiler sees the actual futures, because auto traits leak through
31/// `async fn`: if the handler is `Send`, none of its methods holds anything
32/// else `!Send` (an `Rc`, a `RefCell` borrow) across an `.await`, and the
33/// model's value type is `Send`, then `run_async(handler, model)` is `Send`.
34/// A handler that also borrows nothing makes it `'static`, so the run can be
35/// spawned on a multi-threaded executor. The trait stays usable by handlers
36/// that are not `Send`, which run on a single-threaded executor instead.
37///
38/// Example:
39/// ```rust
40/// # use fugue::*;
41/// # use fugue::runtime::interpreters::PriorHandler;
42/// # use rand::rngs::StdRng;
43/// # use rand::SeedableRng;
44///
45/// // Use a built-in handler, through the synchronous adapter
46/// let mut rng = StdRng::seed_from_u64(42);
47/// let handler = FromSync(PriorHandler {
48///     rng: &mut rng,
49///     trace: Trace::default(),
50/// });
51/// let model = sample(addr!("x"), Normal::new(0.0, 1.0).unwrap());
52///
53/// // Any executor can drive the future; this one is tokio's.
54/// let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();
55/// let (x, trace) = rt.block_on(run_async(handler, model));
56/// assert_eq!(trace.get_f64(&addr!("x")), Some(x));
57/// ```
58#[allow(async_fn_in_trait)]
59pub trait AsyncHandler {
60    /// Handle an f64 sampling operation (continuous distributions).
61    async fn on_sample_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>) -> f64;
62
63    /// Handle a bool sampling operation (Bernoulli).
64    async fn on_sample_bool(&mut self, addr: &Address, dist: &dyn Distribution<bool>) -> bool;
65
66    /// Handle a u64 sampling operation (Poisson, Binomial).
67    async fn on_sample_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>) -> u64;
68
69    /// Handle a usize sampling operation (Categorical).
70    async fn on_sample_usize(&mut self, addr: &Address, dist: &dyn Distribution<usize>) -> usize;
71
72    /// Handle an i64 sampling operation (signed discrete distributions).
73    ///
74    /// Like [`Handler::on_sample_i64`], this defaults to a panic with the same
75    /// message,
76    /// `handler does not implement on_sample_i64 (i64 sample site at {addr})`,
77    /// so a handler only for models without [`Model::SampleI64`] nodes need
78    /// not implement it. [`FromSync`] forwards to the wrapped handler's
79    /// method.
80    async fn on_sample_i64(&mut self, addr: &Address, _dist: &dyn Distribution<i64>) -> i64 {
81        panic!(
82            "handler does not implement on_sample_i64 (i64 sample site at {})",
83            addr
84        )
85    }
86
87    /// Handle an f64 observation operation.
88    async fn on_observe_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>, value: f64);
89
90    /// Handle a bool observation operation.
91    async fn on_observe_bool(&mut self, addr: &Address, dist: &dyn Distribution<bool>, value: bool);
92
93    /// Handle a u64 observation operation.
94    async fn on_observe_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>, value: u64);
95
96    /// Handle a usize observation operation.
97    async fn on_observe_usize(
98        &mut self,
99        addr: &Address,
100        dist: &dyn Distribution<usize>,
101        value: usize,
102    );
103
104    /// Handle an i64 observation operation.
105    ///
106    /// Like [`Handler::on_observe_i64`], this defaults to a panic with the
107    /// same message,
108    /// `handler does not implement on_observe_i64 (i64 observe site at {addr})`.
109    /// [`FromSync`] forwards to the wrapped handler's method.
110    async fn on_observe_i64(&mut self, addr: &Address, _dist: &dyn Distribution<i64>, _value: i64) {
111        panic!(
112            "handler does not implement on_observe_i64 (i64 observe site at {})",
113            addr
114        )
115    }
116
117    /// Handle a factor operation.
118    ///
119    /// Called when the model reaches a `factor` statement; the handler
120    /// typically adds `logw` to the trace's `log_factors`.
121    ///
122    /// # Arguments
123    ///
124    /// * `logw` - Log-weight to add to the model's total weight
125    async fn on_factor(&mut self, logw: f64);
126
127    /// Finalize the handler and return the accumulated trace.
128    ///
129    /// [`run_async`] awaits this once, after the model has returned its
130    /// value. It is never called for a run whose future is dropped first.
131    async fn finish(self) -> Trace
132    where
133        Self: Sized;
134}
135
136/// A synchronous [`Handler`] used as an [`AsyncHandler`].
137///
138/// Each method calls the wrapped handler's method of the same name and
139/// completes without suspending, so `run_async(FromSync(h), m)` yields exactly
140/// what [`run(h, m)`](crate::runtime::handler::run) does: the same value, the
141/// same trace, and the same draws from the handler's RNG. Use it to hand the
142/// built-in interpreters ([`PriorHandler`](crate::PriorHandler),
143/// [`ReplayHandler`](crate::ReplayHandler),
144/// [`ScoreGivenTrace`](crate::ScoreGivenTrace) and the others) to code written
145/// against `AsyncHandler`.
146///
147/// Example:
148/// ```rust
149/// # use fugue::*;
150///
151/// // Score a recorded trace from inside async code
152/// let mut recorded = Trace::default();
153/// recorded.insert_choice(addr!("x"), ChoiceValue::F64(0.5), 0.0);
154/// let scorer = FromSync(ScoreGivenTrace {
155///     base: recorded,
156///     trace: Trace::default(),
157/// });
158/// let model = sample(addr!("x"), Normal::new(0.0, 1.0).unwrap());
159///
160/// let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();
161/// let (x, scored) = rt.block_on(run_async(scorer, model));
162/// assert_eq!(x, 0.5);
163/// assert_eq!(scored.log_prior, Normal::new(0.0, 1.0).unwrap().log_prob(&0.5));
164/// ```
165#[derive(Clone, Copy, Debug, Default)]
166pub struct FromSync<H>(pub H);
167
168impl<H: Handler> AsyncHandler for FromSync<H> {
169    async fn on_sample_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>) -> f64 {
170        self.0.on_sample_f64(addr, dist)
171    }
172
173    async fn on_sample_bool(&mut self, addr: &Address, dist: &dyn Distribution<bool>) -> bool {
174        self.0.on_sample_bool(addr, dist)
175    }
176
177    async fn on_sample_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>) -> u64 {
178        self.0.on_sample_u64(addr, dist)
179    }
180
181    async fn on_sample_usize(&mut self, addr: &Address, dist: &dyn Distribution<usize>) -> usize {
182        self.0.on_sample_usize(addr, dist)
183    }
184
185    async fn on_sample_i64(&mut self, addr: &Address, dist: &dyn Distribution<i64>) -> i64 {
186        self.0.on_sample_i64(addr, dist)
187    }
188
189    async fn on_observe_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>, value: f64) {
190        self.0.on_observe_f64(addr, dist, value)
191    }
192
193    async fn on_observe_bool(
194        &mut self,
195        addr: &Address,
196        dist: &dyn Distribution<bool>,
197        value: bool,
198    ) {
199        self.0.on_observe_bool(addr, dist, value)
200    }
201
202    async fn on_observe_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>, value: u64) {
203        self.0.on_observe_u64(addr, dist, value)
204    }
205
206    async fn on_observe_usize(
207        &mut self,
208        addr: &Address,
209        dist: &dyn Distribution<usize>,
210        value: usize,
211    ) {
212        self.0.on_observe_usize(addr, dist, value)
213    }
214
215    async fn on_observe_i64(&mut self, addr: &Address, dist: &dyn Distribution<i64>, value: i64) {
216        self.0.on_observe_i64(addr, dist, value)
217    }
218
219    async fn on_factor(&mut self, logw: f64) {
220        self.0.on_factor(logw)
221    }
222
223    async fn finish(self) -> Trace {
224        self.0.finish()
225    }
226}
227
228/// Execute a probabilistic model with an asynchronous handler.
229///
230/// The asynchronous counterpart of [`run`](crate::runtime::handler::run): the
231/// same walk over the same model, dispatching each effect to the matching
232/// [`AsyncHandler`] method and awaiting it before the continuation runs, then
233/// awaiting [`AsyncHandler::finish`]. It returns the model's final value and
234/// the accumulated trace.
235///
236/// - **Sequential**: a run awaits its sites one at a time, in program order,
237///   since the next site is not known until the current one has its value.
238///   Concurrency comes from running many models at once.
239/// - **Stack-safe**: like `run` (FG-19), an explicit loop rather than
240///   recursion, so the stack depth of every poll is independent of the number
241///   of sites, whether or not the handler suspends.
242/// - **Executor-agnostic**: it is a plain future; fugue depends on no async
243///   runtime.
244/// - **Cancellable**: dropping the future stops the run at the site it is
245///   waiting on; the handler is dropped without [`AsyncHandler::finish`].
246///
247/// The future is `Send` when the handler, its method futures and `A` are; see
248/// [`AsyncHandler`].
249///
250/// Example:
251/// ```rust
252/// # use fugue::*;
253/// # use fugue::runtime::interpreters::PriorHandler;
254/// # use rand::rngs::StdRng;
255/// # use rand::SeedableRng;
256///
257/// let model = || {
258///     sample(addr!("x"), Normal::new(0.0, 1.0).unwrap())
259///         .bind(|x| observe(addr!("y"), Normal::new(x, 0.1).unwrap(), 1.2).map(move |_| x))
260/// };
261///
262/// let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();
263/// let mut rng = StdRng::seed_from_u64(123);
264/// let (x, trace) = rt.block_on(run_async(
265///     FromSync(PriorHandler { rng: &mut rng, trace: Trace::default() }),
266///     model(),
267/// ));
268///
269/// // Exactly what the synchronous interpreter gives from the same seed
270/// let mut rng = StdRng::seed_from_u64(123);
271/// let (x_sync, trace_sync) = runtime::handler::run(
272///     PriorHandler { rng: &mut rng, trace: Trace::default() },
273///     model(),
274/// );
275/// assert_eq!(x, x_sync);
276/// assert_eq!(trace.total_log_weight(), trace_sync.total_log_weight());
277/// ```
278pub async fn run_async<A, H: AsyncHandler>(mut h: H, m: Model<A>) -> (A, Trace) {
279    // The same iterative trampoline as `run` (FG-19), awaiting each effect.
280    // The model is a CPS-encoded linked list of continuations, so one node is
281    // live at a time: its `addr` and `dist` are locals of the match arm,
282    // borrowed by the handler's future until that future completes, and only
283    // then does the continuation `k(value)` build the next node. Nothing
284    // recurses, so neither the stack depth of a poll nor this future's size
285    // depends on the model's length. Only `Model::Pure` terminates.
286    let mut m = m;
287    let a = loop {
288        m = match m {
289            Model::Pure(a) => break a,
290            Model::SampleF64 { addr, dist, k } => {
291                let x = h.on_sample_f64(&addr, &*dist).await;
292                k(x)
293            }
294            Model::SampleBool { addr, dist, k } => {
295                let x = h.on_sample_bool(&addr, &*dist).await;
296                k(x)
297            }
298            Model::SampleU64 { addr, dist, k } => {
299                let x = h.on_sample_u64(&addr, &*dist).await;
300                k(x)
301            }
302            Model::SampleUsize { addr, dist, k } => {
303                let x = h.on_sample_usize(&addr, &*dist).await;
304                k(x)
305            }
306            Model::SampleI64 { addr, dist, k } => {
307                let x = h.on_sample_i64(&addr, &*dist).await;
308                k(x)
309            }
310            Model::ObserveF64 {
311                addr,
312                dist,
313                value,
314                k,
315            } => {
316                h.on_observe_f64(&addr, &*dist, value).await;
317                k(())
318            }
319            Model::ObserveBool {
320                addr,
321                dist,
322                value,
323                k,
324            } => {
325                h.on_observe_bool(&addr, &*dist, value).await;
326                k(())
327            }
328            Model::ObserveU64 {
329                addr,
330                dist,
331                value,
332                k,
333            } => {
334                h.on_observe_u64(&addr, &*dist, value).await;
335                k(())
336            }
337            Model::ObserveUsize {
338                addr,
339                dist,
340                value,
341                k,
342            } => {
343                h.on_observe_usize(&addr, &*dist, value).await;
344                k(())
345            }
346            Model::ObserveI64 {
347                addr,
348                dist,
349                value,
350                k,
351            } => {
352                h.on_observe_i64(&addr, &*dist, value).await;
353                k(())
354            }
355            Model::Factor { logw, k } => {
356                h.on_factor(logw).await;
357                k(())
358            }
359        };
360    };
361    let t = h.finish().await;
362    (a, t)
363}
364
365#[cfg(test)]
366mod tests {
367    use super::*;
368    use crate::addr;
369    use crate::core::distribution::*;
370    use crate::core::model::{factor, observe, pure, sample, ModelExt};
371    use crate::runtime::handler::run;
372    use crate::runtime::interpreters::{PriorHandler, ReplayHandler, ScoreGivenTrace};
373    use rand::rngs::StdRng;
374    use rand::{RngCore, SeedableRng};
375    use std::future::Future;
376    use std::panic::{catch_unwind, AssertUnwindSafe};
377    use std::pin::Pin;
378    use std::sync::Arc;
379    use std::task::{Context, Poll, Wake, Waker};
380
381    /// A minimal executor: poll the future on this thread, park until woken.
382    /// Enough to drive `run_async` with no async runtime at all.
383    fn block_on<F: Future>(fut: F) -> F::Output {
384        struct ThreadWaker(std::thread::Thread);
385        impl Wake for ThreadWaker {
386            fn wake(self: Arc<Self>) {
387                self.0.unpark();
388            }
389        }
390        let waker = Waker::from(Arc::new(ThreadWaker(std::thread::current())));
391        let mut cx = Context::from_waker(&waker);
392        let mut fut = std::pin::pin!(fut);
393        loop {
394            match fut.as_mut().poll(&mut cx) {
395                Poll::Ready(out) => return out,
396                Poll::Pending => std::thread::park(),
397            }
398        }
399    }
400
401    /// Returns `Pending` once (after waking its task), then `Ready`: a full
402    /// suspend-and-resume round trip through the executor.
403    struct YieldNow(bool);
404
405    impl Future for YieldNow {
406        type Output = ();
407        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
408            if self.0 {
409                Poll::Ready(())
410            } else {
411                self.0 = true;
412                cx.waker().wake_by_ref();
413                Poll::Pending
414            }
415        }
416    }
417
418    /// Suspends once before delegating each effect to the inner handler, so
419    /// every site of a run goes through the executor.
420    struct Yielding<H>(H);
421
422    impl<H: AsyncHandler> AsyncHandler for Yielding<H> {
423        async fn on_sample_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>) -> f64 {
424            YieldNow(false).await;
425            self.0.on_sample_f64(addr, dist).await
426        }
427        async fn on_sample_bool(&mut self, addr: &Address, dist: &dyn Distribution<bool>) -> bool {
428            YieldNow(false).await;
429            self.0.on_sample_bool(addr, dist).await
430        }
431        async fn on_sample_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>) -> u64 {
432            YieldNow(false).await;
433            self.0.on_sample_u64(addr, dist).await
434        }
435        async fn on_sample_usize(
436            &mut self,
437            addr: &Address,
438            dist: &dyn Distribution<usize>,
439        ) -> usize {
440            YieldNow(false).await;
441            self.0.on_sample_usize(addr, dist).await
442        }
443        async fn on_sample_i64(&mut self, addr: &Address, dist: &dyn Distribution<i64>) -> i64 {
444            YieldNow(false).await;
445            self.0.on_sample_i64(addr, dist).await
446        }
447        async fn on_observe_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>, v: f64) {
448            YieldNow(false).await;
449            self.0.on_observe_f64(addr, dist, v).await
450        }
451        async fn on_observe_bool(
452            &mut self,
453            addr: &Address,
454            dist: &dyn Distribution<bool>,
455            v: bool,
456        ) {
457            YieldNow(false).await;
458            self.0.on_observe_bool(addr, dist, v).await
459        }
460        async fn on_observe_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>, v: u64) {
461            YieldNow(false).await;
462            self.0.on_observe_u64(addr, dist, v).await
463        }
464        async fn on_observe_usize(
465            &mut self,
466            addr: &Address,
467            dist: &dyn Distribution<usize>,
468            v: usize,
469        ) {
470            YieldNow(false).await;
471            self.0.on_observe_usize(addr, dist, v).await
472        }
473        async fn on_observe_i64(&mut self, addr: &Address, dist: &dyn Distribution<i64>, v: i64) {
474            YieldNow(false).await;
475            self.0.on_observe_i64(addr, dist, v).await
476        }
477        async fn on_factor(&mut self, logw: f64) {
478            YieldNow(false).await;
479            self.0.on_factor(logw).await
480        }
481        async fn finish(self) -> Trace {
482            YieldNow(false).await;
483            self.0.finish().await
484        }
485    }
486
487    type AllTypes = (f64, bool, u64, usize, i64);
488
489    /// Samples and observes each of the five value types and ends with a
490    /// factor; the observations depend on the sampled values, so the
491    /// likelihood depends on every site.
492    fn all_effects_model() -> Model<AllTypes> {
493        crate::prob! {
494            let x <- sample(addr!("x"), Normal::new(0.0, 1.0).unwrap());
495            let b <- sample(addr!("b"), Bernoulli::new(0.3).unwrap());
496            let n <- sample(addr!("n"), Poisson::new(4.0).unwrap());
497            let c <- sample(addr!("c"), Categorical::new(vec![0.2, 0.5, 0.3]).unwrap());
498            let z <- sample(addr!("z"), DiscreteUniform::new(-3, 3).unwrap());
499            observe(addr!("obs_f64"), Normal::new(x, 0.5).unwrap(), 0.4);
500            observe(addr!("obs_bool"), Bernoulli::new(if b { 0.9 } else { 0.2 }).unwrap(), true);
501            observe(addr!("obs_u64"), Poisson::new(n as f64 + 1.0).unwrap(), 3u64);
502            observe(addr!("obs_usize"), Categorical::uniform(3).unwrap(), c);
503            observe(addr!("obs_i64"), DiscreteUniform::new(z - 1, z + 1).unwrap(), z);
504            factor(-0.5 * x * x);
505            pure((x, b, n, c, z))
506        }
507    }
508
509    /// Bit-for-bit equality of two runs: the value, every choice (address,
510    /// value, logp) and the three log-weight accumulators.
511    fn assert_same_run(expected: &(AllTypes, Trace), actual: &(AllTypes, Trace)) {
512        let ((x, b, n, c, z), t) = expected;
513        let ((ax, ab, an, ac, az), at) = actual;
514        assert_eq!(x.to_bits(), ax.to_bits(), "value x");
515        assert_eq!((b, n, c, z), (ab, an, ac, az), "value");
516        assert_eq!(t.choices.len(), at.choices.len(), "number of choices");
517        for ((addr, choice), (a_addr, a_choice)) in t.choices.iter().zip(&at.choices) {
518            assert_eq!(addr, a_addr);
519            assert_eq!(choice.addr, a_choice.addr);
520            assert_eq!(choice.value, a_choice.value, "value at {addr}");
521            assert_eq!(
522                choice.logp.to_bits(),
523                a_choice.logp.to_bits(),
524                "logp at {addr}"
525            );
526        }
527        assert_eq!(t.log_prior.to_bits(), at.log_prior.to_bits(), "log_prior");
528        assert_eq!(
529            t.log_likelihood.to_bits(),
530            at.log_likelihood.to_bits(),
531            "log_likelihood"
532        );
533        assert_eq!(
534            t.log_factors.to_bits(),
535            at.log_factors.to_bits(),
536            "log_factors"
537        );
538    }
539
540    // #62: `run_async(FromSync(h), m)` is `run(h, m)`, for each fast built-in
541    // handler, on a model with every effect: the same value, the same trace
542    // and the same RNG consumption. The `Yielding` runs add a suspension at
543    // every site, which must not change anything either.
544    #[test]
545    fn from_sync_matches_run_for_builtin_handlers() {
546        // PriorHandler, from the same seed.
547        let mut rngs = [0, 1, 2].map(|_| StdRng::seed_from_u64(62));
548        let [r_sync, r_async, r_yield] = &mut rngs;
549        let prior = |rng| PriorHandler {
550            rng,
551            trace: Trace::default(),
552        };
553        let expected = run(prior(r_sync), all_effects_model());
554        let from_sync = block_on(run_async(FromSync(prior(r_async)), all_effects_model()));
555        let yielding = block_on(run_async(
556            Yielding(FromSync(prior(r_yield))),
557            all_effects_model(),
558        ));
559        assert_same_run(&expected, &from_sync);
560        assert_same_run(&expected, &yielding);
561        let next = rngs.map(|mut rng| rng.next_u64());
562        assert!(next[0] == next[1] && next[0] == next[2], "RNG consumption");
563
564        // The model really exercises every effect.
565        let prior_trace = expected.1;
566        assert_eq!(prior_trace.choices.len(), 5);
567        assert!(prior_trace.log_prior.is_finite());
568        assert!(prior_trace.log_likelihood.is_finite() && prior_trace.log_likelihood < 0.0);
569        assert!(prior_trace.log_factors <= 0.0);
570
571        // ReplayHandler, with one site missing from the base trace, so the
572        // run both replays and draws fresh from its RNG.
573        let mut base = prior_trace.clone();
574        base.choices.remove(&addr!("n"));
575        let mut rngs = [0, 1, 2].map(|_| StdRng::seed_from_u64(7));
576        let [r_sync, r_async, r_yield] = &mut rngs;
577        let replay = |rng| ReplayHandler {
578            rng,
579            base: base.clone(),
580            trace: Trace::default(),
581        };
582        let expected = run(replay(r_sync), all_effects_model());
583        let from_sync = block_on(run_async(FromSync(replay(r_async)), all_effects_model()));
584        let yielding = block_on(run_async(
585            Yielding(FromSync(replay(r_yield))),
586            all_effects_model(),
587        ));
588        assert_same_run(&expected, &from_sync);
589        assert_same_run(&expected, &yielding);
590        let next = rngs.map(|mut rng| rng.next_u64());
591        assert!(next[0] == next[1] && next[0] == next[2], "RNG consumption");
592
593        // ScoreGivenTrace, on the full prior trace.
594        let score = || ScoreGivenTrace {
595            base: prior_trace.clone(),
596            trace: Trace::default(),
597        };
598        let expected = run(score(), all_effects_model());
599        let from_sync = block_on(run_async(FromSync(score()), all_effects_model()));
600        let yielding = block_on(run_async(Yielding(FromSync(score())), all_effects_model()));
601        assert_same_run(&expected, &from_sync);
602        assert_same_run(&expected, &yielding);
603        assert_eq!(
604            expected.1.total_log_weight().to_bits(),
605            prior_trace.total_log_weight().to_bits()
606        );
607    }
608
609    // #62, mirroring `interpretation_is_stack_safe_for_deep_models` (FG-19):
610    // 100_000 sequential sample+bind sites through `run_async` on a 512 KiB
611    // thread, driven by `block_on` on that thread. Once with effects that
612    // complete without suspending (the whole run happens inside one poll), and
613    // once with a suspension at every site (100_000 round trips through the
614    // executor); neither may grow the stack with the model.
615    #[test]
616    fn run_async_is_stack_safe_for_deep_models() {
617        fn build(i: usize, n: usize, acc: f64) -> Model<f64> {
618            if i >= n {
619                pure(acc)
620            } else {
621                sample(addr!("x", i), Normal::new(0.0, 1.0).unwrap())
622                    .bind(move |x| build(i + 1, n, acc + x))
623            }
624        }
625
626        let handle = std::thread::Builder::new()
627            .stack_size(512 * 1024)
628            .spawn(|| {
629                let n = 100_000;
630                let mut rng = StdRng::seed_from_u64(2024);
631                let (sum, trace) = block_on(run_async(
632                    FromSync(PriorHandler {
633                        rng: &mut rng,
634                        trace: Trace::default(),
635                    }),
636                    build(0, n, 0.0),
637                ));
638                assert!(sum.is_finite());
639                assert_eq!(trace.choices.len(), n);
640                assert!(trace.log_prior.is_finite());
641
642                let mut rng = StdRng::seed_from_u64(2024);
643                let (sum_yielding, trace_yielding) = block_on(run_async(
644                    Yielding(FromSync(PriorHandler {
645                        rng: &mut rng,
646                        trace: Trace::default(),
647                    })),
648                    build(0, n, 0.0),
649                ));
650                assert_eq!(sum_yielding.to_bits(), sum.to_bits());
651                assert_eq!(trace_yielding.choices.len(), n);
652                assert_eq!(
653                    trace_yielding.log_prior.to_bits(),
654                    trace.log_prior.to_bits()
655                );
656            })
657            .expect("spawn thread");
658        handle
659            .join()
660            .expect("deep model async interpretation overflowed the stack");
661    }
662
663    /// Implements every required method of both handler traits, and neither
664    /// trait's i64 methods, so their defaults run. The tests below reach only
665    /// i64 sites.
666    struct NoI64;
667
668    impl Handler for NoI64 {
669        fn on_sample_f64(&mut self, _: &Address, _: &dyn Distribution<f64>) -> f64 {
670            unreachable!()
671        }
672        fn on_sample_bool(&mut self, _: &Address, _: &dyn Distribution<bool>) -> bool {
673            unreachable!()
674        }
675        fn on_sample_u64(&mut self, _: &Address, _: &dyn Distribution<u64>) -> u64 {
676            unreachable!()
677        }
678        fn on_sample_usize(&mut self, _: &Address, _: &dyn Distribution<usize>) -> usize {
679            unreachable!()
680        }
681        fn on_observe_f64(&mut self, _: &Address, _: &dyn Distribution<f64>, _: f64) {
682            unreachable!()
683        }
684        fn on_observe_bool(&mut self, _: &Address, _: &dyn Distribution<bool>, _: bool) {
685            unreachable!()
686        }
687        fn on_observe_u64(&mut self, _: &Address, _: &dyn Distribution<u64>, _: u64) {
688            unreachable!()
689        }
690        fn on_observe_usize(&mut self, _: &Address, _: &dyn Distribution<usize>, _: usize) {
691            unreachable!()
692        }
693        fn on_factor(&mut self, _: f64) {
694            unreachable!()
695        }
696        fn finish(self) -> Trace {
697            unreachable!()
698        }
699    }
700
701    impl AsyncHandler for NoI64 {
702        async fn on_sample_f64(&mut self, _: &Address, _: &dyn Distribution<f64>) -> f64 {
703            unreachable!()
704        }
705        async fn on_sample_bool(&mut self, _: &Address, _: &dyn Distribution<bool>) -> bool {
706            unreachable!()
707        }
708        async fn on_sample_u64(&mut self, _: &Address, _: &dyn Distribution<u64>) -> u64 {
709            unreachable!()
710        }
711        async fn on_sample_usize(&mut self, _: &Address, _: &dyn Distribution<usize>) -> usize {
712            unreachable!()
713        }
714        async fn on_observe_f64(&mut self, _: &Address, _: &dyn Distribution<f64>, _: f64) {
715            unreachable!()
716        }
717        async fn on_observe_bool(&mut self, _: &Address, _: &dyn Distribution<bool>, _: bool) {
718            unreachable!()
719        }
720        async fn on_observe_u64(&mut self, _: &Address, _: &dyn Distribution<u64>, _: u64) {
721            unreachable!()
722        }
723        async fn on_observe_usize(&mut self, _: &Address, _: &dyn Distribution<usize>, _: usize) {
724            unreachable!()
725        }
726        async fn on_factor(&mut self, _: f64) {
727            unreachable!()
728        }
729        async fn finish(self) -> Trace {
730            unreachable!()
731        }
732    }
733
734    fn panic_message(f: impl FnOnce()) -> String {
735        let payload = catch_unwind(AssertUnwindSafe(f)).expect_err("expected a panic");
736        match payload.downcast::<String>() {
737            Ok(message) => *message,
738            Err(payload) => payload
739                .downcast_ref::<&str>()
740                .expect("panic payload is a string")
741                .to_string(),
742        }
743    }
744
745    // #62: the i64 methods of `AsyncHandler` default to a panic with the
746    // documented message, which is `Handler`'s: a handler that does not
747    // implement them fails the same way through `run`, through `run_async`,
748    // and through `run_async` over `FromSync`.
749    #[test]
750    fn i64_defaults_panic_with_the_handler_message() {
751        let sample_i64 = || sample(addr!("z"), DiscreteUniform::new(-1, 1).unwrap());
752        let documented = "handler does not implement on_sample_i64 (i64 sample site at z)";
753        assert_eq!(
754            panic_message(|| drop(block_on(run_async(NoI64, sample_i64())))),
755            documented
756        );
757        assert_eq!(panic_message(|| drop(run(NoI64, sample_i64()))), documented);
758        assert_eq!(
759            panic_message(|| drop(block_on(run_async(FromSync(NoI64), sample_i64())))),
760            documented
761        );
762
763        let observe_i64 = || observe(addr!("w"), DiscreteUniform::new(-1, 1).unwrap(), 0);
764        let documented = "handler does not implement on_observe_i64 (i64 observe site at w)";
765        assert_eq!(
766            panic_message(|| drop(block_on(run_async(NoI64, observe_i64())))),
767            documented
768        );
769        assert_eq!(
770            panic_message(|| drop(run(NoI64, observe_i64()))),
771            documented
772        );
773        assert_eq!(
774            panic_message(|| drop(block_on(run_async(FromSync(NoI64), observe_i64())))),
775            documented
776        );
777    }
778}