probl-engine 0.1.0

//! Operators, and applying operations to every outcome of a distribution.

use crate::dist::{Budget, Dist};
use crate::error::{OpError, OpResult};
use crate::value::{EnumValue, Record, Value};
use probl_number::Integer;
use probl_syntax::ast::{BinOp, UnOp};
use std::borrow::Cow;
use std::sync::Arc;

/// Checked integer context: preserve ints, convert exactly integral finite
/// floats, and never round or admit other numeric types implicitly.
pub fn integer<'a>(v: &'a Value, context: &str, budget: &mut Budget) -> OpResult<Cow<'a, Integer>> {
    match v {
        Value::Int(n) => Ok(Cow::Borrowed(n)),
        Value::Float(f) if f.is_finite() && f.fract() == 0.0 => {
            let n = Integer::from_f64(*f).expect("finite integral float fits the integer hard limit");
            budget.integer_allocation(n.bits(), 1)?;
            budget.work(n.bits().div_ceil(64).max(1))?;
            Ok(Cow::Owned(n))
        }
        _ => {
            let found = match v {
                Value::Float(f) => format!("the float {f:?}"),
                _ => article(&v.kind()),
            };
            Err(OpError::new(format!(
                "`{context}` needs an int (or an exactly integral finite float), found {found}"
            )))
        }
    }
}

/// The outcomes of a value: those of a distribution, or the value itself.
pub fn outcomes(v: &Value) -> std::borrow::Cow<'_, [(Value, f64)]> {
    match v {
        Value::Dist(d) => std::borrow::Cow::Borrowed(&d.outcomes),
        other => std::borrow::Cow::Owned(vec![(other.clone(), 1.0)]),
    }
}

fn missing(v: &Value) -> f64 {
    match v {
        Value::Dist(d) => d.missing,
        _ => 0.0,
    }
}

/// Gather the results of applying an operation to each outcome of some
/// distributions into a distribution. Results that are themselves
/// distributions are mixed in. The result is always a distribution, even when
/// only one outcome is possible.
pub fn combine(results: Vec<(Value, f64)>, missing: f64, budget: &mut Budget) -> OpResult<Value> {
    if results.is_empty() {
        return Err(OpError::new("the result would be an empty distribution"));
    }
    let mut flat = Vec::with_capacity(results.len());
    let mut missing = missing;
    for (v, w) in results {
        match v {
            Value::Dist(d) => {
                missing += w * d.missing;
                flat.extend(d.outcomes.iter().map(|(x, p)| (x.clone(), w * p)));
            }
            other => flat.push((other, w)),
        }
    }
    budget.outcomes(flat.len() as u128)?;
    budget.work(flat.len() as u64)?;
    Ok(Dist::from_pairs(flat, missing.min(1.0)).into_value())
}

/// Apply `f` to a value, or to every outcome of a distribution.
pub fn lift1(a: &Value, budget: &mut Budget, f: impl Fn(&Value, &mut Budget) -> OpResult<Value>) -> OpResult<Value> {
    if !a.is_dist() {
        return f(a, budget);
    }
    let mut results = Vec::new();
    for (v, w) in outcomes(a).iter() {
        results.push((f(v, budget)?, *w));
    }
    combine(results, missing(a), budget)
}

/// Apply `f` to two values; distributions are independent draws.
pub fn lift2(
    a: &Value,
    b: &Value,
    budget: &mut Budget,
    f: impl Fn(&Value, &Value, &mut Budget) -> OpResult<Value>,
) -> OpResult<Value> {
    if !a.is_dist() && !b.is_dist() {
        return f(a, b, budget);
    }
    let (oa, ob) = (outcomes(a), outcomes(b));
    budget.outcomes(oa.len() as u128 * ob.len() as u128)?;
    budget.work((oa.len() * ob.len()) as u64)?;
    let mut results = Vec::with_capacity(oa.len() * ob.len());
    for (x, p) in oa.iter() {
        for (y, q) in ob.iter() {
            results.push((f(x, y, budget)?, p * q));
        }
    }
    let m = missing(a) + (1.0 - missing(a)) * missing(b);
    combine(results, m, budget)
}

/// A function lifted by [`lift_n`]; it may need the budget itself (to build
/// distributions).
pub type NaryFn<'a> = &'a dyn Fn(&[Value], &mut Budget) -> OpResult<Value>;

/// Apply `f` to a list of values, lifting over any that are distributions.
pub fn lift_n(args: &[Value], budget: &mut Budget, f: NaryFn) -> OpResult<Value> {
    if !args.iter().any(Value::is_dist) {
        return f(args, budget);
    }
    let size = args
        .iter()
        .fold(1u128, |acc, a| acc.saturating_mul(outcomes(a).len() as u128));
    budget.outcomes(size)?;
    budget.work(size as u64)?;
    let missing = args.iter().fold(0.0, |m, a| m + (1.0 - m) * missing(a));
    let mut results = Vec::new();
    let mut current = Vec::with_capacity(args.len());
    product(args, 0, &mut current, 1.0, f, budget, &mut results)?;
    combine(results, missing, budget)
}

