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