fugue/runtime/

interpreters.rs

#![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/docs/runtime/interpreters.md"))]

use crate::core::address::Address;
use crate::core::distribution::Distribution;
use crate::core::model::Model;
use crate::core::numerical::nan_to_neg_inf;
use crate::error::{ErrorCode, FugueError, FugueResult};
use crate::runtime::handler::{run, Handler};
use crate::runtime::trace::{Choice, ChoiceValue, Trace};

use rand::{RngCore, SeedableRng};
use std::hash::{Hash, Hasher};

/// Panic when a sample site reuses an address already recorded in this
/// execution's output trace (FG-47).
///
/// "Fast" handlers (`PriorHandler`, `ReplayHandler`, `ScoreGivenTrace`) treat a
/// duplicate address as a programming error and panic with a precise message —
/// this is the documented fast/safe split. Detection is O(log n) and needs no
/// extra state: the output trace's `choices` map already contains exactly the
/// addresses visited so far in this run, so a hit means the address was sampled
/// twice. Before this check, the second visit silently double-counted its
/// log-prior while dropping the first choice.
#[inline]
fn assert_no_duplicate_sample(trace: &Trace, addr: &Address, handler: &str) {
    if trace.choices.contains_key(addr) {
        panic!(
            "{handler}: address {addr} was sampled twice in one execution \
             (AddressConflict, ErrorCode::{:?}={}). Every sample site must have \
             a unique address.",
            ErrorCode::AddressConflict,
            ErrorCode::AddressConflict as u32
        );
    }
}

/// Build the [`FugueError`] used for a duplicate sample address, so that "safe"
/// handlers and the fallible scoring paths surface a real
/// [`ErrorCode::AddressConflict`] (FG-47) instead of panicking.
fn address_conflict_error(addr: &Address, handler: &str) -> FugueError {
    FugueError::ModelError {
        address: Some(addr.clone()),
        reason: format!("{handler}: address sampled twice in one execution"),
        code: ErrorCode::AddressConflict,
        context: crate::error::ErrorContext::new(),
    }
}

// =============================================================================
// Internal monomorphization macros (FG-54)
//
// Each handler used to hand-copy four (now five, with i64) near-identical
// `on_sample_*` / `on_observe_*` methods. These `macro_rules!` collapse that
// copy-paste: a handler lists the `(method, type, variant)` rows once and the
// macro expands the shared body for every supported value type. The behavior is
// identical to the old hand-written methods, plus the FG-47 duplicate-address
// check and the FG-48 fresh-logp fix, applied uniformly.
// =============================================================================

/// The five value types every handler supports, as `(sample_method,
/// observe_method, rust_type, ChoiceValue variant, type-name literal,
/// Option-getter, Result-getter)` rows. Passed to the per-handler macros so the
/// row list lives in exactly one place.
macro_rules! for_each_value_type {
    ($m:ident) => {
        $m! {
            (on_sample_f64,   on_observe_f64,   f64,   F64,   "f64",   get_f64,   get_f64_result),
            (on_sample_bool,  on_observe_bool,  bool,  Bool,  "bool",  get_bool,  get_bool_result),
            (on_sample_u64,   on_observe_u64,   u64,   U64,   "u64",   get_u64,   get_u64_result),
            (on_sample_usize, on_observe_usize, usize, Usize, "usize", get_usize, get_usize_result),
            (on_sample_i64,   on_observe_i64,   i64,   I64,   "i64",   get_i64,   get_i64_result),
        }
    };
}

/// Generate the five identical `on_observe_*` methods (every handler scores an
/// observation the same way: add its log-density to `log_likelihood`). A `NaN`
/// log-density (e.g. an observed `NaN`, or a distribution whose parameters
/// went non-finite mid-model) is accumulated as `-∞` (FG-N2).
macro_rules! impl_observe_methods {
    ($(($sample:ident, $observe:ident, $ty:ty, $variant:ident, $tyname:literal, $get:ident, $get_res:ident)),* $(,)?) => {
        $(
            fn $observe(&mut self, _addr: &Address, dist: &dyn Distribution<$ty>, value: $ty) {
                self.trace.log_likelihood += nan_to_neg_inf(dist.log_prob(&value));
            }
        )*
    };
}

/// `PriorHandler`: draw a fresh value, score it, record it. Panics on a
/// duplicate address (fast handler).
macro_rules! impl_prior_sample_methods {
    ($(($sample:ident, $observe:ident, $ty:ty, $variant:ident, $tyname:literal, $get:ident, $get_res:ident)),* $(,)?) => {
        $(
            fn $sample(&mut self, addr: &Address, dist: &dyn Distribution<$ty>) -> $ty {
                assert_no_duplicate_sample(&self.trace, addr, "PriorHandler");
                let x = dist.sample(self.rng);
                let lp = dist.log_prob(&x);
                self.trace.log_prior += lp;
                self.trace.choices.insert(
                    addr.clone(),
                    Choice { addr: addr.clone(), value: ChoiceValue::$variant(x), logp: lp },
                );
                x
            }
        )*
    };
}

/// `ReplayHandler`: reuse the base value if present (panic on type mismatch),
/// otherwise sample fresh; always re-score under the current distribution.
/// Panics on a duplicate address (fast handler).
macro_rules! impl_replay_sample_methods {
    ($(($sample:ident, $observe:ident, $ty:ty, $variant:ident, $tyname:literal, $get:ident, $get_res:ident)),* $(,)?) => {
        $(
            fn $sample(&mut self, addr: &Address, dist: &dyn Distribution<$ty>) -> $ty {
                assert_no_duplicate_sample(&self.trace, addr, "ReplayHandler");
                let x = if let Some(c) = self.base.choices.get(addr) {
                    match c.value {
                        ChoiceValue::$variant(v) => v,
                        _ => panic!("expected {} at {}", $tyname, addr),
                    }
                } else {
                    dist.sample(self.rng)
                };
                let lp = dist.log_prob(&x);
                self.trace.log_prior += lp;
                self.trace.choices.insert(
                    addr.clone(),
                    Choice { addr: addr.clone(), value: ChoiceValue::$variant(x), logp: lp },
                );
                x
            }
        )*
    };
}

