Skip to main content

probl_engine/
builtins.rs

1//! Built-in functions that don't need the interpreter.
2
3use crate::complex::Complex;
4use crate::continuous::{Family, Mixture, Part};
5use crate::dates;
6use crate::dist::{Budget, Counts, Dist};
7use crate::error::{Fault, OpError, OpResult};
8use crate::ops::{self, article, as_index, integer, range_count, range_len, to_prob};
9use crate::value::{Value, fmt_float};
10use probl_number::Integer;
11use probl_sema::Builtin;
12use std::borrow::Cow;
13use std::collections::BTreeMap;
14use std::sync::Arc;
15
16/// Call a built-in on plain (non-distribution) arguments.
17pub fn call_plain(b: Builtin, args: &[Value], budget: &mut Budget) -> OpResult<Value> {
18    for arg in args {
19        if let Value::Str(s) = arg {
20            budget.string_work(s)?;
21        }
22        if let Value::Int(n) = arg {
23            if matches!(b, Builtin::BitLength | Builtin::ILog2) {
24                // The stored length and highest word suffice, even for a bigint.
25                budget.integer_bits(n.bits())?;
26                budget.work(1)?;
27            } else {
28                budget.integer_work(n, if b == Builtin::Str { n } else { &Integer::ONE }, b == Builtin::Str)?;
29            }
30        }
31    }
32    let value = call_plain_inner(b, args, budget)?;
33    if let Value::Int(n) = &value {
34        budget.integer_allocation(n.bits(), 1)?;
35    }
36    Ok(value)
37}
38
39fn call_plain_inner(b: Builtin, args: &[Value], budget: &mut Budget) -> OpResult<Value> {
40    use Builtin as B;
41    let a = |i: usize| &args[i];
42    match b {
43        B::Min | B::Max => min_max(args, b == B::Max, budget),
44        B::Abs => match a(0) {
45            Value::Complex(z) => finite_float(z.abs(), "abs"),
46            Value::Int(n) => Ok(Value::Int(n.abs())),
47            Value::Prob(p) => Ok(Value::Prob(p.abs())),
48            v => float1(v, "abs", |x| Some(x.abs())),
49        },
50        B::Floor => to_int(a(0), f64::floor),
51        B::Ceil => to_int(a(0), f64::ceil),
52        B::Trunc => to_int(a(0), libm::trunc),
53        B::Round => round(a(0), args.get(1), budget),
54        B::Sqrt => elementary1(a(0), "sqrt", Complex::sqrt, |x| (x >= 0.0).then(|| x.sqrt())),
55        B::Cbrt => elementary1(a(0), "cbrt", Complex::cbrt, |x| Some(libm::cbrt(x))),
56        B::Exp => elementary1(a(0), "exp", Complex::exp, |x| {
57            Some(crate::math::exp(x)).filter(|y| y.is_finite())
58        }),
59        B::Exp2 => elementary1(a(0), "exp2", Complex::exp2, |x| Some(libm::exp2(x))),
60        B::Ln => elementary1(a(0), "ln", Complex::ln, |x| (x > 0.0).then(|| libm::log(x))),
61        B::Log10 => elementary1(a(0), "log10", Complex::log10, |x| (x > 0.0).then(|| libm::log10(x))),
62        B::Log2 => elementary1(a(0), "log2", Complex::log2, |x| (x > 0.0).then(|| libm::log2(x))),
63        B::Log1p => elementary1(a(0), "log1p", Complex::log1p, |x| (x > -1.0).then(|| libm::log1p(x))),
64        B::Expm1 => elementary1(a(0), "expm1", Complex::expm1, |x| {
65            Some(libm::expm1(x)).filter(|y| y.is_finite())
66        }),
67        B::Sin => elementary1(a(0), "sin", Complex::sin, |x| Some(libm::sin(x))),
68        B::Cos => elementary1(a(0), "cos", Complex::cos, |x| Some(libm::cos(x))),
69        B::Tan => elementary1(a(0), "tan", Complex::tan, |x| Some(libm::tan(x))),
70        B::Asin => elementary1(a(0), "asin", Complex::asin, |x| {
71            (-1.0..=1.0).contains(&x).then(|| libm::asin(x))
72        }),
73        B::Acos => elementary1(a(0), "acos", Complex::acos, |x| {
74            (-1.0..=1.0).contains(&x).then(|| libm::acos(x))
75        }),
76        B::Atan => elementary1(a(0), "atan", Complex::atan, |x| Some(libm::atan(x))),
77        B::Atan2 => float2(a(0), a(1), "atan2", libm::atan2),
78        B::Hypot => float2(a(0), a(1), "hypot", libm::hypot),
79        B::Sinh => elementary1(a(0), "sinh", Complex::sinh, |x| {
80            Some(libm::sinh(x)).filter(|y| y.is_finite())
81        }),
82        B::Cosh => elementary1(a(0), "cosh", Complex::cosh, |x| {
83            Some(libm::cosh(x)).filter(|y| y.is_finite())
84        }),
85        B::Tanh => elementary1(a(0), "tanh", Complex::tanh, |x| Some(libm::tanh(x))),
86        B::Asinh => elementary1(a(0), "asinh", Complex::asinh, |x| Some(libm::asinh(x))),
87        B::Acosh => elementary1(a(0), "acosh", Complex::acosh, |x| (x >= 1.0).then(|| libm::acosh(x))),
88        B::Atanh => elementary1(a(0), "atanh", Complex::atanh, |x| {
89            (x.abs() < 1.0).then(|| libm::atanh(x))
90        }),
91        B::BitLength => Ok(Value::Int(integer(a(0), "bit_length", budget)?.bits().into())),
92        B::BitAnd => Ok(Value::Int(
93            integer(a(0), b.name(), budget)?.bit_and(integer(a(1), b.name(), budget)?.as_ref())?,
94        )),
95        B::BitOr => Ok(Value::Int(
96            integer(a(0), b.name(), budget)?.bit_or(integer(a(1), b.name(), budget)?.as_ref())?,
97        )),
98        B::BitXor => Ok(Value::Int(
99            integer(a(0), b.name(), budget)?.bit_xor(integer(a(1), b.name(), budget)?.as_ref())?,
100        )),
101        B::BitNot => Ok(Value::Int(integer(a(0), b.name(), budget)?.bit_not()?)),
102        B::BitCount => Ok(Value::Int(integer(a(0), b.name(), budget)?.bit_count().into())),
103        B::ILog2 => {
104            let n = integer(a(0), "ilog2", budget)?;
105            if n.is_zero() || n.is_negative() {
106                return Err(OpError::fault(Fault::DomainError, "`ilog2` needs a positive integer"));
107            }
108            Ok(Value::Int((n.bits() - 1).into()))
109        }
110        B::Choose => choose(
111            nonnegative_int(a(0), "choose", budget)?.as_ref(),
112            nonnegative_int(a(1), "choose", budget)?.as_ref(),
113            budget,
114        ),
115        B::Factorial => factorial(nonnegative_int(a(0), "factorial", budget)?.as_ref(), budget),
116        B::Gcd | B::Lcm => {
117            let (x, y) = (
118                integer(a(0), b.name(), budget)?.abs(),
119                integer(a(1), b.name(), budget)?.abs(),
120            );
121            let d = gcd(x.clone(), y.clone(), budget)?;
122            let result = if b == B::Gcd {
123                d
124            } else if x.is_zero() || y.is_zero() {
125                Integer::ZERO
126            } else {
127                let q = x.div_mod(&d)?.0;
128                budget.integer_work(&q, &y, true)?;
129                q.mul(&y)?
130            };
131            Ok(Value::Int(result))
132        }
133        B::EulerPhi => euler_phi(nonnegative_int(a(0), "euler_phi", budget)?.as_ref(), budget),
134        B::LnGamma => float1(a(0), "ln_gamma", |x| (x > 0.0).then(|| libm::lgamma(x))),
135        B::Erf => float1(a(0), "erf", |x| Some(libm::erf(x))),
136        B::Erfc => float1(a(0), "erfc", |x| Some(libm::erfc(x))),
137        B::Complex => {
138            let z = if args.len() == 1 {
139                complex_number(a(0), "complex")?
140            } else {
141                Complex::new(number(a(0), "complex")?, number(a(1), "complex")?)?
142            };
143            Ok(Value::Complex(z))
144        }
145        B::Real => Ok(Value::Float(complex_number(a(0), "real")?.re())),
146        B::Imag => Ok(Value::Float(complex_number(a(0), "imag")?.im())),
147        B::Conj => Ok(Value::Complex(complex_number(a(0), "conj")?.conjugate())),
148        B::Abs2 => finite_float(complex_number(a(0), "abs2")?.abs2(), "abs2"),
149        B::Arg => Ok(Value::Float(complex_number(a(0), "arg")?.arg())),
150        B::Cis => {
151            let theta = number(a(0), "cis")?;
152            Complex::new(libm::cos(theta), libm::sin(theta)).map(Value::Complex)
153        }
154        B::Clamp => {
155            let (lo, hi) = (a(1), a(2));
156            if ops::compare(lo, hi)?.is_gt() {
157                return Err(OpError::fault(
158                    Fault::DomainError,
159                    "clamp's lower bound is above its upper bound",
160                ));
161            }
162            if ops::compare(a(0), lo)?.is_lt() {
163                Ok(lo.clone())
164            } else if ops::compare(a(0), hi)?.is_gt() {
165                Ok(hi.clone())
166            } else {
167                Ok(a(0).clone())
168            }
169        }
170        B::Str => crate::text::formatted(a(0), budget),
171        B::Upper | B::Lower => crate::text::case(&text(a(0), b.name())?, b == B::Upper, budget),
172        B::Trim | B::TrimStart | B::TrimEnd => {
173            let s = text(a(0), b.name())?;
174            let set = args.get(1).map(|v| text(v, b.name())).transpose()?;
175            let set = set
176                .as_ref()
177                .map(|s| {
178                    let mut set = rustc_hash::FxHashSet::default();
179                    for c in s.chars() {
180                        if !set.contains(&c) {
181                            budget.collection(set.len() as u128 + 1)?;
182                            set.insert(c);
183                        }
184                    }
185                    Ok::<_, OpError>(set)
186                })
187                .transpose()?;
188            let matches = |c: char| set.as_ref().map_or_else(|| c.is_whitespace(), |set| set.contains(&c));
189            let trimmed = match b {
190                B::TrimStart => s.trim_start_matches(matches),
191                B::TrimEnd => s.trim_end_matches(matches),
192                _ => s.trim_matches(matches),
193            };
194            crate::text::value(trimmed, budget)
195        }
196        B::StartsWith | B::EndsWith => {
197            let (s, part) = (text(a(0), b.name())?, text(a(1), b.name())?);
198            Ok(Value::Bool(if b == B::StartsWith {
199                s.starts_with(&*part)
200            } else {
201                s.ends_with(&*part)
202            }))
203        }
204        B::Chars => crate::text::chars(&text(a(0), "chars")?, budget).map(Value::list),
205        B::Split => {
206            let (s, sep) = (text(a(0), "split")?, text(a(1), "split")?);
207            if sep.is_empty() {
208                return crate::text::chars(&s, budget).map(Value::list);
209            }
210            let mut parts = Vec::new();
211            for part in s.split(&*sep) {
212                budget.collection(parts.len() as u128 + 1)?;
213                parts.push(crate::text::value(part, budget)?);
214            }
215            Ok(Value::list(parts))
216        }
217        B::Join => {
218            let items = list(a(0), "join")?;
219            let sep = text(a(1), "join")?;
220            budget.collection(items.len() as u128)?;
221            budget.work(items.len() as u64)?;
222            let mut out = String::new();
223            for (i, item) in items.iter().enumerate() {
224                if i != 0 {
225                    crate::text::push(&mut out, &sep, budget)?;
226                }
227                crate::text::push_value(&mut out, item, budget)?;
228            }
229            Ok(Value::str(&out))
230        }
231        B::Len => len(a(0)),
232        B::Slice => slice(a(0), a(1), args.get(2), budget),
233        B::Sum => sum(a(0), budget),
234        B::Count if args.len() == 1 => len(a(0)),
235        B::Sort | B::SortDesc => {
236            assert_eq!(args.len(), 1, "comparator sorting is handled by the interpreter");
237            let mut items = items(a(0), b.name(), budget)?;
238            sort(&mut items, b == B::SortDesc, budget)?;
239            Ok(Value::list(items))
240        }
241        B::Reverse => match a(0) {
242            Value::Str(s) => {
243                budget.string_allocation(s.len())?;
244                Ok(Value::str(&s.chars().rev().collect::<String>()))
245            }
246            v => {
247                let mut items = items(v, "reverse", budget)?;
248                items.reverse();
249                Ok(Value::list(items))
250            }
251        },
252        B::Keys => match a(0) {
253            Value::Map(m) => Ok(Value::list(m.keys().cloned().collect())),
254            Value::Bag(m) => Ok(Value::list(m.keys().cloned().collect())),
255            v => Err(expected("a map", v, "keys")),
256        },
257        B::Values => match a(0) {
258            Value::Map(m) => Ok(Value::list(m.values().cloned().collect())),
259            v => Err(expected("a map", v, "values")),
260        },
261        B::Get => get(a(0), a(1), args.get(2), budget),
262        B::Contains => ops::contains(a(0), a(1), budget).map(Value::Bool),
263        B::Highest | B::Lowest => extremes(a(0), a(1), b == B::Highest, budget),
264        B::Enumerate => {
265            let items = items(a(0), "enumerate", budget)?;
266            Ok(Value::list(
267                items
268                    .into_iter()
269                    .enumerate()
270                    .map(|(i, x)| Value::list(vec![Value::Int((i as i64).into()), x]))
271                    .collect(),
272            ))
273        }
274        B::Zip => {
275            let (xs, ys) = (items(a(0), "zip", budget)?, items(a(1), "zip", budget)?);
276            Ok(Value::list(
277                xs.into_iter().zip(ys).map(|(x, y)| Value::list(vec![x, y])).collect(),
278            ))
279        }
280        B::Push => {
281            let items = list(a(0), "push")?;
282            budget.collection(items.len() as u128 + 1)?;
283            let mut items = items.to_vec();
284            items.push(a(1).clone());
285            Ok(Value::list(items))
286        }
287        B::Insert => insert(a(0), a(1), a(2), budget),
288        B::Remove => remove(a(0), a(1), budget),
289        B::Pop => Err(
290            OpError::new("`pop` changes a list, so call it as a method on a variable")
291                .help("write `let top = xs.pop()`"),
292        ),
293        B::Last => {
294            let items = list(a(0), "pop")?;
295            items
296                .last()
297                .cloned()
298                .ok_or_else(|| OpError::fault(Fault::EmptyCollection, "can't pop from an empty list"))
299        }
300        B::DropLast => {
301            let items = list(a(0), "pop")?;
302            let n = items.len().saturating_sub(1);
303            Ok(Value::list(items[..n].to_vec()))
304        }
305        B::Prob => ops::make_prob(a(0)),
306        B::BooleanLaw => ops::boolean_law(a(0)),
307        B::Bernoulli | B::ScoreLaw => Ok(Dist::bernoulli(to_prob(a(0))?).into_value()),
308        B::OneOf => one_of(a(0), budget),
309        B::Binomial | B::Poisson | B::Geometric => {
310            let counts = counts(b, args, budget)?.expect("called with plain arguments");
311            Ok(counts.list(budget)?.into_value())
312        }
313        B::Bag => bag(a(0), budget),
314        B::Normal => continuous(Family::normal(number(a(0), "normal")?, number(a(1), "normal")?)),
315        B::Lognormal => continuous(Family::lognormal(
316            number(a(0), "lognormal")?,
317            number(a(1), "lognormal")?,
318        )),
319        B::Uniform => continuous(Family::uniform(number(a(0), "uniform")?, number(a(1), "uniform")?)),
320        B::Beta => continuous(Family::beta(number(a(0), "beta")?, number(a(1), "beta")?)),
321        B::Gamma => continuous(Family::gamma(number(a(0), "gamma")?, number(a(1), "gamma")?)),
322        B::Exponential => continuous(Family::exponential(number(a(0), "exponential")?)),
323        B::Triangular => continuous(Family::triangular(
324            number(a(0), "triangular")?,
325            number(a(1), "triangular")?,
326            number(a(2), "triangular")?,
327        )),
328        B::Pert => continuous(Family::pert(
329            number(a(0), "pert")?,
330            number(a(1), "pert")?,
331            number(a(2), "pert")?,
332        )),
333        B::NormalRange => continuous(Family::normal_range(
334            number(a(0), "normal_range")?,
335            number(a(1), "normal_range")?,
336        )),
337        B::To => continuous(Family::estimate(number(a(0), "to")?, number(a(1), "to")?)),
338        B::Mixture | B::Truncate | B::Bins => {
339            Err(OpError::unsupported(format!("`{}` isn't implemented yet", b.name())))
340        }
341        B::Odds => {
342            let p = to_prob(a(0))?;
343            if p >= 1.0 {
344                return Err(OpError::fault(
345                    Fault::DomainError,
346                    "the odds of a certain event are infinite",
347                ));
348            }
349            Ok(Value::Float(p / (1.0 - p)))
350        }
351        B::Logit => {
352            let p = to_prob(a(0))?;
353            if p <= 0.0 || p >= 1.0 {
354                return Err(OpError::fault(
355                    Fault::DomainError,
356                    "logit needs a probability strictly between 0% and 100%",
357                ));
358            }
359            Ok(Value::Float(libm::log(p / (1.0 - p))))
360        }
361        B::InvLogit => ops::computed_prob(1.0 / (1.0 + crate::math::exp(-number(a(0), "inv_logit")?)), "inv_logit"),
362        B::Date => {
363            let date = match args.len() {
364                1 => dates::parse(&text(a(0), "date")?),
365                3 => dates::from_parts(
366                    date_count(a(0), "date", budget)?,
367                    date_count(a(1), "date", budget)?,
368                    date_count(a(2), "date", budget)?,
369                ),
370                _ => {
371                    return Err(OpError::new(
372                        "`date` takes one ISO string or three integers (year, month, day)",
373                    ));
374                }
375            };
376            date.map(Value::Date).ok_or_else(|| {
377                OpError::fault(
378                    Fault::ConversionError,
379                    "invalid date; use YYYY-MM-DD within 0001-01-01..9999-12-31",
380                )
381            })
382        }
383        B::Days => to_int(a(0), f64::round),
384        B::Weeks if matches!(a(0), Value::Int(_)) => {
385            ops::binary(probl_syntax::ast::BinOp::Mul, a(0), &Value::Int(7.into()), budget)
386        }
387        B::Weeks => {
388            let n = number(a(0), "weeks")?;
389            to_int(&Value::Float(n * 7.0), f64::round)
390        }
391        B::AddWorkdays => {
392            let d = date_value(a(0), b.name())?;
393            let n = date_count(a(1), b.name(), budget)?;
394            let holidays = holiday_calendar(args.get(2), b.name(), budget)?;
395            date_result(dates::add_workdays_with_holidays(d, n, &holidays))
396        }
397        B::IsWorkday => {
398            let d = date_value(a(0), b.name())?;
399            let holidays = holiday_calendar(args.get(1), b.name(), budget)?;
400            Ok(Value::Bool(
401                dates::weekday(d) < 5 && holidays.binary_search(&d).is_err(),
402            ))
403        }
404        B::AddMonths | B::AddYears => {
405            let d = date_value(a(0), b.name())?;
406            let n = date_count(a(1), b.name(), budget)?;
407            date_result(if b == B::AddMonths {
408                dates::add_months(d, n)
409            } else {
410                dates::add_years(d, n)
411            })
412        }
413        B::StartOfMonth | B::EndOfMonth => {
414            let d = date_value(a(0), b.name())?;
415            date_result(if b == B::StartOfMonth {
416                dates::start_of_month(d)
417            } else {
418                dates::end_of_month(d)
419            })
420        }
421        B::Year | B::Month | B::Day => {
422            let (y, m, d) = dates::to_civil(date_value(a(0), b.name())? as i64);
423            Ok(Value::Int(match b {
424                B::Year => y.into(),
425                B::Month => m.into(),
426                _ => d.into(),
427            }))
428        }
429        B::Weekday => crate::text::value(
430            dates::WEEKDAYS[dates::weekday(date_value(a(0), b.name())?) as usize],
431            budget,
432        ),
433        B::IsFalse => Ok(Value::Bool(ops::is_certain(a(0), false))),
434        B::IsTrue => Ok(Value::Bool(ops::is_certain(a(0), true))),
435        B::IsListOfLen => Ok(Value::Bool(
436            matches!((a(0), a(1)), (Value::List(items), Value::Int(n)) if *n == items.len() as i64),
437        )),
438        B::Minimum | B::Maximum => unreachable!("population extrema receive distributions whole"),
439        B::Typeof | B::RunDate | B::Count | B::Map | B::Filter | B::Reduce | B::Print | B::Roll | B::Take => {
440            unreachable!("`{}` is handled by the interpreter", b.name())
441        }
442        B::P
443        | B::Mean
444        | B::Sd
445        | B::Variance
446        | B::Median
447        | B::MedianLow
448        | B::MedianHigh
449        | B::Quantile
450        | B::Support
451        | B::Cdf
452        | B::Pmf
453        | B::IterItems
454        | B::RepeatCount
455        | B::Pdf
456        | B::Settled => unreachable!("`{}` takes distributions as they are", b.name()),
457    }
458}
459
460/// Statistical queries take an explicit population, never an implicit point
461/// distribution. Also called before the interpreter's analytic-value guard so
462/// a scalar has the same type error in enumeration and sampling.
463pub fn check_query_input(b: Builtin, args: &[Value]) -> OpResult<()> {
464    use Builtin as B;
465    let expected = match b {
466        B::Minimum | B::Maximum => "a distribution or nonempty list, range or string",
467        B::P => "a boolean distribution (dist[bool])",
468        B::Pdf => "a continuous distribution",
469        B::Mean
470        | B::Variance
471        | B::Sd
472        | B::Median
473        | B::MedianLow
474        | B::MedianHigh
475        | B::Quantile
476        | B::Cdf
477        | B::Pmf
478        | B::Support => "a distribution, nonempty list or nonempty range",
479        _ => return Ok(()),
480    };
481    let v = &args[0];
482    let valid = match b {
483        B::Minimum | B::Maximum => matches!(
484            v,
485            Value::Dist(_) | Value::Continuous(_) | Value::List(_) | Value::Range(..) | Value::Str(_)
486        ),
487        B::P => matches!(v, Value::Dist(_)),
488        B::Pdf => matches!(v, Value::Dist(_) | Value::Continuous(_)),
489        _ => matches!(
490            v,
491            Value::Dist(_) | Value::Continuous(_) | Value::List(_) | Value::Range(..)
492        ),
493    };
494    if !valid {
495        let help = if b == B::P {
496            "use `report event` to measure a fact across worlds, or `prob(event)` to convert a bool to 0 or 1"
497        } else {
498            "use `report x` to summarize values across worlds, or put the model inside `simulate { ... }` to obtain a distribution"
499        };
500        return Err(OpError::new(format!(
501            "`{}` expects {expected}, found {}",
502            b.name(),
503            article(&v.kind())
504        ))
505        .help(help));
506    }
507    if !matches!(b, B::Minimum | B::Maximum) && matches!(v, Value::List(xs) if xs.is_empty()) {
508        return Err(OpError::fault(
509            Fault::EmptyCollection,
510            format!("`{}` needs a nonempty list", b.name()),
511        ));
512    }
513    Ok(())
514}
515
516/// Built-ins that receive distributions whole.
517pub fn call_raw(b: Builtin, args: &[Value], budget: &mut Budget) -> OpResult<Value> {
518    use Builtin as B;
519    check_query_input(b, args)?;
520    let v = &args[0];
521    if let Value::Range(lo, hi) = v {
522        if matches!(
523            b,
524            B::Mean
525                | B::Variance
526                | B::Sd
527                | B::Median
528                | B::MedianLow
529                | B::MedianHigh
530                | B::Quantile
531                | B::Cdf
532                | B::Pmf
533                | B::Support
534        ) {
535            return range_query(b, lo, hi, args.get(1), budget);
536        }
537    }
538    if matches!(b, B::P | B::Cdf | B::Pmf | B::Pdf) {
539        if let Value::Dist(d) = v {
540            if d.missing > 0.0 {
541                return Err(OpError::new(format!("`{}` cannot return an exact scalar probability while the distribution has unresolved mass", b.name()))
542                    .help("report a distribution comparison to retain probability bounds; scalar queries require a fully resolved distribution"));
543            }
544        }
545    }
546    if matches!(
547        b,
548        B::Mean
549            | B::Variance
550            | B::Sd
551            | B::Median
552            | B::MedianLow
553            | B::MedianHigh
554            | B::Quantile
555            | B::Cdf
556            | B::Pdf
557            | B::Pmf
558            | B::Support
559    ) && continuous_parts(v)
560    {
561        if matches!(b, B::Median | B::MedianLow | B::MedianHigh) {
562            let n = match v {
563                Value::Dist(d) => d.outcomes.len(),
564                _ => 1,
565            };
566            // Median bounds sort component supports before any CDF inversion.
567            budget.work((n as u64).saturating_mul(n.max(1).ilog2() as u64 + 1))?;
568        }
569        return continuous_query(b, args);
570    }
571    match b {
572        B::Minimum | B::Maximum => population_extreme(v, b == B::Maximum, None, budget),
573        B::P => probability_of(v),
574        B::Pdf => Err(OpError::new("pdf needs a continuous distribution")
575            .help("for a distribution whose outcomes can be listed, use `pmf`")),
576        B::Mean => {
577            let d = stat_dist(v, b.name(), budget)?;
578            if d.outcomes.iter().any(|(x, _)| matches!(x, Value::Date(_))) {
579                return date_mean(v, &d);
580            }
581            if d.outcomes.iter().any(|(x, _)| matches!(x, Value::Complex(_))) {
582                budget.work(d.outcomes.len() as u64)?;
583                let mut sum = Complex::new(0.0, 0.0)?;
584                let total = d.total();
585                for (x, p) in &d.outcomes {
586                    let x = complex_number(x, "mean")?;
587                    sum = sum.plus(x.times(Complex::new(p / total, 0.0)?)?)?;
588                }
589                Ok(Value::Complex(sum))
590            } else {
591                check_numeric(&d, "mean")?;
592                finite_float(d.mean().unwrap(), "mean")
593            }
594        }
595        B::Variance => {
596            numeric_dist(v, "variance", budget).and_then(|d| finite_float(d.variance().unwrap(), "variance"))
597        }
598        B::Sd => numeric_dist(v, "sd", budget).and_then(|d| finite_float(d.sd().unwrap(), "sd")),
599        B::Median | B::MedianLow | B::MedianHigh => median(v, b, budget),
600        B::Quantile => {
601            let q = to_prob(&args[1])?;
602            quantile(v, q, "quantile", budget)
603        }
604        B::Support => {
605            let d = stat_dist(v, b.name(), budget)?;
606            budget.collection(d.outcomes.len() as u128)?;
607            Ok(Value::list(d.outcomes.iter().map(|(x, _)| x.clone()).collect()))
608        }
609        B::Cdf => {
610            let d = stat_dist(v, b.name(), budget)?;
611            let mut p = crate::stats::Sum::default();
612            for (x, w) in &d.outcomes {
613                if statistical_compare(x, &args[1])?.is_le() {
614                    p.add(*w);
615                }
616            }
617            ops::computed_prob(p.value() / d.total(), "cdf")
618        }
619        B::Pmf => {
620            let d = stat_dist(v, b.name(), budget)?;
621            let p = crate::stats::sum(d.outcomes.iter().filter(|(x, _)| *x == args[1]).map(|(_, w)| *w));
622            ops::computed_prob(p / d.total(), "pmf")
623        }
624        B::IterItems => iter_items(v, budget),
625        B::RepeatCount => {
626            if v.is_uncertain() {
627                return Err(OpError::new(format!("`repeat` needs a number, not a {}", v.kind()))
628                    .help("draw a value first, like `let n ~ d6`, then `repeat n { … }`"));
629            }
630            let n = integer(v, "repeat count", budget)?;
631            if n.is_negative() {
632                return Err(OpError::fault(
633                    Fault::DomainError,
634                    "`repeat` needs a count of 0 or more",
635                ));
636            }
637            Ok(Value::Int(n.into_owned()))
638        }
639        B::Settled => match v {
640            v if v.is_uncertain() => Err(
641                OpError::new(format!("`match` needs a settled value, not a {}", v.kind()))
642                    .help("draw a value first, like `let x ~ d6`, and match on `x`"),
643            ),
644            other => Ok(other.clone()),
645        },
646        // Internal helpers that take their arguments as they are.
647        _ => call_plain(b, args, budget),
648    }
649}
650
651/// `P(d)`: query an explicitly constructed boolean distribution.
652fn probability_of(v: &Value) -> OpResult<Value> {
653    match v {
654        Value::Dist(d) if d.truth().is_some() => ops::computed_prob(d.truth().unwrap().0 / d.total(), "P"),
655        _ => Err(OpError::new(format!(
656            "P needs a boolean distribution (dist[bool]), found {}",
657            article(&v.kind())
658        ))
659        .help("compare a distribution, like `P(d6 > 4)`; use `report event` for a fact across worlds")),
660    }
661}
662
663/// The parameters of `binomial`, `poisson` or `geometric`, checked; `None`
664/// for another built-in, or when an argument is a distribution.
665pub fn counts(b: Builtin, args: &[Value], budget: &mut Budget) -> OpResult<Option<Counts>> {
666    use Builtin as B;
667    if args.iter().any(Value::is_uncertain) {
668        return Ok(None);
669    }
670    Ok(Some(match b {
671        B::Binomial => {
672            let n = whole(&args[0], "binomial's number of trials", budget)?;
673            if n < 0 {
674                return Err(OpError::fault(
675                    Fault::DomainError,
676                    "binomial needs a number of trials of 0 or more",
677                ));
678            }
679            Counts::Binomial {
680                n: n as u64,
681                p: to_prob(&args[1])?,
682            }
683        }
684        B::Poisson => {
685            let rate = number(&args[0], "poisson")?;
686            if rate < 0.0 || !rate.is_finite() {
687                return Err(OpError::fault(Fault::DomainError, "poisson needs a rate of 0 or more"));
688            }
689            if rate > 1e15 {
690                return Err(OpError::new("poisson's rate is too large to count exactly")
691                    .help("above 10¹⁵, use a normal distribution with the same mean and variance"));
692            }
693            Counts::Poisson { rate }
694        }
695        B::Geometric => {
696            let p = to_prob(&args[0])?;
697            if p <= 0.0 {
698                return Err(OpError::fault(
699                    Fault::DomainError,
700                    "geometric needs a chance of success above 0%",
701                ));
702            }
703            Counts::Geometric { p }
704        }
705        _ => return Ok(None),
706    }))
707}
708
709fn continuous(family: OpResult<Family>) -> OpResult<Value> {
710    family.map(|f| Value::Continuous(Arc::new(f)))
711}
712
713/// Whether a value is, or mixes in, a continuous distribution.
714fn continuous_parts(v: &Value) -> bool {
715    match v {
716        Value::Continuous(_) => true,
717        Value::Dist(d) => d.outcomes.iter().any(|(x, _)| matches!(x, Value::Continuous(_))),
718        _ => false,
719    }
720}
721
722/// Questions about a continuous distribution, or a mixture that includes
723/// one, answered from the formulas (docs/semantics.md, section 13).
724fn continuous_query(b: Builtin, args: &[Value]) -> OpResult<Value> {
725    use Builtin as B;
726    // PMF is a typed-outcome query even when continuous components are mixed
727    // in. Converting atoms to f64 here would collapse distinct outcomes.
728    if b == B::Pmf {
729        return match &args[0] {
730            Value::Dist(d) => {
731                let at = crate::stats::sum(
732                    d.outcomes
733                        .iter()
734                        .filter(|(x, _)| !matches!(x, Value::Continuous(_)) && *x == args[1])
735                        .map(|(_, p)| *p),
736                );
737                ops::computed_prob(at / d.total(), "pmf")
738            }
739            Value::Continuous(_) => ops::computed_prob(0.0, "pmf"),
740            _ => unreachable!("checked continuous input"),
741        };
742    }
743    let parts = match &args[0] {
744        Value::Continuous(f) => vec![(Part::Continuous(**f), 1.0)],
745        Value::Dist(d) => d
746            .outcomes
747            .iter()
748            .map(|(x, p)| {
749                let part = match x {
750                    Value::Continuous(f) => Part::Continuous(**f),
751                    other => Part::Point(number(other, b.name())?),
752                };
753                Ok((part, *p))
754            })
755            .collect::<OpResult<Vec<_>>>()?,
756        _ => unreachable!("checked by `continuous_parts`"),
757    };
758    let m = Mixture { parts };
759    match b {
760        B::Mean => finite_float(m.mean(), "mean"),
761        B::Variance => finite_float(m.variance(), "variance"),
762        B::Sd => finite_float(m.sd(), "sd"),
763        B::Median | B::MedianLow | B::MedianHigh => {
764            let (lo, hi) = m.median_bounds();
765            finite_float(
766                match b {
767                    B::MedianLow => lo,
768                    B::MedianHigh => hi,
769                    _ => crate::stats::midpoint(lo, hi),
770                },
771                b.name(),
772            )
773        }
774        B::Quantile => finite_float(m.quantile(to_prob(&args[1])?), "quantile"),
775        B::Cdf => ops::computed_prob(m.cdf(number(&args[1], "cdf")?), "cdf"),
776        B::Pdf => {
777            let x = number(&args[1], "pdf")?;
778            let mut density = 0.0;
779            for (part, p) in &m.parts {
780                match part {
781                    Part::Continuous(f) => density += p * f.pdf(x),
782                    Part::Analytic(a) => density += p * a.pdf(x),
783                    Part::Point(_) => {
784                        return Err(OpError::new(
785                            "pdf needs a continuous distribution, without single values mixed in",
786                        ));
787                    }
788                }
789            }
790            finite_float(density, "pdf")
791        }
792        _ => Err(OpError::new(format!(
793            "`{}` needs a distribution whose outcomes can be listed, not a continuous one",
794            b.name()
795        ))),
796    }
797}
798
799/// Unit-step integer ranges are uniform populations; scalar queries do not
800/// allocate their support. Rank arithmetic remains exact for arbitrary integers.
801fn range_query(b: Builtin, lo: &Integer, hi: &Integer, arg: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
802    use Builtin as B;
803    if hi < lo {
804        return Err(OpError::fault(
805            Fault::EmptyCollection,
806            format!("`{}` needs a nonempty range", b.name()),
807        ));
808    }
809    budget.integer_work(lo, hi, false)?;
810    let n = range_len(lo, hi)?;
811    budget.integer_allocation(n.bits(), 1)?;
812    let two = Integer::from(2);
813    match b {
814        B::Support => Ok(Value::list(range_items(lo, hi, budget)?)),
815        B::Mean => lo.add(hi)?.ratio(&two).map(Value::Float).ok_or_else(|| {
816            OpError::fault(
817                Fault::NumericOverflow,
818                "`mean` gave a result that isn't a finite number",
819            )
820        }),
821        B::Median | B::MedianLow | B::MedianHigh => {
822            budget.integer_work(&n, &two, true)?;
823            let (half, odd) = n.div_mod(&two)?;
824            let high = Value::Int(lo.add(&half)?);
825            let low = if odd.is_zero() {
826                Value::Int(lo.add(&half)?.sub(&Integer::ONE)?)
827            } else {
828                high.clone()
829            };
830            match b {
831                B::MedianLow => Ok(low),
832                B::MedianHigh => Ok(high),
833                _ => midpoint(&low, &high, budget),
834            }
835        }
836        B::Sd | B::Variance => {
837            let quarter = n.ratio(&Integer::from(4)).unwrap_or(f64::INFINITY);
838            let inv = Integer::ONE.ratio(&n).unwrap_or(0.0);
839            let sd = quarter * (4.0 / libm::sqrt(12.0)) * libm::sqrt((1.0 - inv) * (1.0 + inv));
840            finite_float(if b == B::Sd { sd } else { sd * sd }, b.name())
841        }
842        B::Quantile => {
843            let q = to_prob(arg.expect("quantile argument"))?;
844            if q == 0.0 {
845                return Ok(Value::Int(lo.clone()));
846            }
847            budget.work(n.bits().div_ceil(64).saturating_mul(17))?;
848            let rank = if let Some(count) = n.to_u64().filter(|n| *n <= 1 << 53) {
849                // Match the rounded cumulative weights of a finite uniform
850                // population, including decimal boundaries such as 10% of 10.
851                let mut weight = 1.0 / count as f64;
852                let total = weight * count as f64;
853                if total != 1.0 {
854                    weight *= 1.0 / total;
855                }
856                let target = q * (weight * count as f64);
857                let (mut left, mut right) = (1, count);
858                while left < right {
859                    budget.work(1)?;
860                    let mid = left + (right - left) / 2;
861                    if mid as f64 * weight >= target {
862                        right = mid;
863                    } else {
864                        left = mid + 1;
865                    }
866                }
867                Integer::from(left)
868            } else {
869                // Above exact f64 integer precision, keep the rank arithmetic
870                // exact instead of rounding the population size to a float.
871                n.probability_rank(q)?
872            };
873            Ok(Value::Int(lo.add(&rank.sub(&Integer::ONE)?)?))
874        }
875        B::Pmf => {
876            let present = matches!(arg, Some(Value::Int(k)) if lo <= k && k <= hi);
877            ops::computed_prob(
878                if present {
879                    Integer::ONE.ratio(&n).unwrap_or(0.0)
880                } else {
881                    0.0
882                },
883                "pmf",
884            )
885        }
886        B::Cdf => {
887            let x = arg.expect("cdf argument");
888            if ops::compare(x, &Value::Int(lo.clone()))?.is_lt() {
889                return ops::computed_prob(0.0, "cdf");
890            }
891            if ops::compare(x, &Value::Int(hi.clone()))?.is_ge() {
892                return ops::computed_prob(1.0, "cdf");
893            }
894            let k = match x {
895                Value::Int(k) => k.clone(),
896                _ => Integer::from_f64(number(x, "cdf")?.floor()).expect("finite numeric bound"),
897            };
898            let count = k.sub(lo)?.add(&Integer::ONE)?;
899            budget.integer_work(&count, &n, true)?;
900            ops::computed_prob(count.ratio(&n).unwrap_or(0.0), "cdf")
901        }
902        _ => unreachable!("checked range query"),
903    }
904}
905
906/// A list is an empirical distribution: each element has equal weight,
907/// including repetitions. Its elements remain values, never implicit draws.
908fn stat_dist<'a>(v: &'a Value, what: &str, budget: &mut Budget) -> OpResult<Cow<'a, Dist>> {
909    match v {
910        Value::Dist(d) => {
911            budget.work(d.outcomes.len() as u64)?;
912            Ok(Cow::Borrowed(d))
913        }
914        Value::List(xs) if !xs.is_empty() => {
915            budget.collection(xs.len() as u128)?;
916            let comparisons = (xs.len().ilog2() + 1) as u64;
917            budget.work((xs.len() as u64).saturating_mul(comparisons))?;
918            for x in xs.iter() {
919                match x {
920                    Value::Int(n) => {
921                        budget.integer_work(n, &Integer::ONE, false)?;
922                    }
923                    Value::Str(s) => {
924                        budget.string_work(s)?;
925                    }
926                    x if x.is_uncertain() => {
927                        return Err(OpError::new(format!(
928                            "`{what}` needs list elements that are values, found {}",
929                            article(&x.kind())
930                        ))
931                        .help("draw the elements first, or explicitly build a mixture with `one_of`"));
932                    }
933                    _ => {}
934                }
935            }
936            Ok(Cow::Owned(Dist::uniform(xs.to_vec())))
937        }
938        _ => Err(OpError::new(format!(
939            "`{what}` needs a distribution or a nonempty list"
940        ))),
941    }
942}
943
944fn numeric_dist<'a>(v: &'a Value, what: &str, budget: &mut Budget) -> OpResult<Cow<'a, Dist>> {
945    let d = stat_dist(v, what, budget)?;
946    check_numeric(&d, what)?;
947    Ok(d)
948}
949
950fn check_numeric(d: &Dist, what: &str) -> OpResult<()> {
951    for (x, _) in &d.outcomes {
952        if !matches!(x, Value::Int(_) | Value::Float(_) | Value::Prob(_)) {
953            return Err(OpError::new(format!(
954                "`{what}` needs real numeric elements, found {}",
955                article(&x.kind())
956            )));
957        }
958        if x.as_f64().is_none_or(|x| !x.is_finite()) {
959            return Err(OpError::new(format!(
960                "`{what}` needs numbers representable as finite floats"
961            )));
962        }
963    }
964    if d.outcomes.is_empty() {
965        return Err(OpError::new(format!("`{what}` needs at least one resolved outcome")));
966    }
967    Ok(())
968}
969
970fn statistical_compare(a: &Value, b: &Value) -> OpResult<std::cmp::Ordering> {
971    match (a, b) {
972        (Value::Bool(a), Value::Bool(b)) => Ok(a.cmp(b)),
973        _ => ops::compare(a, b),
974    }
975}
976
977/// Public ordering is separate from the typed total order used for storage.
978/// Never put this reordered vector back into a Dist or a world/cache key.
979fn ordered_outcomes(d: &Dist, budget: &mut Budget) -> OpResult<Vec<(Value, f64)>> {
980    budget.collection(d.outcomes.len() as u128)?;
981    for (x, _) in &d.outcomes {
982        budget.work(1)?;
983        statistical_compare(x, x)?;
984    }
985    crate::ordering::reserve_sort(d.outcomes.len(), budget)?;
986    let mut outcomes = d.outcomes.clone();
987    crate::ordering::try_sort_by(&mut outcomes, |(a, _), (b, _)| {
988        budget.work(1)?;
989        statistical_compare(a, b)
990    })?;
991    Ok(outcomes)
992}
993
994fn quantile(v: &Value, q: f64, what: &str, budget: &mut Budget) -> OpResult<Value> {
995    let d = stat_dist(v, what, budget)?;
996    let outcomes = ordered_outcomes(&d, budget)?;
997    crate::stats::quantile(&outcomes, q)
998        .cloned()
999        .ok_or_else(|| OpError::new(format!("`{what}` needs at least one resolved outcome")))
1000}
1001
1002fn median(v: &Value, b: Builtin, budget: &mut Budget) -> OpResult<Value> {
1003    let d = stat_dist(v, b.name(), budget)?;
1004    for (x, _) in &d.outcomes {
1005        budget.work(1)?;
1006        // Validate even singletons. Being sortable internally does not give a
1007        // record, complex number or recipe a mathematical ordering.
1008        statistical_compare(x, x)?;
1009        if b == Builtin::Median && !matches!(x, Value::Int(_) | Value::Float(_) | Value::Prob(_) | Value::Date(_)) {
1010            return Err(OpError::new(format!(
1011                "`median` needs real numeric or date elements, found {}",
1012                article(&x.kind())
1013            ))
1014            .help("use `median_low` or `median_high` for ordered values such as strings"));
1015        }
1016    }
1017    let outcomes = ordered_outcomes(&d, budget)?;
1018    let (lo, hi) = crate::stats::median_bounds(&outcomes)
1019        .ok_or_else(|| OpError::new(format!("`{}` needs at least one resolved outcome", b.name())))?;
1020    match b {
1021        Builtin::MedianLow => Ok(lo.clone()),
1022        Builtin::MedianHigh => Ok(hi.clone()),
1023        _ if lo == hi => Ok(lo.clone()),
1024        _ => midpoint(lo, hi, budget),
1025    }
1026}
1027
1028fn date_mean(v: &Value, d: &Dist) -> OpResult<Value> {
1029    // Lists permit exact integer arithmetic, so rounding cannot depend on
1030    // repeated values, list order, or a date's distance from the epoch.
1031    if let Value::List(xs) = v {
1032        let total = xs
1033            .iter()
1034            .try_fold(0i128, |sum, x| Ok::<_, OpError>(sum + date_value(x, "mean")? as i128))?;
1035        let n = xs.len() as i128;
1036        let day = total.div_euclid(n) + i128::from(2 * total.rem_euclid(n) > n);
1037        return Ok(Value::Date(day as i32));
1038    }
1039    let origin = date_value(&d.outcomes[0].0, "mean")?;
1040    let mut offsets = crate::stats::Sum::default();
1041    let mut weights = crate::stats::Sum::default();
1042    for (x, w) in &d.outcomes {
1043        offsets.add((date_value(x, "mean")? - origin) as f64 * w);
1044        weights.add(*w);
1045    }
1046    let days = offsets.value() / weights.value();
1047    let floor = days.floor();
1048    let tie_error = 4.0 * f64::EPSILON * days.abs().max(1.0);
1049    let rounded = floor + f64::from(days - floor > 0.5 + tie_error);
1050    date_result(
1051        i32::try_from(origin as i64 + rounded as i64)
1052            .ok()
1053            .filter(|d| dates::valid(*d)),
1054    )
1055}
1056
1057/// Preserve exact integral midpoints, including bigints beyond float range.
1058pub(crate) fn midpoint(a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
1059    if let (Value::Date(a), Value::Date(b)) = (a, b) {
1060        return Ok(Value::Date(((*a as i64 + *b as i64).div_euclid(2)) as i32));
1061    }
1062    if let (Value::Int(a), Value::Int(b)) = (a, b) {
1063        budget.integer_work(a, b, false)?;
1064        let sum = a.add(b)?;
1065        budget.integer_allocation(sum.bits(), 1)?;
1066        let two = Integer::from(2);
1067        budget.integer_work(&sum, &two, true)?;
1068        let (whole, remainder) = sum.div_mod(&two)?;
1069        if remainder.is_zero() {
1070            return Ok(Value::Int(whole));
1071        }
1072        return sum.ratio(&two).map(Value::Float).ok_or_else(|| {
1073            OpError::fault(
1074                Fault::NumericOverflow,
1075                "median's fractional midpoint is too large for a finite float",
1076            )
1077        });
1078    }
1079    let (a, b) = (number(a, "median")?, number(b, "median")?);
1080    // Same-sign subtraction and opposite-sign addition avoid overflow; this
1081    // also preserves equal subnormal values instead of halving both to zero.
1082    finite_float(crate::stats::midpoint(a, b), "median")
1083}
1084
1085fn iter_items(v: &Value, budget: &mut Budget) -> OpResult<Value> {
1086    match v {
1087        Value::List(_) | Value::Range(..) => Ok(v.clone()),
1088        Value::Map(m) => Ok(Value::list(
1089            m.iter().map(|(k, x)| Value::list(vec![k.clone(), x.clone()])).collect(),
1090        )),
1091        Value::Bag(b) => {
1092            let n: u128 = b.values().map(|n| *n as u128).sum();
1093            budget.collection(n)?;
1094            Ok(Value::list(
1095                b.iter()
1096                    .flat_map(|(k, n)| std::iter::repeat_n(k.clone(), *n as usize))
1097                    .collect(),
1098            ))
1099        }
1100        Value::Str(s) => crate::text::chars(s, budget).map(Value::list),
1101        v if v.is_uncertain() => Err(OpError::new(format!("can't loop over a {}", v.kind()))
1102            .help("draw a value first with `~`, or loop over `support(…)`")),
1103        other => Err(OpError::new(format!("can't loop over {}", article(&other.kind())))),
1104    }
1105}
1106
1107fn date_count(v: &Value, func: &str, budget: &mut Budget) -> OpResult<i64> {
1108    integer(v, func, budget)?
1109        .to_i64()
1110        .ok_or_else(|| OpError::fault(Fault::NumericOverflow, "date out of range"))
1111}
1112
1113fn date_value(v: &Value, func: &str) -> OpResult<i32> {
1114    match v {
1115        Value::Date(d) if dates::valid(*d) => Ok(*d),
1116        Value::Date(_) => Err(OpError::fault(Fault::NumericOverflow, "date out of range")),
1117        _ => Err(expected("a date", v, func)),
1118    }
1119}
1120
1121fn date_result(date: Option<i32>) -> OpResult<Value> {
1122    date.map(Value::Date)
1123        .ok_or_else(|| OpError::fault(Fault::NumericOverflow, "date out of range (0001-01-01..9999-12-31)"))
1124}
1125
1126fn holiday_calendar(v: Option<&Value>, func: &str, budget: &mut Budget) -> OpResult<Vec<i32>> {
1127    let Some(v) = v else {
1128        return Ok(Vec::new());
1129    };
1130    let values = list(v, func)?;
1131    budget.collection(values.len() as u128)?;
1132    let log = (values.len() as u64).checked_ilog2().unwrap_or(0) as u64 + 1;
1133    budget.work((values.len() as u64).saturating_mul(log))?;
1134    let mut holidays = Vec::new();
1135    for value in values.iter() {
1136        let d = date_value(value, func)?;
1137        if dates::weekday(d) < 5 {
1138            holidays.push(d);
1139        }
1140    }
1141    holidays.sort_unstable();
1142    holidays.dedup();
1143    Ok(holidays)
1144}
1145
1146fn expected(what: &str, v: &Value, func: &str) -> OpError {
1147    OpError::new(format!("`{func}` needs {what}, found {}", article(&v.kind())))
1148}
1149
1150fn number(v: &Value, func: &str) -> OpResult<f64> {
1151    match v {
1152        Value::Bool(_) => Err(expected("a number", v, func)),
1153        Value::Continuous(_) => {
1154            Err(expected("a number", v, func).help("draw a value first, like `let x ~ normal(0, 1)`"))
1155        }
1156        Value::Int(n) => n.to_f64().ok_or_else(|| {
1157            OpError::fault(
1158                Fault::NumericOverflow,
1159                format!("`{func}` needs an integer that fits in a finite float"),
1160            )
1161        }),
1162        _ => v.as_f64().ok_or_else(|| expected("a number", v, func)),
1163    }
1164}
1165
1166fn complex_number(v: &Value, func: &str) -> OpResult<Complex> {
1167    v.as_complex()
1168        .ok_or_else(|| expected("a finite real or complex number", v, func))
1169}
1170
1171fn finite_float(x: f64, func: &str) -> OpResult<Value> {
1172    if x.is_finite() {
1173        Ok(Value::Float(x))
1174    } else {
1175        Err(OpError::new(format!(
1176            "`{func}` gave a result that isn't a finite number"
1177        )))
1178    }
1179}
1180
1181fn whole(v: &Value, what: &str, budget: &mut Budget) -> OpResult<i64> {
1182    integer(v, what, budget)?
1183        .to_i64()
1184        .ok_or_else(|| OpError::new(format!("{what} is outside the supported count range")))
1185}
1186
1187fn nonnegative_int<'a>(v: &'a Value, func: &str, budget: &mut Budget) -> OpResult<Cow<'a, Integer>> {
1188    let n = integer(v, func, budget)?;
1189    if n.is_negative() {
1190        Err(OpError::fault(
1191            Fault::DomainError,
1192            format!("`{func}` needs nonnegative integers"),
1193        ))
1194    } else {
1195        Ok(n)
1196    }
1197}
1198
1199fn choose(n: &Integer, k: &Integer, budget: &mut Budget) -> OpResult<Value> {
1200    if k > n {
1201        return Ok(Value::Int(Integer::ZERO));
1202    }
1203    let k = k.min(&n.sub(k)?).clone();
1204    let steps = k
1205        .to_u64()
1206        .ok_or_else(|| OpError::limit("choose needs too many iterations"))?;
1207    budget.work(steps)?;
1208    let offset = n.sub(&k)?;
1209    let mut result = Integer::ONE;
1210    for i in 1..=steps {
1211        let numerator = offset.add(&i.into())?;
1212        let divisor = Integer::from(i);
1213        // Cancel first: an intermediate product must not exceed the integer
1214        // ceiling when the final binomial coefficient fits.
1215        let d = gcd(numerator.clone(), divisor.clone(), budget)?;
1216        let numerator = numerator.div_mod(&d)?.0;
1217        let divisor = divisor.div_mod(&d)?.0;
1218        budget.integer_work(&result, &divisor, true)?;
1219        result = result.div_mod(&divisor)?.0;
1220        budget.integer_work(&result, &numerator, true)?;
1221        result = result.mul(&numerator)?;
1222        budget.integer_bits(result.bits())?;
1223    }
1224    Ok(Value::Int(result))
1225}
1226
1227fn factorial(n: &Integer, budget: &mut Budget) -> OpResult<Value> {
1228    let n = n
1229        .to_u64()
1230        .ok_or_else(|| OpError::limit("factorial exceeds the integer size limit"))?;
1231    // n! contains at least n/2 factors of n/2 or more.
1232    let half = n / 2;
1233    let lower_bits = half.saturating_mul(63u64.saturating_sub(u64::from(half.leading_zeros())));
1234    budget.integer_bits(lower_bits)?;
1235    budget.work(n)?;
1236    let mut result = Integer::ONE;
1237    for i in 2..=n {
1238        let factor = Integer::from(i);
1239        budget.integer_work(&result, &factor, true)?;
1240        result = result.mul(&factor)?;
1241        budget.integer_bits(result.bits())?;
1242    }
1243    Ok(Value::Int(result))
1244}
1245
1246fn gcd(mut a: Integer, mut b: Integer, budget: &mut Budget) -> OpResult<Integer> {
1247    while !b.is_zero() {
1248        budget.integer_work(&a, &b, true)?;
1249        let r = a.div_mod(&b)?.1;
1250        a = b;
1251        b = r;
1252    }
1253    Ok(a)
1254}
1255
1256fn euler_phi(n: &Integer, budget: &mut Budget) -> OpResult<Value> {
1257    if n.is_zero() {
1258        return Err(OpError::fault(
1259            Fault::DomainError,
1260            "`euler_phi` needs a positive integer",
1261        ));
1262    }
1263    let mut n = n.clone();
1264    let mut result = n.clone();
1265    let mut divisor = Integer::from(2);
1266    loop {
1267        budget.integer_work(&n, &divisor, true)?;
1268        let (q, rem) = n.div_mod(&divisor)?;
1269        if divisor > q {
1270            break;
1271        }
1272        if rem.is_zero() {
1273            budget.integer_work(&result, &divisor, true)?;
1274            result = result.sub(&result.div_mod(&divisor)?.0)?;
1275            n = q;
1276            loop {
1277                budget.integer_work(&n, &divisor, true)?;
1278                let (q, rem) = n.div_mod(&divisor)?;
1279                if !rem.is_zero() {
1280                    break;
1281                }
1282                n = q;
1283            }
1284        }
1285        divisor = divisor.add(&if divisor == 2 { Integer::ONE } else { 2.into() })?;
1286    }
1287    if n > 1 {
1288        budget.integer_work(&result, &n, true)?;
1289        result = result.sub(&result.div_mod(&n)?.0)?;
1290    }
1291    Ok(Value::Int(result))
1292}
1293
1294fn text(v: &Value, func: &str) -> OpResult<Arc<str>> {
1295    match v {
1296        Value::Str(s) => Ok(s.clone()),
1297        other => Err(expected("a string", other, func)),
1298    }
1299}
1300
1301fn list<'a>(v: &'a Value, func: &str) -> OpResult<&'a [Value]> {
1302    match v {
1303        Value::List(items) => Ok(items),
1304        other => Err(expected("a list", other, func)),
1305    }
1306}
1307
1308fn range_items(lo: &Integer, hi: &Integer, budget: &mut Budget) -> OpResult<Vec<Value>> {
1309    let n = range_count(lo, hi)?;
1310    budget.collection(n)?;
1311    let n = usize::try_from(n).map_err(|_| OpError::limit("range has too many elements"))?;
1312    budget.integer_allocation(lo.bits().max(hi.bits()), n as u64)?;
1313    budget.work((n as u64).saturating_mul(lo.bits().max(hi.bits()).div_ceil(64).max(1)))?;
1314    (0..n)
1315        .map(|i| lo.add(&Integer::from(i)).map(Value::Int).map_err(OpError::from))
1316        .collect()
1317}
1318
1319/// The elements of a list-like value, checking that a range isn't too long
1320/// to spell out.
1321pub fn items(v: &Value, func: &str, budget: &mut Budget) -> OpResult<Vec<Value>> {
1322    match v {
1323        Value::List(items) => {
1324            budget.collection(items.len() as u128)?;
1325            budget.work(items.len() as u64)?;
1326            Ok(items.to_vec())
1327        }
1328        Value::Range(lo, hi) => range_items(lo, hi, budget),
1329        Value::Str(s) => crate::text::chars(s, budget),
1330        other => Err(expected("a list, range or string", other, func)),
1331    }
1332}
1333
1334fn sort(items: &mut [Value], descending: bool, budget: &mut Budget) -> OpResult<()> {
1335    // Even a singleton must have a language ordering; internal storage order
1336    // doesn't make complex values, records or distributions sortable.
1337    for x in items.iter() {
1338        budget.work(1)?;
1339        ops::compare(x, x)?;
1340    }
1341    crate::ordering::reserve_sort(items.len(), budget)?;
1342    crate::ordering::try_sort_by(items, |x, y| {
1343        budget.work(1)?;
1344        ops::compare(x, y).map(|c| if descending { c.reverse() } else { c })
1345    })
1346}
1347
1348/// Real inputs keep their real domains; an explicit complex input requests the
1349/// principal complex extension. Distribution lifting happens in the interpreter.
1350fn elementary1(
1351    v: &Value,
1352    func: &str,
1353    complex: fn(Complex) -> OpResult<Complex>,
1354    real: impl Fn(f64) -> Option<f64>,
1355) -> OpResult<Value> {
1356    match v {
1357        Value::Complex(z) => complex(*z).map(Value::Complex).map_err(|_| {
1358            OpError::fault(
1359                Fault::DomainError,
1360                format!("`{func}` isn't defined for {v} or its result isn't finite"),
1361            )
1362        }),
1363        _ => float1(v, func, real),
1364    }
1365}
1366
1367fn float1(v: &Value, func: &str, f: impl Fn(f64) -> Option<f64>) -> OpResult<Value> {
1368    let x = number(v, func)?;
1369    f(x).filter(|y| x.is_finite() && y.is_finite())
1370        .map(Value::Float)
1371        .ok_or_else(|| {
1372            OpError::fault(
1373                Fault::DomainError,
1374                format!("`{func}` isn't defined for {}", fmt_float(x)),
1375            )
1376        })
1377}
1378
1379fn float2(a: &Value, b: &Value, func: &str, f: fn(f64, f64) -> f64) -> OpResult<Value> {
1380    let (a, b) = (number(a, func)?, number(b, func)?);
1381    if !a.is_finite() || !b.is_finite() {
1382        return Err(OpError::new(format!("`{func}` needs finite numbers")));
1383    }
1384    // Value equality merges signed zeros. In particular, atan2 must not
1385    // distinguish worlds (or memoized arguments) that the engine considers equal.
1386    let a = if a == 0.0 { 0.0 } else { a };
1387    let b = if b == 0.0 { 0.0 } else { b };
1388    let value = f(a, b);
1389    if !value.is_finite() {
1390        return Err(OpError::new(format!(
1391            "`{func}` gave a result that isn't a finite number"
1392        )));
1393    }
1394    Ok(Value::Float(value))
1395}
1396
1397fn to_int(v: &Value, f: fn(f64) -> f64) -> OpResult<Value> {
1398    match v {
1399        Value::Int(i) => Ok(Value::Int(i.clone())),
1400        other => {
1401            let x = number(other, "rounding")?;
1402            let r = f(x);
1403            Integer::from_f64(r).map(Value::Int).ok_or_else(|| {
1404                OpError::fault(
1405                    Fault::NumericOverflow,
1406                    format!("{} cannot be rounded to a finite int", fmt_float(x)),
1407                )
1408            })
1409        }
1410    }
1411}
1412
1413fn round(v: &Value, digits: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1414    let Some(digits) = digits else {
1415        return to_int(v, f64::round);
1416    };
1417    let digits = integer(digits, "round's digits", budget)?;
1418    if let Value::Int(n) = v {
1419        if *digits >= 0 {
1420            return Ok(v.clone());
1421        }
1422        let places = digits.abs();
1423        if places > Integer::from(probl_number::MAX_INTEGER_DIGITS) {
1424            return Ok(Value::Int(Integer::ZERO));
1425        }
1426        budget.integer_work(n, n, true)?;
1427        return Ok(Value::Int(n.round_decimal(places.to_u64().unwrap() as u32)?));
1428    }
1429    // Once outside this interval the exact magnitude of digits is irrelevant.
1430    let digits = if *digits > 323 {
1431        324
1432    } else if *digits < -308 {
1433        -309
1434    } else {
1435        digits.to_i64().unwrap()
1436    };
1437    float1(v, "round", |x| {
1438        // Beyond these bounds a decimal place cannot change a finite f64,
1439        // or every finite f64 rounds to zero. No unbounded powers or loops.
1440        if digits > 323 || x == 0.0 {
1441            return Some(x);
1442        }
1443        if digits < -308 {
1444            return Some(0.0);
1445        }
1446        let rounded = if digits >= 0 {
1447            // Splitting the scale supports subnormals (up to 323 places)
1448            // without forming an infinite power of ten.
1449            let high = libm::pow(10.0, digits.min(308) as f64);
1450            let low = libm::pow(10.0, (digits - 308).max(0) as f64);
1451            let scaled = (x * high) * low;
1452            // At this precision the decimal adjustment is smaller than half
1453            // an f64 step; scaling back could only introduce a new error.
1454            if scaled.abs() >= 1e16 || x.fract() == 0.0 {
1455                return Some(x);
1456            }
1457            (scaled.round() / low) / high
1458        } else {
1459            let scale = libm::pow(10.0, -digits as f64);
1460            let scaled = x / scale;
1461            if scaled.abs() >= 1e16 {
1462                return Some(x);
1463            }
1464            scaled.round() * scale
1465        };
1466        Some(rounded)
1467    })
1468}
1469
1470fn min_max(args: &[Value], want_max: bool, budget: &mut Budget) -> OpResult<Value> {
1471    let name = if want_max { "max" } else { "min" };
1472    let mut best: Option<Value> = None;
1473    for v in args {
1474        // Validate candidate outcomes, then lift each comparison. Reducing incrementally avoids a Cartesian product of
1475        // every candidate and preserves missing mass without drawing.
1476        let v = ops::lift1(v, budget, |v, budget| {
1477            budget.work(1)?;
1478            ops::compare(v, v)?;
1479            Ok(v.clone())
1480        })?;
1481        best = Some(match best {
1482            None => v,
1483            Some(b) => ops::lift2(&b, &v, budget, |b, v, _| {
1484                let ord = ops::compare(v, b)?;
1485                if (want_max && ord.is_gt()) || (!want_max && ord.is_lt()) {
1486                    Ok(v.clone())
1487                } else {
1488                    Ok(b.clone())
1489                }
1490            })?,
1491        });
1492    }
1493    best.ok_or_else(|| OpError::fault(Fault::EmptyCollection, format!("`{name}` of an empty list")))
1494}
1495
1496fn slice(v: &Value, start: &Value, end: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1497    let start = integer(start, "slice", budget)?;
1498    let end = end.map(|v| integer(v, "slice", budget)).transpose()?;
1499    let length: Integer = match v {
1500        Value::List(xs) => xs.len().into(),
1501        Value::Str(s) => s.chars().count().into(),
1502        Value::Range(lo, hi) => {
1503            budget.integer_work(lo, hi, false)?;
1504            range_len(lo, hi)?
1505        }
1506        other => return Err(expected("a list, range or string", other, "slice")),
1507    };
1508    let end = end.as_deref().unwrap_or(&length);
1509    let start = start.as_ref();
1510    if start.is_negative() || start > end || end > &length {
1511        return Err(OpError::fault(
1512            Fault::IndexOutOfBounds,
1513            "`slice` needs 0 <= start <= end <= length",
1514        ));
1515    }
1516    Ok(match v {
1517        Value::Range(lo, _) => {
1518            // Empty slices have a canonical empty range. Subtract one from the
1519            // offset before adding, so a slice ending at the largest int works.
1520            let (first, last) = if start == end {
1521                (Integer::ZERO, (-1).into())
1522            } else {
1523                (lo.add(start)?, lo.add(&end.sub(&Integer::ONE)?)?)
1524            };
1525            budget.integer_allocation(first.bits(), 1)?;
1526            budget.integer_allocation(last.bits(), 1)?;
1527            Value::Range(first, last)
1528        }
1529        Value::List(xs) => {
1530            let (start, end) = (start.to_u64().unwrap() as usize, end.to_u64().unwrap() as usize);
1531            budget.collection((end - start) as u128)?;
1532            budget.work((end - start) as u64)?;
1533            Value::list(xs[start..end].to_vec())
1534        }
1535        Value::Str(s) => {
1536            let (start, end) = (start.to_u64().unwrap() as usize, end.to_u64().unwrap() as usize);
1537            let boundary = |n| s.char_indices().nth(n).map_or(s.len(), |(i, _)| i);
1538            crate::text::value(&s[boundary(start)..boundary(end)], budget)?
1539        }
1540        _ => unreachable!(),
1541    })
1542}
1543
1544fn len(v: &Value) -> OpResult<Value> {
1545    let n: u128 = match v {
1546        Value::List(items) => items.len() as u128,
1547        Value::Str(s) => s.chars().count() as u128,
1548        Value::Map(m) => m.len() as u128,
1549        Value::Bag(b) => b.values().map(|n| *n as u128).sum(),
1550        Value::Range(lo, hi) => return Ok(Value::Int(range_len(lo, hi)?)),
1551        other => return Err(expected("a collection", other, "len")),
1552    };
1553    Ok(Value::Int(n.into()))
1554}
1555
1556fn sum(v: &Value, budget: &mut Budget) -> OpResult<Value> {
1557    let items = items(v, "sum", budget)?;
1558    let mut acc = Value::Int(0.into());
1559    for x in &items {
1560        acc = ops::binary(probl_syntax::ast::BinOp::Add, &acc, x, budget)?;
1561    }
1562    Ok(acc)
1563}
1564
1565fn get(coll: &Value, key: &Value, default: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1566    let found = match coll {
1567        Value::Map(m) => m.get(key).cloned(),
1568        Value::List(items) => integer(key, "index", budget)?
1569            .to_u64()
1570            .and_then(|n| usize::try_from(n).ok())
1571            .and_then(|n| items.get(n))
1572            .cloned(),
1573        Value::Str(s) => {
1574            let index = integer(key, "index", budget)?
1575                .to_u64()
1576                .and_then(|n| usize::try_from(n).ok());
1577            budget.string_work(s)?;
1578            index
1579                .and_then(|n| s.chars().nth(n))
1580                .map(|c| crate::text::value(c.encode_utf8(&mut [0; 4]), budget))
1581                .transpose()?
1582        }
1583        Value::Range(lo, hi) => {
1584            let index = integer(key, "index", budget)?;
1585            budget.integer_work(lo, hi, false)?;
1586            let len = range_len(lo, hi)?;
1587            if index.is_negative() || *index >= len {
1588                None
1589            } else {
1590                Some(Value::Int(lo.add(&index)?))
1591            }
1592        }
1593        Value::Bag(b) => Some(Value::Int(b.get(key).copied().unwrap_or(0).into())),
1594        other => return Err(expected("a map, bag or sequence", other, "get")),
1595    };
1596    match (found, default) {
1597        (Some(v), _) => Ok(v),
1598        (None, Some(d)) => Ok(d.clone()),
1599        (None, None) => Err(
1600            OpError::fault(Fault::MissingKey, format!("the key {key:?} isn't there"))
1601                .help("give a default: `get(key, default)`"),
1602        ),
1603    }
1604}
1605
1606/// Select an existing element; recipes inside collections are never lifted.
1607pub fn population_extreme(v: &Value, want_max: bool, default: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1608    let name = if want_max { "maximum" } else { "minimum" };
1609    match v {
1610        Value::Range(lo, hi) => {
1611            budget.integer_work(lo, hi, false)?;
1612            if hi < lo {
1613                return empty_extreme(name, default);
1614            }
1615            Ok(Value::Int(if want_max { hi.clone() } else { lo.clone() }))
1616        }
1617        Value::Continuous(f) => {
1618            let (lo, hi) = f.support();
1619            finite_bound(if want_max { hi } else { lo }, name)
1620        }
1621        Value::Dist(d) => {
1622            if d.missing > 0.0 {
1623                return Err(OpError::new(format!(
1624                    "`{name}` cannot determine a support bound while the distribution has unresolved mass"
1625                )));
1626            }
1627            let mut best = None;
1628            for (x, p) in &d.outcomes {
1629                if *p <= 0.0 {
1630                    continue;
1631                }
1632                let x = match x {
1633                    Value::Continuous(_) | Value::Dist(_) => population_extreme(x, want_max, None, budget)?,
1634                    x => x.clone(),
1635                };
1636                select_extreme(&mut best, x, want_max, budget)?;
1637            }
1638            best.ok_or_else(|| OpError::new(format!("`{name}` needs a nonempty distribution")))
1639        }
1640        Value::List(_) | Value::Str(_) => {
1641            let mut best = None;
1642            for x in items(v, name, budget)? {
1643                select_extreme(&mut best, x, want_max, budget)?;
1644            }
1645            best.map_or_else(|| empty_extreme(name, default), Ok)
1646        }
1647        other => Err(expected(
1648            "a distribution or nonempty list, range or string",
1649            other,
1650            name,
1651        )),
1652    }
1653}
1654
1655pub fn empty_extreme(name: &str, default: Option<&Value>) -> OpResult<Value> {
1656    default.cloned().ok_or_else(|| {
1657        OpError::fault(
1658            Fault::EmptyCollection,
1659            format!("`{name}` needs a nonempty collection or a default"),
1660        )
1661        .help(format!("use `{name}(xs, default: value)` to handle empty collections"))
1662    })
1663}
1664
1665fn finite_bound(x: f64, name: &str) -> OpResult<Value> {
1666    if !x.is_finite() {
1667        return Err(OpError::new(format!("`{name}` has no finite support bound")));
1668    }
1669    Ok(Value::Float(x))
1670}
1671
1672fn select_extreme(best: &mut Option<Value>, x: Value, want_max: bool, budget: &mut Budget) -> OpResult<()> {
1673    budget.work(1)?;
1674    ops::compare(&x, &x)?;
1675    let replace = match best {
1676        None => true,
1677        Some(b) => {
1678            let ord = ops::compare(&x, b)?;
1679            if want_max { ord.is_gt() } else { ord.is_lt() }
1680        }
1681    };
1682    if replace {
1683        *best = Some(x);
1684    }
1685    Ok(())
1686}
1687
1688pub fn extreme_count(n: &Value, budget: &mut Budget) -> OpResult<usize> {
1689    let n = integer(n, "the count", budget)?;
1690    if n.is_negative() {
1691        return Err(OpError::fault(Fault::DomainError, "the count must be nonnegative"));
1692    }
1693    Ok(n.to_u64().and_then(|n| usize::try_from(n).ok()).unwrap_or(usize::MAX))
1694}
1695
1696fn extremes(v: &Value, n: &Value, highest: bool, budget: &mut Budget) -> OpResult<Value> {
1697    let n = extreme_count(n, budget)?;
1698    let mut items = items(v, if highest { "highest" } else { "lowest" }, budget)?;
1699    sort(&mut items, highest, budget)?;
1700    items.truncate(n);
1701    Ok(Value::list(items))
1702}
1703
1704fn insert(coll: &Value, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
1705    match coll {
1706        Value::List(items) => {
1707            budget.collection(items.len() as u128 + 1)?;
1708            let mut items = items.to_vec();
1709            let i = as_index(a, items.len() as u128 + 1, budget)? as usize;
1710            items.insert(i, b.clone());
1711            Ok(Value::list(items))
1712        }
1713        Value::Map(m) => {
1714            budget.collection(m.len() as u128 + 1)?;
1715            let mut m = BTreeMap::clone(m);
1716            m.insert(a.clone(), b.clone());
1717            Ok(Value::map(m))
1718        }
1719        other => Err(expected("a list or a map", other, "insert")),
1720    }
1721}
1722
1723fn remove(coll: &Value, key: &Value, budget: &mut Budget) -> OpResult<Value> {
1724    match coll {
1725        Value::List(items) => {
1726            let mut items = items.to_vec();
1727            let i = as_index(key, items.len() as u128, budget)? as usize;
1728            items.remove(i);
1729            Ok(Value::list(items))
1730        }
1731        Value::Map(m) => {
1732            let mut m = BTreeMap::clone(m);
1733            m.remove(key);
1734            Ok(Value::map(m))
1735        }
1736        Value::Bag(b) => match b.without(key) {
1737            Some(rest) => Ok(Value::multiset(rest)),
1738            None => Err(OpError::fault(Fault::MissingKey, format!("{key:?} isn't in the bag"))),
1739        },
1740        other => Err(expected("a list, map or bag", other, "remove")),
1741    }
1742}
1743
1744/// A choice among options. Options that are distributions are mixed in, so
1745/// that a value drawn from the result is settled (docs/semantics.md, section 2).
1746fn one_of(v: &Value, budget: &mut Budget) -> OpResult<Value> {
1747    match v {
1748        Value::List(items) if !items.is_empty() => {
1749            budget.outcomes(items.len() as u128)?;
1750            let p = 1.0 / items.len() as f64;
1751            ops::combine(items.iter().map(|x| (x.clone(), p)).collect(), 0.0, budget)
1752        }
1753        Value::Range(lo, hi) if hi >= lo => {
1754            budget.outcomes(range_count(lo, hi)?)?;
1755            Ok(Dist::uniform(range_items(lo, hi, budget)?).into_value())
1756        }
1757        Value::Map(m) if !m.is_empty() => {
1758            let all_probs = m.values().all(|w| matches!(w, Value::Prob(_)));
1759            if !all_probs && m.values().any(|w| matches!(w, Value::Prob(_))) {
1760                return Err(
1761                    OpError::new("one_of can't mix probabilities and relative numeric weights")
1762                        .help("use prob(...) for every absolute probability, or numbers for every relative weight"),
1763                );
1764            }
1765            let mut pairs = Vec::new();
1766            for (k, w) in m.iter() {
1767                let w = match w {
1768                    Value::Bool(_) => None,
1769                    _ => w.as_f64().filter(|w| *w >= 0.0 && w.is_finite()),
1770                }
1771                .ok_or_else(|| {
1772                    OpError::new(format!(
1773                        "one_of needs weights that are numbers of 0 or more, found {w:?}"
1774                    ))
1775                })?;
1776                pairs.push((k.clone(), w));
1777            }
1778            let largest = pairs.iter().map(|(_, w)| *w).fold(0.0, f64::max);
1779            if largest == 0.0 {
1780                return Err(OpError::new("one_of needs at least one positive weight"));
1781            }
1782            // Absolute probabilities are checked before normalization. Relative
1783            // weights may have an unrepresentable sum even though each is valid.
1784            let total = crate::stats::sum(pairs.iter().map(|(_, w)| *w));
1785            if all_probs && (total - 1.0).abs() > 1e-9 {
1786                return Err(OpError::new(format!(
1787                    "the chances add up to {}, not 100%",
1788                    crate::value::fmt_prob(total)
1789                ))
1790                .help("use plain numbers for relative weights, like [\"a\": 3, \"b\": 1]"));
1791            }
1792            for (_, w) in &mut pairs {
1793                *w /= largest;
1794            }
1795            let total = crate::stats::sum(pairs.iter().map(|(_, w)| *w));
1796            ops::combine(pairs.into_iter().map(|(k, w)| (k, w / total)).collect(), 0.0, budget)
1797        }
1798        Value::Bag(b) => {
1799            let total: u128 = b.values().map(|n| *n as u128).sum();
1800            if total == 0 {
1801                return Err(OpError::fault(Fault::EmptyCollection, "the bag is empty"));
1802            }
1803            let pairs = b.iter().map(|(k, n)| (k.clone(), *n as f64 / total as f64)).collect();
1804            ops::combine(pairs, 0.0, budget)
1805        }
1806        Value::List(_) | Value::Range(..) | Value::Map(_) => Err(OpError::fault(
1807            Fault::EmptyCollection,
1808            "one_of needs at least one option",
1809        )),
1810        other => Err(expected("a list, range, map or bag", other, "one_of")),
1811    }
1812}
1813
1814fn bag(v: &Value, budget: &mut Budget) -> OpResult<Value> {
1815    let mut counts = BTreeMap::new();
1816    match v {
1817        Value::Map(m) => {
1818            for (k, n) in m.iter() {
1819                let n = integer(n, "bag count", budget)?;
1820                if n.is_negative() {
1821                    return Err(OpError::new("bag counts must be whole numbers of 0 or more"));
1822                }
1823                let n = n
1824                    .to_u64()
1825                    .ok_or_else(|| OpError::new("bag count exceeds the supported count range"))?;
1826                if n > 0 {
1827                    counts.insert(k.clone(), n);
1828                }
1829            }
1830        }
1831        Value::List(items) => {
1832            for item in items.iter() {
1833                *counts.entry(item.clone()).or_insert(0) += 1;
1834            }
1835        }
1836        other => return Err(expected("a map of counts or a list", other, "bag")),
1837    }
1838    Ok(Value::bag(counts))
1839}