fn product(
    args: &[Value],
    i: usize,
    current: &mut Vec<Value>,
    weight: f64,
    f: NaryFn,
    budget: &mut Budget,
    results: &mut Vec<(Value, f64)>,
) -> OpResult<()> {
    if i == args.len() {
        results.push((f(current, budget)?, weight));
        return Ok(());
    }
    for (v, w) in outcomes(&args[i]).iter() {
        current.push(v.clone());
        product(args, i + 1, current, weight * w, f, budget, results)?;
        current.pop();
    }
    Ok(())
}

// ── Probabilities, conditions and facts ──────────────────────────────────

pub fn article(kind: &str) -> String {
    let vowel = kind.starts_with(['a', 'e', 'i', 'o', 'u']);
    format!("{} {kind}", if vowel { "an" } else { "a" })
}

/// A value used as a probability (chance weights, `bernoulli`, `binomial`…):
/// a `prob` or a number checked at this boundary. Booleans need `prob(b)`.
pub fn to_prob(v: &Value) -> OpResult<f64> {
    if crate::analytic::contains(v) {
        return Err(crate::analytic::unsupported(
            "using an analytic outcome as a probability parameter",
        ));
    }
    match v {
        Value::Prob(_) | Value::Float(_) | Value::Int(_) => match make_prob(v)? {
            Value::Prob(p) => Ok(p),
            _ => unreachable!("make_prob returns a probability"),
        },
        Value::Bool(_) => Err(OpError::new("expected a probability, found a fact (true or false)")
            .help("convert a boolean explicitly with `prob(fact)`")),
        Value::Dist(_) => Err(OpError::new(format!("expected a probability, found a {}", v.kind()))
            .help("`P(…)` gives the probability that a distribution of facts is true")),
        other => Err(OpError::new(format!(
            "expected a probability, found {}",
            article(&other.kind())
        ))),
    }
}

/// Explicit, checked construction. This never draws, lifts or clamps.
pub fn make_prob(v: &Value) -> OpResult<Value> {
    let p = match v {
        Value::Bool(b) => {
            if *b {
                1.0
            } else {
                0.0
            }
        }
        Value::Prob(p) | Value::Float(p) => *p,
        Value::Int(n) if *n == 0 => 0.0,
        Value::Int(n) if *n == 1 => 1.0,
        Value::Int(_) => return Err(OpError::new("prob needs a finite number between 0 and 1")),
        _ => {
            return Err(
                OpError::new(format!("prob needs a number or bool, found {}", article(&v.kind())))
                    .help("draw distribution outcomes explicitly; use `P(d)` to query a boolean distribution"),
            );
        }
    };
    if !p.is_finite() || !(0.0..=1.0).contains(&p) {
        return Err(OpError::new("prob needs a finite number between 0 and 1"));
    }
    Ok(Value::Prob(p))
}

/// Computed chances may stray a few ulps outside [0, 1] through rounding.
/// This correction is only for algorithms known to produce probabilities;
/// user inputs always go through the strict `make_prob` boundary.
pub(crate) fn computed_prob(p: f64, what: &str) -> OpResult<Value> {
    const ROUNDING: f64 = 8.0 * f64::EPSILON;
    if !p.is_finite() || !(-ROUNDING..=1.0 + ROUNDING).contains(&p) {
        return Err(OpError::new(format!(
            "`{what}` could not compute a finite probability between 0 and 1"
        )));
    }
    Ok(Value::Prob(if p <= 0.0 { 0.0 } else { p.min(1.0) }))
}

pub fn fact(v: &Value, context: &str) -> OpResult<bool> {
    match v {
        Value::Bool(b) => Ok(*b),
        _ => Err(
            OpError::new(format!("`{context}` needs a bool, found {}", article(&v.kind())))
                .help("draw an outcome first with `~`; for a probability, draw from `bernoulli(p)`"),
        ),
    }
}

/// How likely a condition is to hold (docs/semantics.md, section 3).
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Condition {
    /// The probability that it's true.
    pub yes: f64,
    /// The probability that it's false.
    pub no: f64,
    /// The probability missing from a distribution of facts: it could be
    /// either. The three add up to one, up to rounding.
    pub missing: f64,
}

/// A fact, or a distribution of facts.
pub enum Truth {
    Fact(bool),
    Analytic(Arc<crate::analytic::Event>),
    Probability(f64),
    Uncertain(Arc<Dist>),
}