/// `ScoreGivenTrace`: read the fixed value from the base trace (panic if
/// missing or wrong type), score it under the current distribution, and store a
/// FRESH choice carrying that newly computed logp (FG-48). Panics on a duplicate
/// address (fast handler).
macro_rules! impl_score_sample_methods {
    ($(($sample:ident, $observe:ident, $ty:ty, $variant:ident, $tyname:literal, $get:ident, $get_res:ident)),* $(,)?) => {
        $(
            fn $sample(&mut self, addr: &Address, dist: &dyn Distribution<$ty>) -> $ty {
                assert_no_duplicate_sample(&self.trace, addr, "ScoreGivenTrace");
                let c = self
                    .base
                    .choices
                    .get(addr)
                    .unwrap_or_else(|| panic!("missing value for site {} in base trace", addr));
                let x = match c.value {
                    ChoiceValue::$variant(v) => v,
                    _ => panic!("expected {} at {}", $tyname, addr),
                };
                let lp = dist.log_prob(&x);
                self.trace.log_prior += lp;
                // FG-48: store the freshly computed logp, not the stale base one.
                self.trace.choices.insert(
                    addr.clone(),
                    Choice { addr: addr.clone(), value: ChoiceValue::$variant(x), logp: lp },
                );
                x
            }
        )*
    };
}

/// `SafeReplayHandler`: like `ReplayHandler` but recovers from missing/mismatched
/// base values by sampling fresh (optionally warning). A duplicate address is a
/// programming error even in the safe handler, so it invalidates the trace with
/// `-inf` and warns rather than silently double-counting (FG-47).
macro_rules! impl_safe_replay_sample_methods {
    ($(($sample:ident, $observe:ident, $ty:ty, $variant:ident, $tyname:literal, $get:ident, $get_res:ident)),* $(,)?) => {
        $(
            fn $sample(&mut self, addr: &Address, dist: &dyn Distribution<$ty>) -> $ty {
                if self.trace.choices.contains_key(addr) {
                    if self.warn_on_mismatch {
                        eprintln!("Warning: {}", address_conflict_error(addr, "SafeReplayHandler"));
                    }
                    self.trace.log_prior += f64::NEG_INFINITY;
                    return dist.sample(self.rng);
                }
                let x = match self.base.$get(addr) {
                    Some(v) => v,
                    None => {
                        if self.warn_on_mismatch && self.base.choices.contains_key(addr) {
                            if let Some(choice) = self.base.choices.get(addr) {
                                eprintln!(
                                    "Warning: Type mismatch at {}: expected {}, found {}",
                                    addr, $tyname, choice.value.type_name()
                                );
                            }
                        }
                        dist.sample(self.rng)
                    }
                };
                let lp = dist.log_prob(&x);
                self.trace.log_prior += lp;
                self.trace.choices.insert(
                    addr.clone(),
                    Choice { addr: addr.clone(), value: ChoiceValue::$variant(x), logp: lp },
                );
                x
            }
        )*
    };
}

/// The value a non-panicking scorer hands back to the program at a site it
/// could NOT score (missing from the base trace, type-mismatched, or visited
/// twice): a **deterministic draw from the site's own prior**, seeded by the
/// address.
///
/// The score is already invalid at this point (the caller sees `-inf` or an
/// `Err`), so the value carries no statistical meaning. What matters is that
/// execution continues *inside the program's own support*. The previous
/// behaviour, `Default::default()`, did not: a missing `Bool` came back as
/// `false`, which for a grammar prior means "function node, recurse" at every
/// depth (an unbounded recursion that overflows the stack); a missing `f64`
/// scale came back as `0.0`, which `Normal::new(mu, 0.0).unwrap()` turns into a
/// panic inside the likelihood. A prior draw is, by construction, a value the
/// program can be handed.
///
/// Seeding from the address (its cached SipHash, fixed keys) keeps the scorers
/// pure functions of `(base, model)` — same input, same output — while giving
/// different sites independent draws, so a chain of missing `#leaf` decisions
/// terminates the way the prior does rather than repeating one outcome forever.
fn fallback_prior_draw<T>(addr: &Address, dist: &dyn Distribution<T>) -> T {
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    addr.hash(&mut hasher);
    let mut rng = rand::rngs::StdRng::seed_from_u64(hasher.finish());
    dist.sample(&mut rng)
}

/// `SafeScoreGivenTrace`: like `ScoreGivenTrace` but returns an invalid (`-inf`)
/// trace instead of panicking on a missing/mismatched address, and stores a
/// FRESH choice with the newly computed logp (FG-48). A duplicate address
/// invalidates the trace (FG-47). At a site it cannot score it hands the
/// program a deterministic prior draw (see [`fallback_prior_draw`]) so the
/// program keeps executing inside its own support instead of diverging on a
/// `Default::default()` value.
macro_rules! impl_safe_score_sample_methods {
    ($(($sample:ident, $observe:ident, $ty:ty, $variant:ident, $tyname:literal, $get:ident, $get_res:ident)),* $(,)?) => {
        $(
            fn $sample(&mut self, addr: &Address, dist: &dyn Distribution<$ty>) -> $ty {
                if self.trace.choices.contains_key(addr) {
                    if self.warn_on_error {
                        eprintln!("Warning: {}", address_conflict_error(addr, "SafeScoreGivenTrace"));
                    }
                    self.trace.log_prior += f64::NEG_INFINITY;
                    // Invalid already; keep the program inside its support.
                    return self
                        .base
                        .$get(addr)
                        .unwrap_or_else(|| fallback_prior_draw(addr, dist));
                }
                match self.base.$get_res(addr) {
                    Ok(x) => {
                        let lp = dist.log_prob(&x);
                        self.trace.log_prior += lp;
                        // FG-48: fresh logp consistent with what we accumulated.
                        self.trace.choices.insert(
                            addr.clone(),
                            Choice { addr: addr.clone(), value: ChoiceValue::$variant(x), logp: lp },
                        );
                        x
                    }
                    Err(e) => {
                        if self.warn_on_error {
                            eprintln!("Warning: Failed to get {} at {}: {}", $tyname, addr, e);
                        }
                        // The score is invalid (-inf sentinel). Hand the program a
                        // deterministic prior draw so it keeps executing inside its
                        // own support, and record it so the returned trace is a
                        // complete assignment a caller can inspect.
                        self.trace.log_prior += f64::NEG_INFINITY;
                        let x = fallback_prior_draw(addr, dist);
                        self.trace.choices.insert(
                            addr.clone(),
                            Choice { addr: addr.clone(), value: ChoiceValue::$variant(x), logp: dist.log_prob(&x) },
                        );
                        x
                    }
                }
            }
        )*
    };
}

