Skip to main content

probl_engine/
ops.rs

1//! Operators, and applying operations to every outcome of a distribution.
2
3use crate::dist::{Budget, Dist};
4use crate::error::{Fault, OpError, OpResult};
5use crate::value::{EnumValue, Record, Value};
6use probl_number::Integer;
7use probl_syntax::ast::{BinOp, UnOp};
8use std::borrow::Cow;
9use std::sync::Arc;
10
11/// Checked integer context: preserve ints, convert exactly integral finite
12/// floats, and never round or admit other numeric types implicitly.
13pub fn integer<'a>(v: &'a Value, context: &str, budget: &mut Budget) -> OpResult<Cow<'a, Integer>> {
14    match v {
15        Value::Int(n) => Ok(Cow::Borrowed(n)),
16        Value::Float(f) if f.is_finite() && f.fract() == 0.0 => {
17            let n = Integer::from_f64(*f).expect("finite integral float fits the integer hard limit");
18            budget.integer_allocation(n.bits(), 1)?;
19            budget.work(n.bits().div_ceil(64).max(1))?;
20            Ok(Cow::Owned(n))
21        }
22        Value::Analytic(a) if a.int => Err(crate::analytic::unsupported(&format!(
23            "an int outcome of a continuous draw as `{context}`"
24        ))),
25        _ => {
26            let found = match v {
27                Value::Float(f) => format!("the float {f:?}"),
28                _ => article(&v.kind()),
29            };
30            Err(OpError::new(format!(
31                "`{context}` needs an int (or an exactly integral finite float), found {found}"
32            )))
33        }
34    }
35}
36
37/// A count law with its outcomes listed, for what needs them; anything
38/// else as it is. (Out of line: a lazy law is rare, and the lifting
39/// functions, generic and hot, must stay small enough to inline.)
40#[cold]
41#[inline(never)]
42pub fn listed<'a>(v: &'a Value, budget: &mut Budget) -> OpResult<std::borrow::Cow<'a, Value>> {
43    match v {
44        Value::Counts(c) => Ok(std::borrow::Cow::Owned(c.list(budget)?.into_value())),
45        v => Ok(std::borrow::Cow::Borrowed(v)),
46    }
47}
48
49/// The outcomes of a value: those of a distribution, or the value itself.
50pub fn outcomes(v: &Value) -> std::borrow::Cow<'_, [(Value, f64)]> {
51    match v {
52        Value::Dist(d) => std::borrow::Cow::Borrowed(&d.outcomes),
53        other => std::borrow::Cow::Owned(vec![(other.clone(), 1.0)]),
54    }
55}
56
57fn missing(v: &Value) -> f64 {
58    match v {
59        Value::Dist(d) => d.missing,
60        _ => 0.0,
61    }
62}
63
64/// Gather the results of applying an operation to each outcome of some
65/// distributions into a distribution. Results that are themselves
66/// distributions are mixed in. The result is always a distribution, even when
67/// only one outcome is possible.
68pub fn combine(results: Vec<(Value, f64)>, missing: f64, budget: &mut Budget) -> OpResult<Value> {
69    if results.is_empty() {
70        return Err(OpError::new("the result would be an empty distribution"));
71    }
72    let mut flat = Vec::with_capacity(results.len());
73    let mut missing = missing;
74    for (v, w) in results {
75        let v = if let Value::Counts(_) = v {
76            listed(&v, budget)?.into_owned()
77        } else {
78            v
79        };
80        match v {
81            Value::Dist(d) => {
82                missing += w * d.missing;
83                flat.extend(d.outcomes.iter().map(|(x, p)| (x.clone(), w * p)));
84            }
85            other => flat.push((other, w)),
86        }
87    }
88    budget.outcomes(flat.len() as u128)?;
89    budget.work(flat.len() as u64)?;
90    Ok(Dist::from_pairs(flat, missing.min(1.0)).into_value())
91}
92
93/// Apply `f` to a value, or to every outcome of a distribution.
94pub fn lift1(a: &Value, budget: &mut Budget, f: impl Fn(&Value, &mut Budget) -> OpResult<Value>) -> OpResult<Value> {
95    if !a.is_dist() && !matches!(a, Value::Counts(_)) {
96        return f(a, budget);
97    }
98    let a = listed(a, budget)?;
99    let a = a.as_ref();
100    let mut results = Vec::new();
101    for (v, w) in outcomes(a).iter() {
102        results.push((f(v, budget)?, *w));
103    }
104    combine(results, missing(a), budget)
105}
106
107/// Apply `f` to two values; distributions are independent draws.
108pub fn lift2(
109    a: &Value,
110    b: &Value,
111    budget: &mut Budget,
112    f: impl Fn(&Value, &Value, &mut Budget) -> OpResult<Value>,
113) -> OpResult<Value> {
114    let plain = |v: &Value| !v.is_dist() && !matches!(v, Value::Counts(_));
115    if plain(a) && plain(b) {
116        return f(a, b, budget);
117    }
118    let (a, b) = (listed(a, budget)?, listed(b, budget)?);
119    let (a, b) = (a.as_ref(), b.as_ref());
120    let (oa, ob) = (outcomes(a), outcomes(b));
121    budget.outcomes(oa.len() as u128 * ob.len() as u128)?;
122    budget.work((oa.len() * ob.len()) as u64)?;
123    let mut results = Vec::with_capacity(oa.len() * ob.len());
124    for (x, p) in oa.iter() {
125        for (y, q) in ob.iter() {
126            results.push((f(x, y, budget)?, p * q));
127        }
128    }
129    let m = missing(a) + (1.0 - missing(a)) * missing(b);
130    combine(results, m, budget)
131}
132
133/// A function lifted by [`lift_n`]; it may need the budget itself (to build
134/// distributions).
135pub type NaryFn<'a> = &'a dyn Fn(&[Value], &mut Budget) -> OpResult<Value>;
136
137/// Apply `f` to a list of values, lifting over any that are distributions.
138pub fn lift_n(args: &[Value], budget: &mut Budget, f: NaryFn) -> OpResult<Value> {
139    if !args.iter().any(|a| a.is_dist() || matches!(a, Value::Counts(_))) {
140        return f(args, budget);
141    }
142    let listed_args;
143    let args = if args.iter().any(|a| matches!(a, Value::Counts(_))) {
144        listed_args = args
145            .iter()
146            .map(|a| Ok(listed(a, budget)?.into_owned()))
147            .collect::<OpResult<Vec<_>>>()?;
148        &listed_args[..]
149    } else {
150        args
151    };
152    let size = args
153        .iter()
154        .fold(1u128, |acc, a| acc.saturating_mul(outcomes(a).len() as u128));
155    budget.outcomes(size)?;
156    budget.work(size as u64)?;
157    let missing = args.iter().fold(0.0, |m, a| m + (1.0 - m) * missing(a));
158    let mut results = Vec::new();
159    let mut current = Vec::with_capacity(args.len());
160    product(args, 0, &mut current, 1.0, f, budget, &mut results)?;
161    combine(results, missing, budget)
162}
163
164fn product(
165    args: &[Value],
166    i: usize,
167    current: &mut Vec<Value>,
168    weight: f64,
169    f: NaryFn,
170    budget: &mut Budget,
171    results: &mut Vec<(Value, f64)>,
172) -> OpResult<()> {
173    if i == args.len() {
174        results.push((f(current, budget)?, weight));
175        return Ok(());
176    }
177    for (v, w) in outcomes(&args[i]).iter() {
178        current.push(v.clone());
179        product(args, i + 1, current, weight * w, f, budget, results)?;
180        current.pop();
181    }
182    Ok(())
183}
184
185// ── Probabilities, conditions and facts ──────────────────────────────────
186
187pub fn article(kind: &str) -> String {
188    let vowel = kind.starts_with(['a', 'e', 'i', 'o', 'u']);
189    format!("{} {kind}", if vowel { "an" } else { "a" })
190}
191
192/// A value used as a probability (chance weights, `bernoulli`, `binomial`…):
193/// a `prob` or a number checked at this boundary. Booleans need `prob(b)`.
194pub fn to_prob(v: &Value) -> OpResult<f64> {
195    if crate::analytic::contains(v) {
196        return Err(crate::analytic::unsupported(
197            "using an analytic outcome as a probability parameter",
198        ));
199    }
200    match v {
201        // A parameter out of range is outside what its function is defined for.
202        Value::Prob(_) | Value::Float(_) | Value::Int(_) => {
203            match make_prob(v).map_err(|e| e.as_fault(Fault::DomainError))? {
204                Value::Prob(p) => Ok(p),
205                _ => unreachable!("make_prob returns a probability"),
206            }
207        }
208        Value::Bool(_) => Err(OpError::new("expected a probability, found a fact (true or false)")
209            .help("convert a boolean explicitly with `prob(fact)`")),
210        Value::Dist(_) => Err(OpError::new(format!("expected a probability, found a {}", v.kind()))
211            .help("`P(…)` gives the probability that a distribution of facts is true")),
212        other => Err(OpError::new(format!(
213            "expected a probability, found {}",
214            article(&other.kind())
215        ))),
216    }
217}
218
219/// Explicit, checked construction. This never draws, lifts or clamps.
220pub fn make_prob(v: &Value) -> OpResult<Value> {
221    let p = match v {
222        Value::Bool(b) => {
223            if *b {
224                1.0
225            } else {
226                0.0
227            }
228        }
229        Value::Prob(p) | Value::Float(p) => *p,
230        Value::Int(n) if *n == 0 => 0.0,
231        Value::Int(n) if *n == 1 => 1.0,
232        Value::Int(_) => {
233            return Err(OpError::fault(
234                Fault::ConversionError,
235                "prob needs a finite number between 0 and 1",
236            ));
237        }
238        _ => {
239            return Err(
240                OpError::new(format!("prob needs a number or bool, found {}", article(&v.kind())))
241                    .help("draw distribution outcomes explicitly; use `P(d)` to query a boolean distribution"),
242            );
243        }
244    };
245    if !p.is_finite() || !(0.0..=1.0).contains(&p) {
246        return Err(OpError::fault(
247            Fault::ConversionError,
248            "prob needs a finite number between 0 and 1",
249        ));
250    }
251    Ok(Value::Prob(p))
252}
253
254/// Computed chances may stray a few ulps outside [0, 1] through rounding.
255/// This correction is only for algorithms known to produce probabilities;
256/// user inputs always go through the strict `make_prob` boundary.
257pub(crate) fn computed_prob(p: f64, what: &str) -> OpResult<Value> {
258    const ROUNDING: f64 = 8.0 * f64::EPSILON;
259    if !p.is_finite() || !(-ROUNDING..=1.0 + ROUNDING).contains(&p) {
260        return Err(OpError::new(format!(
261            "`{what}` could not compute a finite probability between 0 and 1"
262        )));
263    }
264    Ok(Value::Prob(if p <= 0.0 { 0.0 } else { p.min(1.0) }))
265}
266
267pub fn fact(v: &Value, context: &str) -> OpResult<bool> {
268    match v {
269        Value::Bool(b) => Ok(*b),
270        _ => Err(
271            OpError::new(format!("`{context}` needs a bool, found {}", article(&v.kind())))
272                .help("draw an outcome first with `~`; for a probability, draw from `bernoulli(p)`"),
273        ),
274    }
275}
276
277/// How likely a condition is to hold (docs/semantics.md, section 3).
278#[derive(Clone, Copy, Debug, PartialEq)]
279pub struct Condition {
280    /// The probability that it's true.
281    pub yes: f64,
282    /// The probability that it's false.
283    pub no: f64,
284    /// The probability missing from a distribution of facts: it could be
285    /// either. The three add up to one, up to rounding.
286    pub missing: f64,
287}
288
289/// A fact, or a distribution of facts.
290pub enum Truth {
291    Fact(bool),
292    Analytic(Arc<crate::analytic::Event>),
293    Probability(f64),
294    Uncertain(Arc<Dist>),
295}
296
297/// The operand of `and`, `or` or `not` (docs/semantics.md, section 2).
298pub fn truth(v: &Value, op: &str) -> OpResult<Truth> {
299    match v {
300        Value::Bool(b) => Ok(Truth::Fact(*b)),
301        Value::Event(e) => Ok(Truth::Analytic(e.clone())),
302        Value::Dist(d) if d.truth().is_some() => Ok(Truth::Uncertain(d.clone())),
303        Value::Prob(_) => Ok(Truth::Probability(to_prob(v)?)),
304        other => Err(OpError::new(format!(
305            "`{op}` needs bool, prob or dist[bool], found {}",
306            article(&other.kind())
307        ))),
308    }
309}
310
311pub fn not(v: &Value, budget: &mut Budget) -> OpResult<Value> {
312    match truth(v, "not")? {
313        Truth::Fact(b) => Ok(Value::Bool(!b)),
314        Truth::Analytic(e) => {
315            budget.work(e.size() as u64)?;
316            Ok(e.complement()?.value())
317        }
318        Truth::Probability(p) => computed_prob(1.0 - p, "not"),
319        Truth::Uncertain(d) => lift1(&Value::Dist(d), budget, |x, _| match x {
320            Value::Bool(b) => Ok(Value::Bool(!b)),
321            _ => unreachable!("checked by `truth`"),
322        }),
323    }
324}
325
326/// Compose independent boolean recipes. Bound facts retain their identity.
327pub fn logic(and: bool, a: Truth, b: Truth, budget: &mut Budget) -> OpResult<Value> {
328    if matches!(a, Truth::Analytic(_)) || matches!(b, Truth::Analytic(_)) {
329        let value = |t| match t {
330            Truth::Fact(b) => Value::Bool(b),
331            Truth::Probability(p) => Value::Prob(p),
332            Truth::Uncertain(d) => Value::Dist(d),
333            Truth::Analytic(e) => Value::Event(e),
334        };
335        let (a, b) = (value(a), value(b));
336        for v in [&a, &b] {
337            if let Value::Event(e) = v {
338                budget.collection(e.size() as u128)?;
339                budget.work(e.size() as u64)?;
340            }
341        }
342        return crate::analytic::logic(and, &a, &b);
343    }
344    let op = |x: bool, y: bool| if and { x && y } else { x || y };
345    match (a, b) {
346        (Truth::Fact(x), Truth::Fact(y)) => Ok(Value::Bool(op(x, y))),
347        (Truth::Probability(p), Truth::Probability(q)) => computed_prob(
348            if and { p * q } else { p + (1.0 - p) * q },
349            if and { "and" } else { "or" },
350        ),
351        (Truth::Probability(p), Truth::Fact(b)) | (Truth::Fact(b), Truth::Probability(p)) => {
352            Ok(Value::Prob(if b == and {
353                p
354            } else if b {
355                1.0
356            } else {
357                0.0
358            }))
359        }
360        (Truth::Fact(x), Truth::Uncertain(d)) | (Truth::Uncertain(d), Truth::Fact(x)) => {
361            lift1(&Value::Dist(d), budget, |v, _| match v {
362                Value::Bool(y) => Ok(Value::Bool(op(x, *y))),
363                _ => unreachable!("checked by `truth`"),
364            })
365        }
366        (a, b) => {
367            let law = |v| match v {
368                Truth::Uncertain(d) => Value::Dist(d),
369                Truth::Probability(p) => Dist::bernoulli(p).into_value(),
370                Truth::Fact(b) => Value::Bool(b),
371                Truth::Analytic(_) => unreachable!("handled above"),
372            };
373            lift2(&law(a), &law(b), budget, |x, y, _| match (x, y) {
374                (Value::Bool(x), Value::Bool(y)) => Ok(Value::Bool(op(*x, *y))),
375                _ => unreachable!("checked by `truth`"),
376            })
377        }
378    }
379}
380
381/// The boolean law requested by an observed anonymous draw.
382pub fn boolean_law(v: &Value) -> OpResult<Value> {
383    match truth(v, "observe ~")? {
384        Truth::Analytic(e) => Ok(Value::Event(e)),
385        Truth::Fact(b) => Ok(Dist::bernoulli(if b { 1.0 } else { 0.0 }).into_value()),
386        Truth::Probability(p) => Ok(Dist::bernoulli(p).into_value()),
387        Truth::Uncertain(d) => Ok(Value::Dist(d)),
388    }
389}
390
391pub fn condition(v: &Value) -> OpResult<Condition> {
392    if matches!(v, Value::Analytic(_)) {
393        return Err(crate::analytic::unsupported(
394            "using a continuous outcome as a probability condition",
395        ));
396    }
397    if matches!(v, Value::Float(_) | Value::Int(_)) {
398        let p = to_prob(v)?;
399        return Ok(Condition {
400            yes: p,
401            no: 1.0 - p,
402            missing: 0.0,
403        });
404    }
405    match truth(v, "condition")? {
406        Truth::Analytic(e) => Ok(Condition {
407            yes: e.probability(),
408            no: 1.0 - e.probability(),
409            missing: 0.0,
410        }),
411        Truth::Fact(b) => Ok(Condition {
412            yes: if b { 1.0 } else { 0.0 },
413            no: if b { 0.0 } else { 1.0 },
414            missing: 0.0,
415        }),
416        Truth::Probability(p) => Ok(Condition {
417            yes: p,
418            no: 1.0 - p,
419            missing: 0.0,
420        }),
421        Truth::Uncertain(d) => {
422            let (yes, no) = d.truth().expect("checked by truth");
423            Ok(Condition {
424                yes,
425                no,
426                missing: d.missing,
427            })
428        }
429    }
430}
431
432// ── Operators ────────────────────────────────────────────────────────────
433
434pub fn unary(op: UnOp, v: &Value, budget: &mut Budget) -> OpResult<Value> {
435    match op {
436        UnOp::Neg => lift1(v, budget, |x, budget| match x {
437            Value::Int(i) => {
438                budget.integer_work(i, &Integer::ZERO, false)?;
439                {
440                    let n = i.negated();
441                    budget.integer_allocation(n.bits(), 1)?;
442                    Ok(Value::Int(n))
443                }
444            }
445            Value::Float(f) | Value::Prob(f) => Ok(Value::Float(-f)),
446            Value::Complex(z) => Ok(Value::Complex(z.negated())),
447            Value::Analytic(a) => crate::analytic::negate(a),
448            other => Err(OpError::new(format!("can't negate {}", article(&other.kind())))),
449        }),
450        UnOp::Not => not(v, budget),
451        UnOp::Typeof => unreachable!("typeof is lowered to an interpreter intrinsic"),
452    }
453}
454
455pub fn binary(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
456    match op {
457        BinOp::And | BinOp::Or => unreachable!("`and` and `or` are evaluated lazily by the interpreter"),
458        BinOp::To => unreachable!("`to` is lowered to a built-in"),
459        BinOp::Range | BinOp::RangeExcl => range(op, a, b, budget),
460        BinOp::In => lift2(a, b, budget, |x, coll, budget| {
461            contains(coll, x, budget).map(Value::Bool)
462        }),
463        BinOp::NotIn => lift2(a, b, budget, |x, coll, budget| {
464            contains(coll, x, budget).map(|c| Value::Bool(!c))
465        }),
466        _ => lift2(a, b, budget, |x, y, budget| binary_plain(op, x, y, budget)),
467    }
468}
469
470fn binary_plain(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
471    // Outcomes of draws and continuous distributions are rare: one test
472    // each, for all of them.
473    let unusual = |v: &Value| matches!(v, Value::Analytic(_) | Value::Continuous(_) | Value::Joint(_));
474    if unusual(a) || unusual(b) {
475        return unusual_binary(op, a, b, budget);
476    }
477    if is_comparison(op) && (crate::analytic::contains(a) || crate::analytic::contains(b)) {
478        return analytic_comparison(op, a, b, budget);
479    }
480    if matches!(
481        op,
482        BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge
483    ) {
484        for v in [a, b] {
485            if let Value::Int(n) = v {
486                budget.integer_work(n, &Integer::ONE, false)?;
487            }
488            if let Value::Str(s) = v {
489                budget.string_work(s)?;
490            }
491        }
492    }
493    match op {
494        BinOp::Eq => Ok(Value::Bool(equals(a, b))),
495        BinOp::Ne => Ok(Value::Bool(!equals(a, b))),
496        BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge => {
497            let ord = compare(a, b)?;
498            Ok(Value::Bool(match op {
499                BinOp::Lt => ord.is_lt(),
500                BinOp::Le => ord.is_le(),
501                BinOp::Gt => ord.is_gt(),
502                _ => ord.is_ge(),
503            }))
504        }
505        BinOp::Add => add(a, b, budget),
506        BinOp::Sub => sub(a, b, budget),
507        _ => arith(op, a, b, budget),
508    }
509}
510
511fn add(a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
512    match (a, b) {
513        (Value::Str(x), Value::Str(y)) => {
514            let size = x
515                .len()
516                .checked_add(y.len())
517                .ok_or_else(|| OpError::limit("string size overflow"))?;
518            budget.string_size(size)?;
519            let mut out = String::new();
520            crate::text::push(&mut out, x, budget)?;
521            crate::text::push(&mut out, y, budget)?;
522            Ok(Value::str(&out))
523        }
524        (Value::List(x), Value::List(y)) => {
525            budget.collection(x.len() as u128 + y.len() as u128)?;
526            budget.work(x.len() as u64 + y.len() as u64)?;
527            let mut items = x.to_vec();
528            items.extend(y.iter().cloned());
529            Ok(Value::list(items))
530        }
531        (Value::Date(d), n @ (Value::Int(_) | Value::Float(_)))
532        | (n @ (Value::Int(_) | Value::Float(_)), Value::Date(d)) => {
533            let n = integer(n, "date offset", budget)?;
534            date_plus(
535                *d,
536                n.to_i64()
537                    .ok_or_else(|| OpError::fault(Fault::NumericOverflow, "date out of range"))?,
538            )
539        }
540        (Value::Str(_), _) | (_, Value::Str(_)) => {
541            Err(
542                OpError::new(format!("can't add {} and {}", article(&a.kind()), article(&b.kind())))
543                    .help("to build text, use interpolation: \"total: {x}\""),
544            )
545        }
546        _ => arith(BinOp::Add, a, b, budget),
547    }
548}
549
550fn sub(a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
551    match (a, b) {
552        (Value::Date(x), Value::Date(y)) => Ok(Value::Int((*x as i64 - *y as i64).into())),
553        (Value::Date(d), n @ (Value::Int(_) | Value::Float(_))) => {
554            match integer(n, "date offset", budget)?.negated().to_i64() {
555                Some(m) => date_plus(*d, m),
556                None => Err(OpError::fault(Fault::NumericOverflow, "date out of range")),
557            }
558        }
559        _ => arith(BinOp::Sub, a, b, budget),
560    }
561}
562
563fn date_plus(d: i32, n: i64) -> OpResult<Value> {
564    crate::dates::add_days(d, n)
565        .map(Value::Date)
566        .ok_or_else(|| OpError::fault(Fault::NumericOverflow, "date out of range"))
567}
568
569/// Numbers for arithmetic: ints, floats and probabilities (not facts).
570fn number(v: &Value) -> Option<f64> {
571    match v {
572        Value::Bool(_) => None,
573        _ => v.as_f64(),
574    }
575}
576
577fn arith(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
578    let bad = || {
579        OpError::new(format!(
580            "can't use `{}` with {} and {}",
581            op.symbol(),
582            article(&a.kind()),
583            article(&b.kind())
584        ))
585    };
586    if let (BinOp::Pow, Value::Complex(z), Value::Int(n)) = (op, a, b) {
587        budget.integer_work(n, &Integer::ONE, false)?;
588        budget.work(n.bits())?;
589        return Ok(Value::Complex(z.pow_integer(n)?));
590    }
591    if matches!(a, Value::Complex(_)) || matches!(b, Value::Complex(_)) {
592        let (Some(x), Some(y)) = (a.as_complex(), b.as_complex()) else {
593            return Err(bad());
594        };
595        let z = match op {
596            BinOp::Add => x.plus(y),
597            BinOp::Sub => x.minus(y),
598            BinOp::Mul => x.times(y),
599            BinOp::Div => x.divided_by(y),
600            BinOp::Pow => match b {
601                Value::Int(n) => {
602                    budget.integer_work(n, &Integer::ONE, false)?;
603                    budget.work(n.bits())?;
604                    x.pow_integer(n)
605                }
606                _ => Err(OpError::new("a complex power needs an int exponent")),
607            },
608            _ => Err(bad()),
609        }?;
610        return Ok(Value::Complex(z));
611    }
612    if let (Value::Int(x), Value::Int(y)) = (a, b) {
613        budget.integer_work(x, y, matches!(op, BinOp::Mul | BinOp::Div | BinOp::IntDiv | BinOp::Mod))?;
614        let n = match op {
615            BinOp::Add => x.add(y)?,
616            BinOp::Sub => x.sub(y)?,
617            BinOp::Mul => {
618                if !x.is_zero() && !y.is_zero() {
619                    budget.integer_bits((x.bits() + y.bits()).saturating_sub(1))?;
620                }
621                x.mul(y)?
622            }
623            BinOp::Div => {
624                if y.is_zero() {
625                    return Err(division_by_zero());
626                }
627                return x.ratio(y).map(Value::Float).ok_or_else(|| {
628                    OpError::fault(
629                        Fault::NumericOverflow,
630                        "division result is too large for a finite float",
631                    )
632                });
633            }
634            BinOp::IntDiv => x.div_mod(y)?.0,
635            BinOp::Mod => x.div_mod(y)?.1,
636            BinOp::Pow => return int_power(x, y, budget),
637            _ => return Err(bad()),
638        };
639        budget.integer_allocation(n.bits(), 1)?;
640        return Ok(Value::Int(n));
641    }
642    let (Some(x), Some(y)) = (number(a), number(b)) else {
643        return Err(bad());
644    };
645    let v = match op {
646        BinOp::Add => x + y,
647        BinOp::Sub => x - y,
648        BinOp::Mul => x * y,
649        BinOp::Div => {
650            if y == 0.0 {
651                return Err(division_by_zero());
652            }
653            x / y
654        }
655        BinOp::IntDiv => {
656            if y == 0.0 {
657                return Err(division_by_zero());
658            }
659            let q = (x / y).floor();
660            let n = Integer::from_f64(q)
661                .ok_or_else(|| OpError::fault(Fault::NumericOverflow, "integer division result is not finite"))?;
662            budget.integer_allocation(n.bits(), 1)?;
663            return Ok(Value::Int(n));
664        }
665        BinOp::Mod => {
666            if y == 0.0 {
667                return Err(division_by_zero());
668            }
669            x - (x / y).floor() * y
670        }
671        BinOp::Pow => libm::pow(x, y),
672        _ => return Err(bad()),
673    };
674    if !v.is_finite() {
675        return Err(OpError::fault(
676            Fault::NumericOverflow,
677            format!("`{}` gave a result that isn't a finite number", op.symbol()),
678        ));
679    }
680    Ok(Value::Float(v))
681}
682
683/// A continuous distribution compared with a number: a distribution of
684/// facts, from its CDF. Nothing else is defined on one yet
685/// (docs/semantics.md, section 13).
686fn is_comparison(op: BinOp) -> bool {
687    matches!(
688        op,
689        BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge
690    )
691}
692
693/// `op` with an outcome of a draw, or a continuous distribution, or a
694/// distribution whose outcomes own draws.
695#[inline(never)]
696fn unusual_binary(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
697    if matches!(a, Value::Analytic(_)) || matches!(b, Value::Analytic(_)) {
698        return crate::analytic::binary(op, a, b, budget);
699    }
700    if is_comparison(op) && (crate::analytic::contains(a) || crate::analytic::contains(b)) {
701        return analytic_comparison(op, a, b, budget);
702    }
703    match continuous_binary(op, a, b, budget)? {
704        Some(v) => Ok(v),
705        None => unreachable!("a continuous operand"),
706    }
707}
708
709/// A comparison with a value that holds outcomes of draws: of an event with
710/// a fact or another event.
711#[inline(never)]
712fn analytic_comparison(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
713    if matches!(op, BinOp::Eq | BinOp::Ne) {
714        if let (Value::Event(x), Value::Bool(b)) | (Value::Bool(b), Value::Event(x)) = (a, b) {
715            return if *b == (op == BinOp::Eq) {
716                Ok(Value::Event(x.clone()))
717            } else {
718                not(&Value::Event(x.clone()), budget)
719            };
720        }
721    }
722    if let (Value::Event(x), Value::Event(y)) = (a, b) {
723        if matches!(op, BinOp::Eq | BinOp::Ne) {
724            budget.work((x.size() + y.size()) as u64)?;
725            let equal = x.combine(y, BinOp::Eq)?;
726            return Ok(if op == BinOp::Eq { equal } else { equal.complement()? }.value());
727        }
728    }
729    Err(crate::analytic::unsupported(
730        "comparing aggregate or boolean analytic outcomes",
731    ))
732}
733
734fn continuous_binary(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Option<Value>> {
735    let (family, other, flipped) = match (a, b) {
736        (Value::Continuous(f), other) => (f, other, false),
737        (other, Value::Continuous(f)) => (f, other, true),
738        (Value::Joint(_), _) | (_, Value::Joint(_)) => return crate::joint::binary(op, a, b, budget).map(Some),
739        _ => return Ok(None),
740    };
741    let needs_value = || {
742        OpError::new(format!(
743            "`{}` needs a value, not a {} distribution",
744            op.symbol(),
745            family.name()
746        ))
747        .help("draw a value first, like `let x ~ normal(0, 1)`; comparing a distribution with a number works too")
748    };
749    let x = match number(other) {
750        Some(x) if !x.is_nan() => x,
751        _ => return Err(needs_value()),
752    };
753    // For a continuous X, P(X < x) = P(X ≤ x) = F(x).
754    let below = family.cdf(x);
755    let yes = match (op, flipped) {
756        (BinOp::Lt | BinOp::Le, false) | (BinOp::Gt | BinOp::Ge, true) => below,
757        (BinOp::Gt | BinOp::Ge, false) | (BinOp::Lt | BinOp::Le, true) => 1.0 - below,
758        (BinOp::Eq, _) => 0.0,
759        (BinOp::Ne, _) => 1.0,
760        // Arithmetic would make a new recipe, which isn't built yet.
761        _ => {
762            return Err(OpError::unsupported(format!(
763                "`{}` on an undrawn {} distribution isn't supported yet",
764                op.symbol(),
765                family.name()
766            ))
767            .help("draw a value first, like `let x ~ normal(0, 1)`, and compute with it"));
768        }
769    };
770    Ok(Some(Dist::bernoulli(yes).into_value()))
771}
772
773fn division_by_zero() -> OpError {
774    OpError::fault(Fault::DivisionByZero, "division by zero")
775}
776
777fn int_power(base: &Integer, exponent: &Integer, budget: &mut Budget) -> OpResult<Value> {
778    if exponent.is_zero() {
779        return Ok(Value::Int(Integer::ONE));
780    }
781    if base.is_zero() {
782        return if exponent.is_negative() {
783            Err(division_by_zero())
784        } else {
785            Ok(Value::Int(Integer::ZERO))
786        };
787    }
788    if exponent.is_negative() {
789        let reciprocal = Integer::ONE.ratio(&base.abs()).unwrap_or(0.0);
790        let power = exponent.abs().to_f64().unwrap_or(f64::INFINITY);
791        let result = libm::pow(reciprocal, power);
792        return Ok(Value::Float(if base.is_negative() && exponent.is_odd() {
793            -result
794        } else {
795            result
796        }));
797    }
798    if *base == 1 || *base == -1 {
799        return Ok(Value::Int(if *base == -1 && exponent.is_odd() {
800            (-1).into()
801        } else {
802            Integer::ONE
803        }));
804    }
805    let mut n = exponent
806        .to_u64()
807        .ok_or_else(|| OpError::limit("integer power exceeds the integer size limit"))?;
808    budget.integer_bits(base.bits().saturating_sub(1).saturating_mul(n).saturating_add(1))?;
809    let mut result = Integer::ONE;
810    let mut b = base.clone();
811    while n != 0 {
812        if n & 1 != 0 {
813            budget.integer_work(&result, &b, true)?;
814            result = result.mul(&b)?;
815            budget.integer_bits(result.bits())?;
816        }
817        n >>= 1;
818        if n != 0 {
819            budget.integer_work(&b, &b, true)?;
820            b = b.mul(&b)?;
821            budget.integer_bits(b.bits())?;
822        }
823    }
824    budget.integer_allocation(result.bits(), 1)?;
825    Ok(Value::Int(result))
826}
827
828fn numeric_compare(a: &Value, b: &Value) -> Option<std::cmp::Ordering> {
829    match (a, b) {
830        (Value::Int(a), Value::Int(b)) => Some(a.cmp(b)),
831        (Value::Int(a), Value::Float(b) | Value::Prob(b)) => a.cmp_f64(*b),
832        (Value::Float(a) | Value::Prob(a), Value::Int(b)) => b.cmp_f64(*a).map(std::cmp::Ordering::reverse),
833        (Value::Float(a) | Value::Prob(a), Value::Float(b) | Value::Prob(b)) => a.partial_cmp(b),
834        _ => None,
835    }
836}
837
838/// The language's `==`: numbers compare by value across int, float and prob.
839pub fn equals(a: &Value, b: &Value) -> bool {
840    if let (Value::Complex(x), Value::Complex(y)) = (a, b) {
841        return x == y;
842    }
843    if let (Value::Complex(z), real) | (real, Value::Complex(z)) = (a, b) {
844        if z.im() != 0.0 {
845            return false;
846        }
847        return match real {
848            // Compare integers without rounding them through f64.
849            Value::Int(n) => {
850                let x = z.re();
851                n.cmp_f64(x).is_some_and(|c| c.is_eq())
852            }
853            Value::Float(x) | Value::Prob(x) => z.re() == *x,
854            _ => false,
855        };
856    }
857    if let Some(c) = numeric_compare(a, b) {
858        return c.is_eq();
859    }
860    if matches!(a, Value::Float(x) | Value::Prob(x) if x.is_nan())
861        || matches!(b, Value::Float(x) | Value::Prob(x) if x.is_nan())
862    {
863        return false;
864    }
865    a == b
866}
867
868/// Ordering for `<`, `>` and friends.
869pub fn compare(a: &Value, b: &Value) -> OpResult<std::cmp::Ordering> {
870    if matches!(a, Value::Complex(_)) || matches!(b, Value::Complex(_)) {
871        return Err(OpError::new("complex values have no ordering").help("compare `abs(z)`, `real(z)` or `imag(z)`"));
872    }
873    if let Some(c) = numeric_compare(a, b) {
874        return Ok(c);
875    }
876    if matches!(a, Value::Float(x) | Value::Prob(x) if x.is_nan())
877        || matches!(b, Value::Float(x) | Value::Prob(x) if x.is_nan())
878    {
879        return Err(OpError::new("can't order NaN"));
880    }
881    match (a, b) {
882        (Value::Str(x), Value::Str(y)) => Ok(x.cmp(y)),
883        (Value::Date(x), Value::Date(y)) => Ok(x.cmp(y)),
884        (Value::Enum(x), Value::Enum(y)) if x.ty == y.ty => Ok(x.variant.cmp(&y.variant)),
885        (Value::List(x), Value::List(y)) => {
886            for (p, q) in x.iter().zip(y.iter()) {
887                let c = compare(p, q)?;
888                if c.is_ne() {
889                    return Ok(c);
890                }
891            }
892            Ok(x.len().cmp(&y.len()))
893        }
894        _ => Err(OpError::new(format!(
895            "can't compare {} with {}",
896            article(&a.kind()),
897            article(&b.kind())
898        ))),
899    }
900}
901
902/// Whether `item` is in `coll` (for `in`).
903pub fn contains(coll: &Value, item: &Value, budget: &mut Budget) -> OpResult<bool> {
904    if crate::analytic::contains(coll) || crate::analytic::contains(item) {
905        return Err(crate::analytic::unsupported("membership involving analytic outcomes"));
906    }
907    if let Value::Str(s) = coll {
908        budget.string_work(s)?;
909        if let Value::Str(s) = item {
910            budget.string_work(s)?;
911        }
912    }
913    Ok(match coll {
914        Value::List(items) => items.iter().any(|x| equals(x, item)),
915        Value::Map(m) => m.contains_key(item),
916        Value::Bag(b) => b.get(item).is_some_and(|n| *n > 0),
917        Value::Range(lo, hi) => match item {
918            Value::Int(n) => n >= lo && n <= hi,
919            Value::Float(x) | Value::Prob(x) if x.is_finite() && x.fract() == 0.0 => {
920                lo.cmp_f64(*x).is_some_and(|c| !c.is_gt()) && hi.cmp_f64(*x).is_some_and(|c| !c.is_lt())
921            }
922            _ => false,
923        },
924        Value::Str(s) => match item {
925            Value::Str(sub) => s.contains(&**sub),
926            _ => {
927                return Err(OpError::new(format!(
928                    "can't look for {} in a string",
929                    article(&item.kind())
930                )));
931            }
932        },
933        other => return Err(OpError::new(format!("can't look inside {}", article(&other.kind())))),
934    })
935}
936
937fn range(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
938    if a.is_uncertain() || b.is_uncertain() {
939        return Err(OpError::new("a range needs plain whole numbers, not distributions")
940            .help("draw a value first, like `let n ~ d6`"));
941    }
942    let lo = integer(a, "range bound", budget)?;
943    let hi = integer(b, "range bound", budget)?;
944    let hi = if op == BinOp::RangeExcl {
945        hi.sub(&Integer::ONE)?
946    } else {
947        hi.into_owned()
948    };
949    budget.integer_work(&lo, &hi, false)?;
950    if op == BinOp::RangeExcl {
951        budget.integer_allocation(hi.bits(), 1)?;
952    }
953    Ok(Value::Range(lo.into_owned(), hi))
954}
955
956/// The number of integers in `lo..=hi`, without overflowing.
957pub fn range_len(lo: &Integer, hi: &Integer) -> OpResult<Integer> {
958    if hi < lo {
959        Ok(Integer::ZERO)
960    } else {
961        Ok(hi.sub(lo)?.add(&Integer::ONE)?)
962    }
963}
964
965/// Bounded materialization count, never truncating a large length.
966pub fn range_count(lo: &Integer, hi: &Integer) -> OpResult<u128> {
967    range_len(lo, hi)?
968        .to_u128()
969        .ok_or_else(|| OpError::limit("range has too many elements to materialize"))
970}
971
972// ── Collections ──────────────────────────────────────────────────────────
973
974pub fn field(v: &Value, name: &str, budget: &mut Budget) -> OpResult<Value> {
975    lift1(v, budget, |x, _| match x {
976        Value::Record(r) => r.get(name).cloned().ok_or_else(|| {
977            let known: Vec<String> = r.fields.iter().map(|(n, _)| format!("`{n}`")).collect();
978            OpError::new(format!("{} has no field `{name}`", article(&x.kind())))
979                .help(format!("its fields are {}", known.join(", ")))
980        }),
981        other => Err(OpError::new(format!(
982            "can't read the field `{name}` of {}",
983            article(&other.kind())
984        ))),
985    })
986}
987
988pub fn index(coll: &Value, i: &Value, budget: &mut Budget) -> OpResult<Value> {
989    lift2(coll, i, budget, |c, i, budget| {
990        let v = index_plain(c, i, budget)?;
991        if let (Value::Range(lo, hi), Value::Int(n)) = (c, &v) {
992            budget.integer_work(lo, hi, false)?;
993            budget.integer_allocation(n.bits(), 1)?;
994        }
995        Ok(v)
996    })
997}
998
999pub fn index_plain(coll: &Value, i: &Value, budget: &mut Budget) -> OpResult<Value> {
1000    match coll {
1001        Value::List(items) => {
1002            let k = as_index(i, items.len() as u128, budget)?;
1003            Ok(items[k as usize].clone())
1004        }
1005        Value::Range(lo, hi) => {
1006            let k = integer(i, "index", budget)?;
1007            if k.is_negative() {
1008                return Err(OpError::fault(Fault::IndexOutOfBounds, "index out of range"));
1009            }
1010            let n = lo.add(&k)?;
1011            if &n > hi {
1012                return Err(OpError::fault(Fault::IndexOutOfBounds, "index out of range"));
1013            }
1014            Ok(Value::Int(n))
1015        }
1016        Value::Str(s) => {
1017            budget.string_work(s)?;
1018            let k = as_index(i, s.chars().count() as u128, budget)?;
1019            let c = s.chars().nth(k as usize).expect("checked scalar index");
1020            crate::text::value(c.encode_utf8(&mut [0; 4]), budget)
1021        }
1022        Value::Map(m) => m.get(i).cloned().ok_or_else(|| {
1023            OpError::fault(Fault::MissingKey, format!("the key {i:?} isn't in the map"))
1024                .help("use `get(key, default)` for keys that may be missing")
1025        }),
1026        other => Err(OpError::new(format!("can't index {}", article(&other.kind())))),
1027    }
1028}
1029
1030pub fn as_index(i: &Value, len: u128, budget: &mut Budget) -> OpResult<u128> {
1031    let k = integer(i, "index", budget)?.to_u128();
1032    k.filter(|k| *k < len).ok_or_else(|| {
1033        OpError::fault(
1034            Fault::IndexOutOfBounds,
1035            format!("index {i} is out of range for a length of {len}"),
1036        )
1037        .help("indices start at 0")
1038    })
1039}
1040
1041pub fn make_record(ty: Option<Arc<str>>, mut fields: Vec<(Arc<str>, Value)>) -> Value {
1042    fields.sort_by(|a, b| a.0.cmp(&b.0));
1043    Value::record(Record { ty, fields })
1044}
1045
1046pub fn with_fields(base: &Value, updates: &[(Arc<str>, Value)]) -> OpResult<Value> {
1047    let Value::Record(r) = base else {
1048        return Err(OpError::new(format!(
1049            "`with` needs a record, found {}",
1050            article(&base.kind())
1051        )));
1052    };
1053    let mut r = Record::clone(r);
1054    for (name, v) in updates {
1055        match r.get_mut(name) {
1056            Some(slot) => *slot = v.clone(),
1057            None => {
1058                return Err(OpError::new(format!("{} has no field `{name}`", article(&base.kind()))));
1059            }
1060        }
1061    }
1062    Ok(Value::record(r))
1063}
1064
1065pub fn enum_value(ty: u32, variant: u32, name: &str) -> Value {
1066    Value::Enum(Arc::new(EnumValue {
1067        ty,
1068        variant,
1069        name: Arc::from(name),
1070    }))
1071}
1072
1073/// Is `v` certainly `truth`? (For `and`/`or` whose right side has statements.)
1074pub fn is_certain(v: &Value, truth: bool) -> bool {
1075    matches!(v, Value::Bool(b) if *b == truth)
1076}
1077
1078#[cfg(test)]
1079mod probability_tests {
1080    use super::*;
1081
1082    #[test]
1083    fn computed_probabilities_only_correct_boundary_roundoff() {
1084        assert_eq!(computed_prob(1.0 + f64::EPSILON, "test").unwrap(), Value::Prob(1.0));
1085        assert_eq!(computed_prob(-f64::EPSILON, "test").unwrap(), Value::Prob(0.0));
1086        assert_eq!(computed_prob(1e-300, "test").unwrap(), Value::Prob(1e-300));
1087        for p in [1.001, -0.001, f64::INFINITY, f64::NEG_INFINITY, f64::NAN] {
1088            assert!(computed_prob(p, "test").is_err());
1089        }
1090    }
1091}