/// The operand of `and`, `or` or `not` (docs/semantics.md, section 2).
pub fn truth(v: &Value, op: &str) -> OpResult<Truth> {
    match v {
        Value::Bool(b) => Ok(Truth::Fact(*b)),
        Value::Event(e) => Ok(Truth::Analytic(e.clone())),
        Value::Dist(d) if d.truth().is_some() => Ok(Truth::Uncertain(d.clone())),
        Value::Prob(_) => Ok(Truth::Probability(to_prob(v)?)),
        other => Err(OpError::new(format!(
            "`{op}` needs bool, prob or dist[bool], found {}",
            article(&other.kind())
        ))),
    }
}

pub fn not(v: &Value, budget: &mut Budget) -> OpResult<Value> {
    match truth(v, "not")? {
        Truth::Fact(b) => Ok(Value::Bool(!b)),
        Truth::Analytic(e) => Ok(crate::analytic::Event {
            draw: e.draw.clone(),
            yes: e.yes.complement(),
        }
        .value()),
        Truth::Probability(p) => computed_prob(1.0 - p, "not"),
        Truth::Uncertain(d) => lift1(&Value::Dist(d), budget, |x, _| match x {
            Value::Bool(b) => Ok(Value::Bool(!b)),
            _ => unreachable!("checked by `truth`"),
        }),
    }
}

/// Compose independent boolean recipes. Bound facts retain their identity.
pub fn logic(and: bool, a: Truth, b: Truth, budget: &mut Budget) -> OpResult<Value> {
    if matches!(a, Truth::Analytic(_)) || matches!(b, Truth::Analytic(_)) {
        let value = |t| match t {
            Truth::Fact(b) => Value::Bool(b),
            Truth::Probability(p) => Value::Prob(p),
            Truth::Uncertain(d) => Value::Dist(d),
            Truth::Analytic(e) => Value::Event(e),
        };
        let (a, b) = (value(a), value(b));
        for v in [&a, &b] {
            if let Value::Event(e) = v {
                budget.collection(e.yes.0.len() as u128 + e.draw.domain.0.len() as u128)?;
                budget.work((e.yes.0.len() + e.draw.domain.0.len()) as u64)?;
            }
        }
        return crate::analytic::logic(and, &a, &b);
    }
    let op = |x: bool, y: bool| if and { x && y } else { x || y };
    match (a, b) {
        (Truth::Fact(x), Truth::Fact(y)) => Ok(Value::Bool(op(x, y))),
        (Truth::Probability(p), Truth::Probability(q)) => computed_prob(
            if and { p * q } else { p + (1.0 - p) * q },
            if and { "and" } else { "or" },
        ),
        (Truth::Probability(p), Truth::Fact(b)) | (Truth::Fact(b), Truth::Probability(p)) => {
            Ok(Value::Prob(if b == and {
                p
            } else if b {
                1.0
            } else {
                0.0
            }))
        }
        (Truth::Fact(x), Truth::Uncertain(d)) | (Truth::Uncertain(d), Truth::Fact(x)) => {
            lift1(&Value::Dist(d), budget, |v, _| match v {
                Value::Bool(y) => Ok(Value::Bool(op(x, *y))),
                _ => unreachable!("checked by `truth`"),
            })
        }
        (a, b) => {
            let law = |v| match v {
                Truth::Uncertain(d) => Value::Dist(d),
                Truth::Probability(p) => Dist::bernoulli(p).into_value(),
                Truth::Fact(b) => Value::Bool(b),
                Truth::Analytic(_) => unreachable!("handled above"),
            };
            lift2(&law(a), &law(b), budget, |x, y, _| match (x, y) {
                (Value::Bool(x), Value::Bool(y)) => Ok(Value::Bool(op(*x, *y))),
                _ => unreachable!("checked by `truth`"),
            })
        }
    }
}

/// The boolean law requested by an observed anonymous draw.
pub fn boolean_law(v: &Value) -> OpResult<Value> {
    match truth(v, "observe ~")? {
        Truth::Analytic(e) => Ok(Value::Event(e)),
        Truth::Fact(b) => Ok(Dist::bernoulli(if b { 1.0 } else { 0.0 }).into_value()),
        Truth::Probability(p) => Ok(Dist::bernoulli(p).into_value()),
        Truth::Uncertain(d) => Ok(Value::Dist(d)),
    }
}

pub fn condition(v: &Value) -> OpResult<Condition> {
    if matches!(v, Value::Analytic(_)) {
        return Err(crate::analytic::unsupported(
            "using a continuous outcome as a probability condition",
        ));
    }
    if matches!(v, Value::Float(_) | Value::Int(_)) {
        let p = to_prob(v)?;
        return Ok(Condition {
            yes: p,
            no: 1.0 - p,
            missing: 0.0,
        });
    }
    match truth(v, "condition")? {
        Truth::Analytic(e) => Ok(Condition {
            yes: e.probability(),
            no: 1.0 - e.probability(),
            missing: 0.0,
        }),
        Truth::Fact(b) => Ok(Condition {
            yes: if b { 1.0 } else { 0.0 },
            no: if b { 0.0 } else { 1.0 },
            missing: 0.0,
        }),
        Truth::Probability(p) => Ok(Condition {
            yes: p,
            no: 1.0 - p,
            missing: 0.0,
        }),
        Truth::Uncertain(d) => {
            let (yes, no) = d.truth().expect("checked by truth");
            Ok(Condition {
                yes,
                no,
                missing: d.missing,
            })
        }
    }
}