/// `StrictScoreGivenTrace`: the fallible, structure-checking scoring path
/// (FG-20/FG-21). It records the FIRST structural problem into `self.error`
/// instead of panicking: an address absent from the base trace or a type
/// mismatch yields `ErrorCode::UnexpectedModelStructure`; a duplicate address
/// yields `ErrorCode::AddressConflict`. On success it stores a fresh, correctly
/// scored choice. After a failure the program is handed a deterministic prior
/// draw (see [`fallback_prior_draw`]) so it runs to completion inside its own
/// support; the driver discards the trace and returns the recorded error.
macro_rules! impl_strict_score_sample_methods {
    ($(($sample:ident, $observe:ident, $ty:ty, $variant:ident, $tyname:literal, $get:ident, $get_res:ident)),* $(,)?) => {
        $(
            fn $sample(&mut self, addr: &Address, dist: &dyn Distribution<$ty>) -> $ty {
                if self.trace.choices.contains_key(addr) {
                    if self.error.is_none() {
                        *self.error = Some(address_conflict_error(addr, "StrictScoreGivenTrace"));
                    }
                    return self
                        .base
                        .$get(addr)
                        .unwrap_or_else(|| fallback_prior_draw(addr, dist));
                }
                match self.base.$get_res(addr) {
                    Ok(x) => {
                        let lp = dist.log_prob(&x);
                        self.trace.log_prior += lp;
                        self.trace.choices.insert(
                            addr.clone(),
                            Choice { addr: addr.clone(), value: ChoiceValue::$variant(x), logp: lp },
                        );
                        x
                    }
                    Err(cause) => {
                        if self.error.is_none() {
                            *self.error = Some(FugueError::ModelError {
                                address: Some(addr.clone()),
                                reason: format!(
                                    "model visited address {} not present (as {}) in the base \
                                     trace; structure varies between the base trace and this model",
                                    addr, $tyname
                                ),
                                code: ErrorCode::UnexpectedModelStructure,
                                context: crate::error::ErrorContext::new().with_cause(cause),
                            });
                        }
                        // Keep executing inside the program's support; the trace
                        // is discarded by the driver, which returns the error.
                        let x = fallback_prior_draw(addr, dist);
                        self.trace.choices.insert(
                            addr.clone(),
                            Choice { addr: addr.clone(), value: ChoiceValue::$variant(x), logp: dist.log_prob(&x) },
                        );
                        x
                    }
                }
            }
        )*
    };
}

/// `ReconcilingScoreGivenTrace`: the reconciling scoring path (FG-20/FG-21).
/// Addresses present in the base trace are replayed and re-scored; NEW addresses
/// (absent, or present with the wrong type) are sampled fresh from the prior and
/// their log-prior accumulated (the RJMCMC-correct treatment of prior-proposed
/// fresh dimensions) and recorded in `fresh`. Vanished addresses are computed
/// after the run by the driver. A duplicate address is still an error.
macro_rules! impl_reconciling_score_sample_methods {
    ($(($sample:ident, $observe:ident, $ty:ty, $variant:ident, $tyname:literal, $get:ident, $get_res:ident)),* $(,)?) => {
        $(
            fn $sample(&mut self, addr: &Address, dist: &dyn Distribution<$ty>) -> $ty {
                if self.trace.choices.contains_key(addr) {
                    if self.error.is_none() {
                        *self.error =
                            Some(address_conflict_error(addr, "ReconcilingScoreGivenTrace"));
                    }
                    // Keep the program inside its support until the driver
                    // returns the error (same hazard as the strict scorer).
                    return self
                        .base
                        .$get(addr)
                        .unwrap_or_else(|| dist.sample(self.rng));
                }
                let x = match self.base.$get(addr) {
                    Some(v) => v,
                    None => {
                        // New (or type-changed) dimension: propose from the prior.
                        self.fresh.push(addr.clone());
                        dist.sample(self.rng)
                    }
                };
                let lp = dist.log_prob(&x);
                self.trace.log_prior += lp;
                self.trace.choices.insert(
                    addr.clone(),
                    Choice { addr: addr.clone(), value: ChoiceValue::$variant(x), logp: lp },
                );
                x
            }
        )*
    };
}

/// Handler for prior sampling - generates fresh random values from distributions.
///
/// This is the foundational interpreter that implements standard "forward sampling"
/// from probabilistic models. It draws fresh values from distributions and accumulates
/// log-probabilities in the trace.
///
/// Example:
/// ```rust
/// # use fugue::*;
/// # use fugue::runtime::interpreters::PriorHandler;
/// # use rand::rngs::StdRng;
/// # use rand::SeedableRng;
///
/// let model = sample(addr!("x"), Normal::new(0.0, 1.0).unwrap())
///     .bind(|x| observe(addr!("y"), Normal::new(x, 0.5).unwrap(), 1.2)
///         .map(move |_| x));
///
/// let mut rng = StdRng::seed_from_u64(42);
/// let (result, trace) = runtime::handler::run(
///     PriorHandler { rng: &mut rng, trace: Trace::default() },
///     model
/// );
///
/// assert!(result.is_finite());
/// assert!(trace.log_likelihood.is_finite());
/// ```
pub struct PriorHandler<'r, R: RngCore> {
    /// Random number generator for sampling.
    pub rng: &'r mut R,
    /// Trace to accumulate execution history.
    pub trace: Trace,
}
impl<'r, R: RngCore> Handler for PriorHandler<'r, R> {
    for_each_value_type!(impl_prior_sample_methods);
    for_each_value_type!(impl_observe_methods);