// ── Operators ────────────────────────────────────────────────────────────

pub fn unary(op: UnOp, v: &Value, budget: &mut Budget) -> OpResult<Value> {
    match op {
        UnOp::Neg => lift1(v, budget, |x, budget| match x {
            Value::Int(i) => {
                budget.integer_work(i, &Integer::ZERO, false)?;
                {
                    let n = i.negated();
                    budget.integer_allocation(n.bits(), 1)?;
                    Ok(Value::Int(n))
                }
            }
            Value::Float(f) | Value::Prob(f) => Ok(Value::Float(-f)),
            Value::Complex(z) => Ok(Value::Complex(z.negated())),
            Value::Analytic(a) => {
                let mut a = (**a).clone();
                a.scale = -a.scale;
                a.offset = -a.offset;
                a.value()
            }
            other => Err(OpError::new(format!("can't negate {}", article(&other.kind())))),
        }),
        UnOp::Not => not(v, budget),
        UnOp::Typeof => unreachable!("typeof is lowered to an interpreter intrinsic"),
    }
}

pub fn binary(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
    match op {
        BinOp::And | BinOp::Or => unreachable!("`and` and `or` are evaluated lazily by the interpreter"),
        BinOp::To => unreachable!("`to` is lowered to a built-in"),
        BinOp::Range | BinOp::RangeExcl => range(op, a, b, budget),
        BinOp::In => lift2(a, b, budget, |x, coll, budget| {
            contains(coll, x, budget).map(Value::Bool)
        }),
        BinOp::NotIn => lift2(a, b, budget, |x, coll, budget| {
            contains(coll, x, budget).map(|c| Value::Bool(!c))
        }),
        _ => lift2(a, b, budget, |x, y, budget| binary_plain(op, x, y, budget)),
    }
}

fn binary_plain(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
    if matches!(a, Value::Analytic(_)) || matches!(b, Value::Analytic(_)) {
        return crate::analytic::binary(op, a, b);
    }
    if matches!(
        op,
        BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge
    ) && (crate::analytic::contains(a) || crate::analytic::contains(b))
    {
        if matches!(op, BinOp::Eq | BinOp::Ne) {
            if let (Value::Event(x), Value::Bool(b)) | (Value::Bool(b), Value::Event(x)) = (a, b) {
                return if *b == (op == BinOp::Eq) {
                    Ok(Value::Event(x.clone()))
                } else {
                    not(&Value::Event(x.clone()), budget)
                };
            }
        }
        if let (Value::Event(x), Value::Event(y)) = (a, b) {
            if x.draw.id == y.draw.id && matches!(op, BinOp::Eq | BinOp::Ne) {
                let both = x.yes.intersect(&y.yes);
                let neither = x.yes.complement().intersect(&y.yes.complement());
                let equal = both.complement().intersect(&neither.complement()).complement();
                return Ok(crate::analytic::Event {
                    draw: x.draw.clone(),
                    yes: if op == BinOp::Eq { equal } else { equal.complement() },
                }
                .value());
            }
        }
        return Err(crate::analytic::unsupported(
            "comparing aggregate or boolean analytic outcomes",
        ));
    }
    if let Some(v) = continuous_binary(op, a, b)? {
        return Ok(v);
    }
    if matches!(
        op,
        BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge
    ) {
        for v in [a, b] {
            if let Value::Int(n) = v {
                budget.integer_work(n, &Integer::ONE, false)?;
            }
            if let Value::Str(s) = v {
                budget.string_work(s)?;
            }
        }
    }
    match op {
        BinOp::Eq => Ok(Value::Bool(equals(a, b))),
        BinOp::Ne => Ok(Value::Bool(!equals(a, b))),
        BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge => {
            let ord = compare(a, b)?;
            Ok(Value::Bool(match op {
                BinOp::Lt => ord.is_lt(),
                BinOp::Le => ord.is_le(),
                BinOp::Gt => ord.is_gt(),
                _ => ord.is_ge(),
            }))
        }
        BinOp::Add => add(a, b, budget),
        BinOp::Sub => sub(a, b, budget),
        _ => arith(op, a, b, budget),
    }
}

fn add(a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
    match (a, b) {
        (Value::Str(x), Value::Str(y)) => {
            let size = x
                .len()
                .checked_add(y.len())
                .ok_or_else(|| OpError::limit("string size overflow"))?;
            budget.string_size(size)?;
            let mut out = String::new();
            crate::text::push(&mut out, x, budget)?;
            crate::text::push(&mut out, y, budget)?;
            Ok(Value::str(&out))
        }
        (Value::List(x), Value::List(y)) => {
            budget.collection(x.len() as u128 + y.len() as u128)?;
            budget.work(x.len() as u64 + y.len() as u64)?;
            let mut items = x.to_vec();
            items.extend(y.iter().cloned());
            Ok(Value::list(items))
        }
        (Value::Date(d), n @ (Value::Int(_) | Value::Float(_)))
        | (n @ (Value::Int(_) | Value::Float(_)), Value::Date(d)) => {
            let n = integer(n, "date offset", budget)?;
            date_plus(*d, n.to_i64().ok_or_else(|| OpError::new("date out of range"))?)
        }
        (Value::Str(_), _) | (_, Value::Str(_)) => {
            Err(
                OpError::new(format!("can't add {} and {}", article(&a.kind()), article(&b.kind())))
                    .help("to build text, use interpolation: \"total: {x}\""),
            )
        }
        _ => arith(BinOp::Add, a, b, budget),
    }
}

fn sub(a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
    match (a, b) {
        (Value::Date(x), Value::Date(y)) => Ok(Value::Int((*x as i64 - *y as i64).into())),
        (Value::Date(d), n @ (Value::Int(_) | Value::Float(_))) => {
            match integer(n, "date offset", budget)?.negated().to_i64() {
                Some(m) => date_plus(*d, m),
                None => Err(OpError::new("date out of range")),
            }
        }
        _ => arith(BinOp::Sub, a, b, budget),
    }
}

fn date_plus(d: i32, n: i64) -> OpResult<Value> {
    crate::dates::add_days(d, n)
        .map(Value::Date)
        .ok_or_else(|| OpError::new("date out of range"))
}

/// Numbers for arithmetic: ints, floats and probabilities (not facts).
fn number(v: &Value) -> Option<f64> {
    match v {
        Value::Bool(_) => None,
        _ => v.as_f64(),
    }
}

fn arith(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
    let bad = || {
        OpError::new(format!(
            "can't use `{}` with {} and {}",
            op.symbol(),
            article(&a.kind()),
            article(&b.kind())
        ))
    };
    if let (BinOp::Pow, Value::Complex(z), Value::Int(n)) = (op, a, b) {
        budget.integer_work(n, &Integer::ONE, false)?;
        budget.work(n.bits())?;
        return Ok(Value::Complex(z.pow_integer(n)?));
    }
    if matches!(a, Value::Complex(_)) || matches!(b, Value::Complex(_)) {
        let (Some(x), Some(y)) = (a.as_complex(), b.as_complex()) else {
            return Err(bad());
        };
        let z = match op {
            BinOp::Add => x.plus(y),
            BinOp::Sub => x.minus(y),
            BinOp::Mul => x.times(y),
            BinOp::Div => x.divided_by(y),
            BinOp::Pow => match b {
                Value::Int(n) => {
                    budget.integer_work(n, &Integer::ONE, false)?;
                    budget.work(n.bits())?;
                    x.pow_integer(n)
                }
                _ => Err(OpError::new("a complex power needs an int exponent")),
            },
            _ => Err(bad()),
        }?;
        return Ok(Value::Complex(z));
    }
    if let (Value::Int(x), Value::Int(y)) = (a, b) {
        budget.integer_work(x, y, matches!(op, BinOp::Mul | BinOp::Div | BinOp::IntDiv | BinOp::Mod))?;
        let n = match op {
            BinOp::Add => x.add(y)?,
            BinOp::Sub => x.sub(y)?,
            BinOp::Mul => {
                if !x.is_zero() && !y.is_zero() {
                    budget.integer_bits((x.bits() + y.bits()).saturating_sub(1))?;
                }
                x.mul(y)?
            }
            BinOp::Div => {
                if y.is_zero() {
                    return Err(division_by_zero());
                }
                return x
                    .ratio(y)
                    .map(Value::Float)
                    .ok_or_else(|| OpError::new("division result is too large for a finite float"));
            }
            BinOp::IntDiv => x.div_mod(y)?.0,
            BinOp::Mod => x.div_mod(y)?.1,
            BinOp::Pow => return int_power(x, y, budget),
            _ => return Err(bad()),
        };
        budget.integer_allocation(n.bits(), 1)?;
        return Ok(Value::Int(n));
    }
    let (Some(x), Some(y)) = (number(a), number(b)) else {
        return Err(bad());
    };
    let v = match op {
        BinOp::Add => x + y,
        BinOp::Sub => x - y,
        BinOp::Mul => x * y,
        BinOp::Div => {
            if y == 0.0 {
                return Err(division_by_zero());
            }
            x / y
        }
        BinOp::IntDiv => {
            if y == 0.0 {
                return Err(division_by_zero());
            }
            let q = (x / y).floor();
            let n = Integer::from_f64(q).ok_or_else(|| OpError::new("integer division result is not finite"))?;
            budget.integer_allocation(n.bits(), 1)?;
            return Ok(Value::Int(n));
        }
        BinOp::Mod => {
            if y == 0.0 {
                return Err(division_by_zero());
            }
            x - (x / y).floor() * y
        }
        BinOp::Pow => libm::pow(x, y),
        _ => return Err(bad()),
    };
    if !v.is_finite() {
        return Err(OpError::new(format!(
            "`{}` gave a result that isn't a finite number",
            op.symbol()
        )));
    }
    Ok(Value::Float(v))
}