    fn on_factor(&mut self, logw: f64) {
        self.trace.log_factors += nan_to_neg_inf(logw);
    }

    fn finish(self) -> Trace {
        self.trace
    }
}

/// Handler for replaying models with values from an existing trace.
///
/// ReplayHandler replays a model execution using stored trace values. When a sampling
/// site is encountered: if the address exists in the base trace, use that value;
/// if missing, sample fresh. Essential for MCMC where you replay most choices
/// but sample new values at specific sites.
///
/// Example:
/// ```rust
/// # use fugue::*;
/// # use fugue::runtime::interpreters::*;
/// # use rand::rngs::StdRng;
/// # use rand::SeedableRng;
///
/// // Create base trace
/// let mut rng = StdRng::seed_from_u64(42);
/// let model_fn = || sample(addr!("x"), Normal::new(0.0, 1.0).unwrap());
/// let (original, base_trace) = runtime::handler::run(
///     PriorHandler { rng: &mut rng, trace: Trace::default() },
///     model_fn()
/// );
///
/// // Replay using base trace values
/// let (replayed, _) = runtime::handler::run(
///     ReplayHandler {
///         rng: &mut rng,
///         base: base_trace,
///         trace: Trace::default()
///     },
///     model_fn()
/// );
///
/// assert_eq!(original, replayed); // Same value replayed
/// ```
pub struct ReplayHandler<'r, R: RngCore> {
    /// Random number generator for sampling at addresses not in base trace.
    pub rng: &'r mut R,
    /// Base trace containing values to replay.
    pub base: Trace,
    /// New trace to accumulate the replay execution.
    pub trace: Trace,
}
impl<'r, R: RngCore> Handler for ReplayHandler<'r, R> {
    for_each_value_type!(impl_replay_sample_methods);
    for_each_value_type!(impl_observe_methods);

    fn on_factor(&mut self, logw: f64) {
        self.trace.log_factors += nan_to_neg_inf(logw);
    }

    fn finish(self) -> Trace {
        self.trace
    }
}

/// Handler for scoring a model given a complete trace of fixed choices.
///
/// ScoreGivenTrace computes log-probability of a model execution where all random
/// choices are predetermined. No sampling occurs - values are looked up from the
/// base trace and their log-probabilities computed. Essential for MCMC acceptance
/// ratios, importance sampling, and model comparison.
///
/// Example:
/// ```rust
/// # use fugue::*;
/// # use fugue::runtime::interpreters::*;
/// # use rand::rngs::StdRng;
/// # use rand::SeedableRng;
///
/// // Create a complete trace
/// let mut rng = StdRng::seed_from_u64(42);
/// let (_, complete_trace) = runtime::handler::run(
///     PriorHandler { rng: &mut rng, trace: Trace::default() },
///     sample(addr!("x"), Normal::new(0.0, 1.0).unwrap())
/// );
///
/// // Score under different model parameters  
/// let (value, score_trace) = runtime::handler::run(
///     ScoreGivenTrace {
///         base: complete_trace,
///         trace: Trace::default()
///     },
///     sample(addr!("x"), Normal::new(1.0, 2.0).unwrap()) // Different parameters
/// );
///
/// assert!(score_trace.total_log_weight().is_finite());
/// ```
///
/// # Incompatible traces
///
/// This is the fast path and it **panics** on an address missing from `base`
/// or present with another value type (`"missing value for site …"`). For a
/// trace of uncertain provenance use [`score_given_trace_strict`] (an `Err`
/// naming the first offending site), [`SafeScoreGivenTrace`] (a `-inf`
/// `log_prior` sentinel), or [`score_given_trace_reconciled`] (fresh prior draws
/// at the new sites, reported). All three keep executing the program past the
/// failure with a prior draw at the unscorable site, so a program whose control
/// flow depends on that site terminates inside its own support.
pub struct ScoreGivenTrace {
    /// Base trace containing the fixed choices to score.
    pub base: Trace,
    /// New trace to accumulate log-probabilities.
    pub trace: Trace,
}
impl Handler for ScoreGivenTrace {
    for_each_value_type!(impl_score_sample_methods);
    for_each_value_type!(impl_observe_methods);

    fn on_factor(&mut self, logw: f64) {
        self.trace.log_factors += nan_to_neg_inf(logw);
    }

    fn finish(self) -> Trace {
        self.trace
    }
}

/// Safe version of ReplayHandler that gracefully handles trace inconsistencies.
///
/// SafeReplayHandler replays model execution like ReplayHandler, but handles type
/// mismatches and missing addresses gracefully by logging warnings and sampling
/// fresh values instead of panicking. Essential for production systems where
/// trace consistency cannot be guaranteed.
///
/// Example:
/// ```rust
/// # use fugue::*;
/// # use fugue::runtime::interpreters::*;
/// # use rand::rngs::StdRng;
/// # use rand::SeedableRng;
///
/// // Create trace with potential inconsistencies
/// let mut rng = StdRng::seed_from_u64(42);
/// let (_, base_trace) = runtime::handler::run(
///     PriorHandler { rng: &mut rng, trace: Trace::default() },
///     sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()) // f64 value
/// );
///
/// // Safe replay handles type mismatch gracefully
/// let (result, trace) = runtime::handler::run(
///     SafeReplayHandler {
///         rng: &mut rng,
///         base: base_trace,
///         trace: Trace::default(),
///         warn_on_mismatch: true, // Enable warnings
///     },
///     sample(addr!("x"), Bernoulli::new(0.5).unwrap()) // Expects bool
/// );
///
/// assert!(trace.total_log_weight().is_finite()); // Continues execution
/// ```
pub struct SafeReplayHandler<'r, R: RngCore> {
    /// Random number generator for sampling at addresses not in base trace.
    pub rng: &'r mut R,
    /// Base trace containing values to replay.
    pub base: Trace,
    /// New trace to accumulate the replay execution.
    pub trace: Trace,
    /// Whether to log warnings on type mismatches (useful for debugging).
    pub warn_on_mismatch: bool,
}
impl<'r, R: RngCore> Handler for SafeReplayHandler<'r, R> {
    for_each_value_type!(impl_safe_replay_sample_methods);
    for_each_value_type!(impl_observe_methods);

    fn on_factor(&mut self, logw: f64) {
        self.trace.log_factors += nan_to_neg_inf(logw);
    }

    fn finish(self) -> Trace {
        self.trace
    }
}

/// Safe version of ScoreGivenTrace that gracefully handles incomplete traces.
///
/// SafeScoreGivenTrace computes log-probability like ScoreGivenTrace, but handles
/// missing addresses or type mismatches by returning negative infinity log-weight
/// instead of panicking. Use it when trace validity cannot be guaranteed.
///
/// # Structurally incompatible traces
///
/// The `-inf` `log_prior` is the *only* signal that the score is invalid; the
/// returned value `A` and the returned trace's choices must not be read as a
/// score of anything. At a site it cannot score the handler hands the program a
/// **deterministic draw from that site's prior** (seeded by the address) and
/// records it, so a program whose control flow depends on the site keeps
/// executing inside its own support: a missing `Bool` no longer arrives as
/// `false` (for a grammar prior: "recurse", at every depth, until the stack
/// overflows) and a missing scale no longer arrives as `0.0` (which panics
/// inside `Normal::new`). Same input, same output — the handler remains a pure
/// function of `(base, model)`.
///
/// Example:
/// ```rust
/// # use fugue::*;
/// # use fugue::runtime::interpreters::*;
/// # use rand::rngs::StdRng;
/// # use rand::SeedableRng;
///
/// // Create incomplete trace
/// let mut rng = StdRng::seed_from_u64(42);
/// let (_, incomplete_trace) = runtime::handler::run(
///     PriorHandler { rng: &mut rng, trace: Trace::default() },
///     sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()) // Only has "x"
/// );
///
/// // Safe scoring handles missing address gracefully
/// let (_, score_trace) = runtime::handler::run(
///     SafeScoreGivenTrace {
///         base: incomplete_trace,
///         trace: Trace::default(),
///         warn_on_error: true, // Enable warnings
///     },
///     sample(addr!("missing"), Normal::new(0.0, 1.0).unwrap()) // Address not in base
/// );
///
/// assert_eq!(score_trace.total_log_weight(), f64::NEG_INFINITY); // Graceful failure
/// ```
pub struct SafeScoreGivenTrace {
    /// Base trace containing the fixed choices to score.
    pub base: Trace,
    /// New trace to accumulate log-probabilities.
    pub trace: Trace,
    /// Whether to log warnings on missing addresses or type mismatches.
    pub warn_on_error: bool,
}
impl Handler for SafeScoreGivenTrace {
    for_each_value_type!(impl_safe_score_sample_methods);
    for_each_value_type!(impl_observe_methods);

    fn on_factor(&mut self, logw: f64) {
        self.trace.log_factors += nan_to_neg_inf(logw);
    }

    fn finish(self) -> Trace {
        self.trace
    }
}

// =============================================================================
// Structure-varying (trans-dimensional) scoring paths (FG-20 / FG-21)
// =============================================================================

/// The strict, fallible sibling of [`ScoreGivenTrace`].
///
/// This handler backs [`score_given_trace_strict`]. Instead of panicking when
/// the model's address structure differs from the base trace, it records the
/// first structural problem so the driver can return a [`FugueError`]:
///
/// - visiting an address absent from the base trace (or present with the wrong
///   value type) -> [`ErrorCode::UnexpectedModelStructure`];
/// - visiting the same address twice -> [`ErrorCode::AddressConflict`].
///
/// On success every visited site stores a fresh, correctly scored choice.
///
/// After the first failure the program still runs to completion — the handler
/// cannot abort `run` — so the value returned at the failed site matters for
/// *termination*, not for the score: it is a deterministic draw from the site's
/// prior (seeded by the address), which keeps a program whose structure depends
/// on that site inside its own support. `Default::default()`, the previous
/// choice, did not: a grammar's missing `Bool` leaf flag read as "function
/// node" at every depth and recursed until the stack overflowed, and a missing
/// `f64` scale read as `0.0` and panicked inside the likelihood.
pub struct StrictScoreGivenTrace<'e> {
    /// Base trace containing the fixed choices to score.
    pub base: Trace,
    /// New trace to accumulate log-probabilities.
    pub trace: Trace,
    /// First structural error encountered, surfaced by the driver.
    error: &'e mut Option<FugueError>,
}
impl<'e> Handler for StrictScoreGivenTrace<'e> {
    for_each_value_type!(impl_strict_score_sample_methods);
    for_each_value_type!(impl_observe_methods);

    fn on_factor(&mut self, logw: f64) {
        self.trace.log_factors += nan_to_neg_inf(logw);
    }

    fn finish(self) -> Trace {
        self.trace
    }
}