/// A continuous distribution compared with a number: a distribution of
/// facts, from its CDF. Nothing else is defined on one yet
/// (docs/semantics.md, section 13).
fn continuous_binary(op: BinOp, a: &Value, b: &Value) -> OpResult<Option<Value>> {
    let (family, other, flipped) = match (a, b) {
        (Value::Continuous(f), other) => (f, other, false),
        (other, Value::Continuous(f)) => (f, other, true),
        _ => return Ok(None),
    };
    let needs_value = || {
        OpError::new(format!(
            "`{}` needs a value, not a {} distribution",
            op.symbol(),
            family.name()
        ))
        .help("draw a value first, like `let x ~ normal(0, 1)`; comparing a distribution with a number works too")
    };
    let x = match number(other) {
        Some(x) if !x.is_nan() => x,
        _ => return Err(needs_value()),
    };
    // For a continuous X, P(X < x) = P(X ≤ x) = F(x).
    let below = family.cdf(x);
    let yes = match (op, flipped) {
        (BinOp::Lt | BinOp::Le, false) | (BinOp::Gt | BinOp::Ge, true) => below,
        (BinOp::Gt | BinOp::Ge, false) | (BinOp::Lt | BinOp::Le, true) => 1.0 - below,
        (BinOp::Eq, _) => 0.0,
        (BinOp::Ne, _) => 1.0,
        _ => return Err(needs_value()),
    };
    Ok(Some(Dist::bernoulli(yes).into_value()))
}

fn division_by_zero() -> OpError {
    OpError::new("division by zero")
}

fn int_power(base: &Integer, exponent: &Integer, budget: &mut Budget) -> OpResult<Value> {
    if exponent.is_zero() {
        return Ok(Value::Int(Integer::ONE));
    }
    if base.is_zero() {
        return if exponent.is_negative() {
            Err(division_by_zero())
        } else {
            Ok(Value::Int(Integer::ZERO))
        };
    }
    if exponent.is_negative() {
        let reciprocal = Integer::ONE.ratio(&base.abs()).unwrap_or(0.0);
        let power = exponent.abs().to_f64().unwrap_or(f64::INFINITY);
        let result = libm::pow(reciprocal, power);
        return Ok(Value::Float(if base.is_negative() && exponent.is_odd() {
            -result
        } else {
            result
        }));
    }
    if *base == 1 || *base == -1 {
        return Ok(Value::Int(if *base == -1 && exponent.is_odd() {
            (-1).into()
        } else {
            Integer::ONE
        }));
    }
    let mut n = exponent
        .to_u64()
        .ok_or_else(|| OpError::limit("integer power exceeds the integer size limit"))?;
    budget.integer_bits(base.bits().saturating_sub(1).saturating_mul(n).saturating_add(1))?;
    let mut result = Integer::ONE;
    let mut b = base.clone();
    while n != 0 {
        if n & 1 != 0 {
            budget.integer_work(&result, &b, true)?;
            result = result.mul(&b)?;
            budget.integer_bits(result.bits())?;
        }
        n >>= 1;
        if n != 0 {
            budget.integer_work(&b, &b, true)?;
            b = b.mul(&b)?;
            budget.integer_bits(b.bits())?;
        }
    }
    budget.integer_allocation(result.bits(), 1)?;
    Ok(Value::Int(result))
}

fn numeric_compare(a: &Value, b: &Value) -> Option<std::cmp::Ordering> {
    match (a, b) {
        (Value::Int(a), Value::Int(b)) => Some(a.cmp(b)),
        (Value::Int(a), Value::Float(b) | Value::Prob(b)) => a.cmp_f64(*b),
        (Value::Float(a) | Value::Prob(a), Value::Int(b)) => b.cmp_f64(*a).map(std::cmp::Ordering::reverse),
        (Value::Float(a) | Value::Prob(a), Value::Float(b) | Value::Prob(b)) => a.partial_cmp(b),
        _ => None,
    }
}