/// Score `model` against `base` strictly, returning an error rather than
/// panicking when the model's address structure does not match the base trace
/// (FG-20 / FG-21).
///
/// Returns `Ok((value, scored_trace))` when every sample site visited by the
/// model is present in `base` with a matching value type. Returns `Err` with
/// [`ErrorCode::UnexpectedModelStructure`] if the model visits an address absent
/// from `base` (a branch opened by a differing latent value), or
/// [`ErrorCode::AddressConflict`] if the model visits the same address twice.
///
/// A structurally incompatible `base` is a well-defined `Err`, never a panic
/// or a divergence: after the first missing/mismatched site the program is
/// handed a deterministic prior draw at each unscorable site so it terminates
/// inside its own support, and the `Err` reports the *first* offending address.
/// If you want the trace scored *with* those fresh draws (and told which sites
/// they were), use [`score_given_trace_reconciled`] instead.
///
/// # Example
///
/// ```rust
/// # use fugue::*;
/// # use fugue::runtime::interpreters::{score_given_trace_strict, PriorHandler};
/// # use rand::rngs::StdRng;
/// # use rand::SeedableRng;
/// let mut rng = StdRng::seed_from_u64(1);
/// let (_, base) = runtime::handler::run(
///     PriorHandler { rng: &mut rng, trace: Trace::default() },
///     sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()),
/// );
///
/// // Same structure: Ok.
/// assert!(score_given_trace_strict(
///     base.clone(),
///     sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()),
/// ).is_ok());
///
/// // Model reaches an address the base trace never recorded: Err.
/// let err = score_given_trace_strict(
///     base,
///     sample(addr!("y"), Normal::new(0.0, 1.0).unwrap()),
/// ).unwrap_err();
/// assert_eq!(err.code(), ErrorCode::UnexpectedModelStructure);
/// ```
pub fn score_given_trace_strict<A>(base: Trace, model: Model<A>) -> FugueResult<(A, Trace)> {
    let mut error: Option<FugueError> = None;
    let handler = StrictScoreGivenTrace {
        base,
        trace: Trace::default(),
        error: &mut error,
    };
    let (a, trace) = run(handler, model);
    match error {
        Some(e) => Err(e),
        None => Ok((a, trace)),
    }
}

/// Report of the structural differences reconciled by
/// [`score_given_trace_reconciled`].
#[derive(Clone, Debug, Default, PartialEq)]
pub struct ReconcileReport {
    /// Addresses visited by the model that were absent from the base trace (or
    /// present with the wrong value type). Each was proposed fresh from its
    /// prior and its `log_prob` accumulated into `log_prior`.
    pub fresh_addresses: Vec<Address>,
    /// Addresses present in the base trace that the model did NOT visit, or
    /// visited with a different value type (in which case the address also
    /// appears in `fresh_addresses`: a type change is a death plus a birth).
    /// Their contribution should be dropped by the caller (e.g. removed from the
    /// reverse-move density in an RJMCMC step).
    pub vanished_addresses: Vec<Address>,
}

/// The reconciling sibling of [`ScoreGivenTrace`], backing
/// [`score_given_trace_reconciled`].
///
/// Addresses present in the base trace are replayed and re-scored. NEW addresses
/// (absent, or present with a different value type) are sampled fresh from the
/// prior, their `log_prob` accumulated into `log_prior`, and recorded in
/// `fresh`. A duplicate address is still an error.
pub struct ReconcilingScoreGivenTrace<'r, 'f, 'e, R: RngCore> {
    /// RNG used to propose fresh values for new addresses.
    pub rng: &'r mut R,
    /// Base trace containing the fixed choices to replay/score.
    pub base: Trace,
    /// New trace accumulating the reconciled execution.
    pub trace: Trace,
    /// Addresses sampled fresh from the prior (not present in `base`), in
    /// visitation order. Borrowed so the driver can read it after `finish`.
    fresh: &'f mut Vec<Address>,
    /// First duplicate-address error encountered, surfaced by the driver.
    error: &'e mut Option<FugueError>,
}
impl<'r, 'f, 'e, R: RngCore> Handler for ReconcilingScoreGivenTrace<'r, 'f, 'e, R> {
    for_each_value_type!(impl_reconciling_score_sample_methods);
    for_each_value_type!(impl_observe_methods);

    fn on_factor(&mut self, logw: f64) {
        self.trace.log_factors += nan_to_neg_inf(logw);
    }

    fn finish(self) -> Trace {
        self.trace
    }
}

/// Score `model` against `base`, reconciling a differing address structure
/// instead of panicking (FG-20 / FG-21).
///
/// Addresses shared with `base` are replayed and re-scored under the current
/// model. Addresses the model introduces that are **not** in `base` (new
/// branches) are sampled fresh from their prior and their log-prior accumulated
/// — the RJMCMC-correct treatment of prior-proposed fresh dimensions. Addresses
/// in `base` that the model does **not** visit are reported as
/// [`ReconcileReport::vanished_addresses`] so the caller can drop their
/// contribution.
///
/// Returns `Err` with [`ErrorCode::AddressConflict`] only if the model visits
/// the same address twice.
///
/// # Example
///
/// ```rust
/// # use fugue::*;
/// # use fugue::runtime::interpreters::{score_given_trace_reconciled, PriorHandler};
/// # use rand::rngs::StdRng;
/// # use rand::SeedableRng;
/// let mut rng = StdRng::seed_from_u64(7);
/// let (_, base) = runtime::handler::run(
///     PriorHandler { rng: &mut rng, trace: Trace::default() },
///     sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()),
/// );
///
/// // Model drops "x" and introduces "y": "y" is proposed fresh, "x" vanished.
/// let (_v, trace, report) = score_given_trace_reconciled(
///     base,
///     &mut rng,
///     sample(addr!("y"), Normal::new(0.0, 1.0).unwrap()),
/// ).unwrap();
/// assert_eq!(report.fresh_addresses, vec![addr!("y")]);
/// assert_eq!(report.vanished_addresses, vec![addr!("x")]);
/// assert!(trace.log_prior.is_finite());
/// ```
pub fn score_given_trace_reconciled<A, R: RngCore>(
    base: Trace,
    rng: &mut R,
    model: Model<A>,
) -> FugueResult<(A, Trace, ReconcileReport)> {
    let mut error: Option<FugueError> = None;
    let mut fresh_addresses: Vec<Address> = Vec::new();
    // Snapshot base addresses (with their value types) up front so we can
    // compute vanished ones after the run consumes the handler.
    let base_addresses: Vec<(Address, &'static str)> = base
        .choices
        .iter()
        .map(|(a, c)| (a.clone(), c.value.type_name()))
        .collect();
    let handler = ReconcilingScoreGivenTrace {
        rng,
        base,
        trace: Trace::default(),
        fresh: &mut fresh_addresses,
        error: &mut error,
    };
    let (a, trace) = run(handler, model);
    if let Some(e) = error {
        return Err(e);
    }
    // Vanished = present in the base trace but not visited by this model run —
    // or visited with a DIFFERENT value type, in which case the handler sampled
    // it fresh (it is in `fresh`) and the old-typed site is gone (FG-N9): a
    // type change is a death plus a birth, and both sides must be reported.
    let vanished_addresses: Vec<Address> = base_addresses
        .into_iter()
        .filter(|(addr, ty)| {
            trace
                .choices
                .get(addr)
                .is_none_or(|c| c.value.type_name() != *ty)
        })
        .map(|(addr, _)| addr)
        .collect();
    Ok((
        a,
        trace,
        ReconcileReport {
            fresh_addresses,
            vanished_addresses,
        },
    ))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::addr;
    use crate::core::distribution::*;
    use crate::core::model::{observe, sample, ModelExt};
    use rand::rngs::StdRng;
    use rand::SeedableRng;

    /// EA-as-PPL F3 acceptance: grafting a prefix block between two executions
    /// of the same model, then re-scoring, yields exactly the weight a
    /// from-scratch replay of the spliced assignment would produce — the
    /// observable contract that "surgery + re-score" is exact.
    #[test]
    fn test_graft_rescore_equality() {
        let model_fn = || {
            sample(addr!("gene", 0), Normal::new(0.0, 1.0).unwrap()).and_then(|g0| {
                sample(addr!("gene", 1), Normal::new(0.0, 1.0).unwrap()).and_then(move |g1| {
                    observe(addr!("y"), Normal::new(g0 + g1, 0.5).unwrap(), 1.0)
                        .map(move |_| (g0, g1))
                })
            })
        };
        let mut rng = StdRng::seed_from_u64(21);
        let (_, ta) = crate::runtime::handler::run(
            PriorHandler {
                rng: &mut rng,
                trace: Trace::default(),
            },
            model_fn(),
        );
        let (_, tb) = crate::runtime::handler::run(
            PriorHandler {
                rng: &mut rng,
                trace: Trace::default(),
            },
            model_fn(),
        );

        // Splice tb's "gene#0" block into ta, then re-score.
        let mut spliced = ta.clone();
        spliced.graft_prefix("gene\u{23}0", &tb); // "gene#0"
        let (_, rescored) = crate::runtime::handler::run(
            ScoreGivenTrace {
                base: spliced,
                trace: Trace::default(),
            },
            model_fn(),
        );

        // From-scratch expectation: hand-build the same assignment and score it.
        let mut manual = Trace::default();
        manual.insert_choice(
            addr!("gene", 0),
            tb.choices[&addr!("gene", 0)].value.clone(),
            0.0,
        );
        manual.insert_choice(
            addr!("gene", 1),
            ta.choices[&addr!("gene", 1)].value.clone(),
            0.0,
        );
        let (_, expected) = crate::runtime::handler::run(
            ScoreGivenTrace {
                base: manual,
                trace: Trace::default(),
            },
            model_fn(),
        );
        assert!((rescored.total_log_weight() - expected.total_log_weight()).abs() < 1e-12);
    }

    #[test]
    fn prior_handler_samples_and_accumulates() {
        let mut rng = StdRng::seed_from_u64(7);
        let (_val, trace) = crate::runtime::handler::run(
            PriorHandler {
                rng: &mut rng,
                trace: Trace::default(),
            },
            sample(addr!("x"), Normal::new(0.0, 1.0).unwrap())
                .and_then(|x| observe(addr!("y"), Normal::new(x, 1.0).unwrap(), 0.5)),
        );
        assert!(trace.choices.contains_key(&addr!("x")));
        assert!(trace.log_prior.is_finite());
        assert!(trace.log_likelihood.is_finite());
    }

    #[test]
    fn replay_handler_reuses_values() {
        let mut rng = StdRng::seed_from_u64(8);
        let ((), base) = crate::runtime::handler::run(
            PriorHandler {
                rng: &mut rng,
                trace: Trace::default(),
            },
            sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()).map(|_| ()),
        );

        let ((), replayed) = crate::runtime::handler::run(
            ReplayHandler {
                rng: &mut rng,
                base: base.clone(),
                trace: Trace::default(),
            },
            sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()).map(|_| ()),
        );

        let x_base = base.get_f64(&addr!("x")).unwrap();
        let x_replay = replayed.get_f64(&addr!("x")).unwrap();
        assert_eq!(x_base, x_replay);
    }

    #[test]
    fn score_given_trace_scores_fixed_values() {
        let mut rng = StdRng::seed_from_u64(9);
        let (_a, base) = crate::runtime::handler::run(
            PriorHandler {
                rng: &mut rng,
                trace: Trace::default(),
            },
            sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()),
        );

        let (_a2, scored) = crate::runtime::handler::run(
            ScoreGivenTrace {
                base: base.clone(),
                trace: Trace::default(),
            },
            sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()),
        );

        // Should have same value and finite log_prior
        assert_eq!(scored.get_f64(&addr!("x")), base.get_f64(&addr!("x")));
        assert!(scored.log_prior.is_finite());
    }

    #[test]
    fn safe_variants_handle_mismatches() {
        // Build base trace with x as f64, then attempt to replay as bool
        let mut rng = StdRng::seed_from_u64(10);
        let (_a, base) = crate::runtime::handler::run(
            PriorHandler {
                rng: &mut rng,
                trace: Trace::default(),
            },
            sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()),
        );

        // SafeReplayHandler should sample fresh value for bool and continue
        let (_b, t1) = crate::runtime::handler::run(
            SafeReplayHandler {
                rng: &mut rng,
                base: base.clone(),
                trace: Trace::default(),
                warn_on_mismatch: true,
            },
            sample(addr!("x"), Bernoulli::new(0.5).unwrap()),
        );
        assert!(t1.log_prior.is_finite());

        // SafeScoreGivenTrace should mark as invalid by adding -inf
        let (_c, t2) = crate::runtime::handler::run(
            SafeScoreGivenTrace {
                base: base.clone(),
                trace: Trace::default(),
                warn_on_error: true,
            },
            sample(addr!("x"), Bernoulli::new(0.5).unwrap()),
        );
        assert!(t2.log_prior.is_infinite());
    }

    #[test]
    fn handlers_cover_all_types_sample_and_observe() {
        // Model with multiple types
        let model = sample(addr!("f"), Normal::new(0.0, 1.0).unwrap())
            .and_then(|_| sample(addr!("b"), Bernoulli::new(0.6).unwrap()))
            .and_then(|_| sample(addr!("u64"), Poisson::new(3.0).unwrap()))
            .and_then(|_| sample(addr!("usz"), Categorical::new(vec![0.3, 0.7]).unwrap()))
            .and_then(|_| observe(addr!("f_obs"), Normal::new(0.0, 1.0).unwrap(), 0.1))
            .and_then(|_| observe(addr!("b_obs"), Bernoulli::new(0.4).unwrap(), true))
            .and_then(|_| observe(addr!("u64_obs"), Poisson::new(2.0).unwrap(), 1))
            .and_then(|_| {
                observe(
                    addr!("usz_obs"),
                    Categorical::new(vec![0.5, 0.5]).unwrap(),
                    1,
                )
            });

        let (_a, t) = crate::runtime::handler::run(
            PriorHandler {
                rng: &mut StdRng::seed_from_u64(100),
                trace: Trace::default(),
            },
            model,
        );
        assert!(t.get_f64(&addr!("f")).is_some());
        assert!(t.get_bool(&addr!("b")).is_some());
        assert!(t.get_u64(&addr!("u64")).is_some());
        assert!(t.get_usize(&addr!("usz")).is_some());
        assert!(t.log_likelihood.is_finite());

        // Build base and score given trace for all types
        let base = t.clone();
        let (_sv, scored) = crate::runtime::handler::run(
            ScoreGivenTrace {
                base: base.clone(),
                trace: Trace::default(),
            },
            sample(addr!("f"), Normal::new(0.0, 1.0).unwrap())
                .and_then(|_| sample(addr!("b"), Bernoulli::new(0.6).unwrap()))
                .and_then(|_| sample(addr!("u64"), Poisson::new(3.0).unwrap()))
                .and_then(|_| sample(addr!("usz"), Categorical::new(vec![0.3, 0.7]).unwrap())),
        );
        assert!(scored.log_prior.is_finite());

        // Safe replay mismatches for integer/categorical types
        let (_sv2, safe) = crate::runtime::handler::run(
            SafeReplayHandler {
                rng: &mut StdRng::seed_from_u64(101),
                base: base.clone(),
                trace: Trace::default(),
                warn_on_mismatch: true,
            },
            sample(addr!("u64"), Bernoulli::new(0.5).unwrap()),
        );
        assert!(safe.log_prior.is_finite());
    }

    #[test]
    fn safe_score_given_trace_warn_flag_branches() {
        let mut rng = StdRng::seed_from_u64(102);
        let (_a, base) = crate::runtime::handler::run(
            PriorHandler {
                rng: &mut rng,
                trace: Trace::default(),
            },
            sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()),
        );
        // warn_on_error = false
        let (_b, t_false) = crate::runtime::handler::run(
            SafeScoreGivenTrace {
                base: base.clone(),
                trace: Trace::default(),
                warn_on_error: false,
            },
            sample(addr!("x"), Bernoulli::new(0.5).unwrap()),
        );
        assert!(t_false.log_prior.is_infinite());

        // warn_on_error = true
        let (_c, t_true) = crate::runtime::handler::run(
            SafeScoreGivenTrace {
                base: base.clone(),
                trace: Trace::default(),
                warn_on_error: true,
            },
            sample(addr!("x"), Bernoulli::new(0.5).unwrap()),
        );
        assert!(t_true.log_prior.is_infinite());
    }

    #[test]
    #[should_panic]
    fn replay_handler_panics_on_type_mismatch() {
        // Base has f64, replay expects bool -> panic as designed
        let mut rng = StdRng::seed_from_u64(103);
        let (_a, base) = crate::runtime::handler::run(
            PriorHandler {
                rng: &mut rng,
                trace: Trace::default(),
            },
            sample(addr!("x"), Normal::new(0.0, 1.0).unwrap()),
        );
        let (_b, _t) = crate::runtime::handler::run(
            ReplayHandler {
                rng: &mut rng,
                base: base.clone(),
                trace: Trace::default(),
            },
            sample(addr!("x"), Bernoulli::new(0.5).unwrap()),
        );
    }

    #[test]
    fn safe_replay_handler_samples_fresh_for_missing_address() {
        let mut rng = StdRng::seed_from_u64(104);
        // Base trace without address "z"
        let base = Trace::default();
        let (_a, t) = crate::runtime::handler::run(
            SafeReplayHandler {
                rng: &mut rng,
                base,
                trace: Trace::default(),
                warn_on_mismatch: true,
            },
            sample(addr!("z"), Normal::new(0.0, 1.0).unwrap()),
        );
        assert!(t.get_f64(&addr!("z")).is_some());
    }
}