/// The language's `==`: numbers compare by value across int, float and prob.
pub fn equals(a: &Value, b: &Value) -> bool {
    if let (Value::Complex(x), Value::Complex(y)) = (a, b) {
        return x == y;
    }
    if let (Value::Complex(z), real) | (real, Value::Complex(z)) = (a, b) {
        if z.im() != 0.0 {
            return false;
        }
        return match real {
            // Compare integers without rounding them through f64.
            Value::Int(n) => {
                let x = z.re();
                n.cmp_f64(x).is_some_and(|c| c.is_eq())
            }
            Value::Float(x) | Value::Prob(x) => z.re() == *x,
            _ => false,
        };
    }
    if let Some(c) = numeric_compare(a, b) {
        return c.is_eq();
    }
    if matches!(a, Value::Float(x) | Value::Prob(x) if x.is_nan())
        || matches!(b, Value::Float(x) | Value::Prob(x) if x.is_nan())
    {
        return false;
    }
    a == b
}

/// Ordering for `<`, `>` and friends.
pub fn compare(a: &Value, b: &Value) -> OpResult<std::cmp::Ordering> {
    if matches!(a, Value::Complex(_)) || matches!(b, Value::Complex(_)) {
        return Err(OpError::new("complex values have no ordering").help("compare `abs(z)`, `real(z)` or `imag(z)`"));
    }
    if let Some(c) = numeric_compare(a, b) {
        return Ok(c);
    }
    if matches!(a, Value::Float(x) | Value::Prob(x) if x.is_nan())
        || matches!(b, Value::Float(x) | Value::Prob(x) if x.is_nan())
    {
        return Err(OpError::new("can't order NaN"));
    }
    match (a, b) {
        (Value::Str(x), Value::Str(y)) => Ok(x.cmp(y)),
        (Value::Date(x), Value::Date(y)) => Ok(x.cmp(y)),
        (Value::Enum(x), Value::Enum(y)) if x.ty == y.ty => Ok(x.variant.cmp(&y.variant)),
        (Value::List(x), Value::List(y)) => {
            for (p, q) in x.iter().zip(y.iter()) {
                let c = compare(p, q)?;
                if c.is_ne() {
                    return Ok(c);
                }
            }
            Ok(x.len().cmp(&y.len()))
        }
        _ => Err(OpError::new(format!(
            "can't compare {} with {}",
            article(&a.kind()),
            article(&b.kind())
        ))),
    }
}

/// Whether `item` is in `coll` (for `in`).
pub fn contains(coll: &Value, item: &Value, budget: &mut Budget) -> OpResult<bool> {
    if crate::analytic::contains(coll) || crate::analytic::contains(item) {
        return Err(crate::analytic::unsupported("membership involving analytic outcomes"));
    }
    if let Value::Str(s) = coll {
        budget.string_work(s)?;
        if let Value::Str(s) = item {
            budget.string_work(s)?;
        }
    }
    Ok(match coll {
        Value::List(items) => items.iter().any(|x| equals(x, item)),
        Value::Map(m) => m.contains_key(item),
        Value::Bag(b) => b.get(item).is_some_and(|n| *n > 0),
        Value::Range(lo, hi) => match item {
            Value::Int(n) => n >= lo && n <= hi,
            Value::Float(x) | Value::Prob(x) if x.is_finite() && x.fract() == 0.0 => {
                lo.cmp_f64(*x).is_some_and(|c| !c.is_gt()) && hi.cmp_f64(*x).is_some_and(|c| !c.is_lt())
            }
            _ => false,
        },
        Value::Str(s) => match item {
            Value::Str(sub) => s.contains(&**sub),
            _ => {
                return Err(OpError::new(format!(
                    "can't look for {} in a string",
                    article(&item.kind())
                )));
            }
        },
        other => return Err(OpError::new(format!("can't look inside {}", article(&other.kind())))),
    })
}

fn range(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
    if a.is_uncertain() || b.is_uncertain() {
        return Err(OpError::new("a range needs plain whole numbers, not distributions")
            .help("draw a value first, like `let n ~ d6`"));
    }
    let lo = integer(a, "range bound", budget)?;
    let hi = integer(b, "range bound", budget)?;
    let hi = if op == BinOp::RangeExcl {
        hi.sub(&Integer::ONE)?
    } else {
        hi.into_owned()
    };
    budget.integer_work(&lo, &hi, false)?;
    if op == BinOp::RangeExcl {
        budget.integer_allocation(hi.bits(), 1)?;
    }
    Ok(Value::Range(lo.into_owned(), hi))
}

/// The number of integers in `lo..=hi`, without overflowing.
pub fn range_len(lo: &Integer, hi: &Integer) -> OpResult<Integer> {
    if hi < lo {
        Ok(Integer::ZERO)
    } else {
        Ok(hi.sub(lo)?.add(&Integer::ONE)?)
    }
}

/// Bounded materialization count, never truncating a large length.
pub fn range_count(lo: &Integer, hi: &Integer) -> OpResult<u128> {
    range_len(lo, hi)?
        .to_u128()
        .ok_or_else(|| OpError::limit("range has too many elements to materialize"))
}

// ── Collections ──────────────────────────────────────────────────────────

pub fn field(v: &Value, name: &str, budget: &mut Budget) -> OpResult<Value> {
    lift1(v, budget, |x, _| match x {
        Value::Record(r) => r.get(name).cloned().ok_or_else(|| {
            let known: Vec<String> = r.fields.iter().map(|(n, _)| format!("`{n}`")).collect();
            OpError::new(format!("{} has no field `{name}`", article(&x.kind())))
                .help(format!("its fields are {}", known.join(", ")))
        }),
        other => Err(OpError::new(format!(
            "can't read the field `{name}` of {}",
            article(&other.kind())
        ))),
    })
}

pub fn index(coll: &Value, i: &Value, budget: &mut Budget) -> OpResult<Value> {
    lift2(coll, i, budget, |c, i, budget| {
        let v = index_plain(c, i, budget)?;
        if let (Value::Range(lo, hi), Value::Int(n)) = (c, &v) {
            budget.integer_work(lo, hi, false)?;
            budget.integer_allocation(n.bits(), 1)?;
        }
        Ok(v)
    })
}

pub fn index_plain(coll: &Value, i: &Value, budget: &mut Budget) -> OpResult<Value> {
    match coll {
        Value::List(items) => {
            let k = as_index(i, items.len() as u128, budget)?;
            Ok(items[k as usize].clone())
        }
        Value::Range(lo, hi) => {
            let k = integer(i, "index", budget)?;
            if k.is_negative() {
                return Err(OpError::new("index out of range"));
            }
            let n = lo.add(&k)?;
            if &n > hi {
                return Err(OpError::new("index out of range"));
            }
            Ok(Value::Int(n))
        }
        Value::Str(s) => {
            budget.string_work(s)?;
            let k = as_index(i, s.chars().count() as u128, budget)?;
            let c = s.chars().nth(k as usize).expect("checked scalar index");
            crate::text::value(c.encode_utf8(&mut [0; 4]), budget)
        }
        Value::Map(m) => m.get(i).cloned().ok_or_else(|| {
            OpError::new(format!("the key {i:?} isn't in the map"))
                .help("use `get(key, default)` for keys that may be missing")
        }),
        other => Err(OpError::new(format!("can't index {}", article(&other.kind())))),
    }
}

pub fn as_index(i: &Value, len: u128, budget: &mut Budget) -> OpResult<u128> {
    let k = integer(i, "index", budget)?.to_u128();
    k.filter(|k| *k < len).ok_or_else(|| {
        OpError::new(format!("index {i} is out of range for a length of {len}")).help("indices start at 0")
    })
}

pub fn make_record(ty: Option<Arc<str>>, mut fields: Vec<(Arc<str>, Value)>) -> Value {
    fields.sort_by(|a, b| a.0.cmp(&b.0));
    Value::record(Record { ty, fields })
}

pub fn with_fields(base: &Value, updates: &[(Arc<str>, Value)]) -> OpResult<Value> {
    let Value::Record(r) = base else {
        return Err(OpError::new(format!(
            "`with` needs a record, found {}",
            article(&base.kind())
        )));
    };
    let mut r = Record::clone(r);
    for (name, v) in updates {
        match r.get_mut(name) {
            Some(slot) => *slot = v.clone(),
            None => {
                return Err(OpError::new(format!("{} has no field `{name}`", article(&base.kind()))));
            }
        }
    }
    Ok(Value::record(r))
}

pub fn enum_value(ty: u32, variant: u32, name: &str) -> Value {
    Value::Enum(Arc::new(EnumValue {
        ty,
        variant,
        name: Arc::from(name),
    }))
}

/// Is `v` certainly `truth`? (For `and`/`or` whose right side has statements.)
pub fn is_certain(v: &Value, truth: bool) -> bool {
    matches!(v, Value::Bool(b) if *b == truth)
}

#[cfg(test)]
mod probability_tests {
    use super::*;

    #[test]
    fn computed_probabilities_only_correct_boundary_roundoff() {
        assert_eq!(computed_prob(1.0 + f64::EPSILON, "test").unwrap(), Value::Prob(1.0));
        assert_eq!(computed_prob(-f64::EPSILON, "test").unwrap(), Value::Prob(0.0));
        assert_eq!(computed_prob(1e-300, "test").unwrap(), Value::Prob(1e-300));
        for p in [1.001, -0.001, f64::INFINITY, f64::NEG_INFINITY, f64::NAN] {
            assert!(computed_prob(p, "test").is_err());
        }
    }
}