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
542        | B::RunDate
543        | B::RunEnv
544        | B::Count
545        | B::Map
546        | B::Filter
547        | B::Reduce
548        | B::Print
549        | B::Roll
550        | B::Take => {
551            unreachable!("`{}` is handled by the interpreter", b.name())
552        }
553        B::P
554        | B::Mean
555        | B::Sd
556        | B::Variance
557        | B::Median
558        | B::MedianLow
559        | B::MedianHigh
560        | B::Quantile
561        | B::Support
562        | B::Cdf
563        | B::Pmf
564        | B::IterItems
565        | B::RepeatCount
566        | B::Pdf
567        | B::Settled => unreachable!("`{}` takes distributions as they are", b.name()),
568    }
569}
570
571/// Statistical queries take an explicit population, never an implicit point
572/// distribution. Also called before the interpreter's analytic-value guard so
573/// a scalar has the same type error in enumeration and sampling.
574pub fn check_query_input(b: Builtin, args: &[Value]) -> OpResult<()> {
575    use Builtin as B;
576    let expected = match b {
577        B::Minimum | B::Maximum => "a distribution or nonempty list, range or string",
578        B::P => "a boolean distribution (dist[bool])",
579        B::Pdf => "a continuous distribution",
580        B::Mean
581        | B::Variance
582        | B::Sd
583        | B::Median
584        | B::MedianLow
585        | B::MedianHigh
586        | B::Quantile
587        | B::Cdf
588        | B::Pmf
589        | B::Support => "a distribution, nonempty list or nonempty range",
590        _ => return Ok(()),
591    };
592    let v = &args[0];
593    let valid = match b {
594        B::Minimum | B::Maximum => matches!(
595            v,
596            Value::Dist(_)
597                | Value::Continuous(_)
598                | Value::Counts(_)
599                | Value::List(_)
600                | Value::Range(..)
601                | Value::Str(_)
602        ),
603        B::P => matches!(v, Value::Dist(_)),
604        B::Pdf => matches!(v, Value::Dist(_) | Value::Continuous(_) | Value::Joint(_)),
605        _ => matches!(
606            v,
607            Value::Dist(_)
608                | Value::Continuous(_)
609                | Value::Counts(_)
610                | Value::Joint(_)
611                | Value::List(_)
612                | Value::Range(..)
613        ),
614    };
615    if !valid {
616        let help = if b == B::P {
617            "use `report event` to measure a fact across worlds, or `prob(event)` to convert a bool to 0 or 1"
618        } else {
619            "use `report x` to summarize values across worlds, or put the model inside `simulate { ... }` to obtain a distribution"
620        };
621        return Err(OpError::new(format!(
622            "`{}` expects {expected}, found {}",
623            b.name(),
624            article(&v.kind())
625        ))
626        .help(help));
627    }
628    if !matches!(b, B::Minimum | B::Maximum) && matches!(v, Value::List(xs) if xs.is_empty()) {
629        return Err(OpError::fault(
630            Fault::EmptyCollection,
631            format!("`{}` needs a nonempty list", b.name()),
632        ));
633    }
634    Ok(())
635}
636
637/// Built-ins that receive distributions whole.
638pub fn call_raw(b: Builtin, args: &[Value], budget: &mut Budget) -> OpResult<Value> {
639    use Builtin as B;
640    check_query_input(b, args)?;
641    let v = &args[0];
642    if let Value::Range(lo, hi) = v {
643        if matches!(
644            b,
645            B::Mean
646                | B::Variance
647                | B::Sd
648                | B::Median
649                | B::MedianLow
650                | B::MedianHigh
651                | B::Quantile
652                | B::Cdf
653                | B::Pmf
654                | B::Support
655        ) {
656            return range_query(b, lo, hi, args.get(1), budget);
657        }
658    }
659    if let (B::Pmf | B::Cdf, Value::Dist(d)) = (b, v) {
660        if d.law.is_some() {
661            return listed_law_query(b, d, &args[1]);
662        }
663    }
664    if let Value::Counts(law) = v {
665        return law_stat(b, law, args, budget);
666    }
667    if matches!(b, B::P | B::Cdf | B::Pmf | B::Pdf) {
668        if let Value::Dist(d) = v {
669            if d.missing > 0.0 {
670                return Err(OpError::new(format!("`{}` cannot return an exact scalar probability while the distribution has unresolved mass", b.name()))
671                    .help("report a distribution comparison to retain probability bounds; scalar queries require a fully resolved distribution"));
672            }
673        }
674    }
675    if matches!(
676        b,
677        B::Mean
678            | B::Variance
679            | B::Sd
680            | B::Median
681            | B::MedianLow
682            | B::MedianHigh
683            | B::Quantile
684            | B::Cdf
685            | B::Pdf
686            | B::Pmf
687            | B::Support
688    ) && continuous_parts(v)
689    {
690        if matches!(b, B::Median | B::MedianLow | B::MedianHigh) {
691            let n = match v {
692                Value::Dist(d) => d.outcomes.len(),
693                _ => 1,
694            };
695            // Median bounds sort component supports before any CDF inversion.
696            budget.work((n as u64).saturating_mul(n.max(1).ilog2() as u64 + 1))?;
697        }
698        return continuous_query(b, args);
699    }
700    match b {
701        B::Minimum | B::Maximum => population_extreme(v, b == B::Maximum, None, budget),
702        B::P => probability_of(v),
703        B::Pdf => Err(OpError::new("pdf needs a continuous distribution")
704            .help("for a distribution whose outcomes can be listed, use `pmf`")),
705        // Within a world, the mean of outcomes of one draw is affine in it.
706        B::Mean if matches!(v, Value::List(xs) if xs.iter().any(|x| matches!(x, Value::Analytic(_)))) => {
707            let n = Value::Float(list(v, "mean")?.len() as f64);
708            ops::binary(probl_syntax::ast::BinOp::Div, &sum(v, budget)?, &n, budget)
709        }
710        B::Mean => {
711            let d = stat_dist(v, b.name(), budget)?;
712            if d.outcomes.iter().any(|(x, _)| matches!(x, Value::Date(_))) {
713                return date_mean(v, &d);
714            }
715            if d.outcomes.iter().any(|(x, _)| matches!(x, Value::Complex(_))) {
716                budget.work(d.outcomes.len() as u64)?;
717                let mut sum = Complex::new(0.0, 0.0)?;
718                let total = d.total();
719                for (x, p) in &d.outcomes {
720                    let x = complex_number(x, "mean")?;
721                    sum = sum.plus(x.times(Complex::new(p / total, 0.0)?)?)?;
722                }
723                Ok(Value::Complex(sum))
724            } else {
725                check_numeric(&d, "mean")?;
726                finite_float(d.mean().unwrap(), "mean")
727            }
728        }
729        B::Variance => {
730            numeric_dist(v, "variance", budget).and_then(|d| finite_float(d.variance().unwrap(), "variance"))
731        }
732        B::Sd => numeric_dist(v, "sd", budget).and_then(|d| finite_float(d.sd().unwrap(), "sd")),
733        B::Median | B::MedianLow | B::MedianHigh => median(v, b, budget),
734        B::Quantile => {
735            let q = to_prob(&args[1])?;
736            quantile(v, q, "quantile", budget)
737        }
738        B::Support => {
739            let d = stat_dist(v, b.name(), budget)?;
740            budget.collection(d.outcomes.len() as u128)?;
741            Ok(Value::list(d.outcomes.iter().map(|(x, _)| x.clone()).collect()))
742        }
743        B::Cdf => {
744            let d = stat_dist(v, b.name(), budget)?;
745            let mut p = crate::stats::Sum::default();
746            for (x, w) in &d.outcomes {
747                if statistical_compare(x, &args[1])?.is_le() {
748                    p.add(*w);
749                }
750            }
751            ops::computed_prob(p.value() / d.total(), "cdf")
752        }
753        B::Pmf => {
754            let d = stat_dist(v, b.name(), budget)?;
755            let p = crate::stats::sum(d.outcomes.iter().filter(|(x, _)| *x == args[1]).map(|(_, w)| *w));
756            ops::computed_prob(p / d.total(), "pmf")
757        }
758        B::IterItems => iter_items(v, budget),
759        B::RepeatCount => {
760            if v.is_uncertain() {
761                return Err(OpError::new(format!("`repeat` needs a number, not a {}", v.kind()))
762                    .help("draw a value first, like `let n ~ d6`, then `repeat n { … }`"));
763            }
764            let n = integer(v, "repeat count", budget)?;
765            if n.is_negative() {
766                return Err(OpError::fault(
767                    Fault::DomainError,
768                    "`repeat` needs a count of 0 or more",
769                ));
770            }
771            Ok(Value::Int(n.into_owned()))
772        }
773        B::Settled => match v {
774            v if v.is_uncertain() => Err(
775                OpError::new(format!("`match` needs a settled value, not a {}", v.kind()))
776                    .help("draw a value first, like `let x ~ d6`, and match on `x`"),
777            ),
778            other => Ok(other.clone()),
779        },
780        // Internal helpers that take their arguments as they are.
781        _ => call_plain(b, args, budget),
782    }
783}
784
785/// A query about a count law too broad to list, from its formulas. What
786/// needs its outcomes lists them, which fails on the outcome limit.
787fn law_stat(b: Builtin, law: &Counts, args: &[Value], budget: &mut Budget) -> OpResult<Value> {
788    use Builtin as B;
789    let int = |k: f64| {
790        Integer::from_f64(k)
791            .map(Value::Int)
792            .ok_or_else(|| OpError::limit("the outcome is too large for an int"))
793    };
794    let exact = matches!(law, Counts::Geometric { .. });
795    match b {
796        B::Pmf | B::Cdf if exact => law_query(b, law, &args[1]),
797        B::Mean => finite_float(law.mean(), "mean"),
798        B::Variance => finite_float(law.variance(), "variance"),
799        B::Sd => finite_float(libm::sqrt(law.variance()), "sd"),
800        // Their CDFs, without listing, aren't accurate this far yet.
801        B::Pmf | B::Cdf | B::Quantile | B::Median | B::MedianLow | B::MedianHigh if !exact => {
802            Err(OpError::unsupported(format!(
803                "`{}` of {law}, which has too many outcomes to list, isn't supported yet",
804                b.name()
805            ))
806            .help("its `mean`, `variance` and `sd` are exact; sampling draws from it directly"))
807        }
808        B::Quantile => {
809            let q = to_prob(&args[1])?;
810            if q <= 0.0 {
811                return law_stat(B::Minimum, law, args, budget);
812            }
813            if q >= 1.0 {
814                return Err(OpError::new("the 100% quantile of this distribution has no bound"));
815            }
816            int(law.quantile(q))
817        }
818        B::Median | B::MedianLow | B::MedianHigh => {
819            let lo = law.quantile(0.5);
820            // Half the probability exactly at or below `lo`: the next
821            // outcome is the high median.
822            let hi = if crate::stats::half_split(law.cdf(lo), 1.0) {
823                lo + 1.0
824            } else {
825                lo
826            };
827            match b {
828                B::MedianLow => int(lo),
829                B::MedianHigh => int(hi),
830                _ if lo == hi => int(lo),
831                _ => midpoint(&int(lo)?, &int(hi)?, budget),
832            }
833        }
834        B::Minimum => int(match law {
835            Counts::Geometric { .. } => 1.0,
836            _ => 0.0,
837        }),
838        _ => {
839            let v = ops::listed(&args[0], budget)?.into_owned();
840            let mut args = args.to_vec();
841            args[0] = v;
842            call_raw(b, &args, budget)
843        }
844    }
845}
846
847/// `pmf` or `cdf` of a count law's listed outcomes. What they leave out is
848/// the law's tail, less than 1e-18: unlike unresolved weight, it can't make
849/// a probability uncertain at any precision a float has.
850fn listed_law_query(b: Builtin, d: &Dist, x: &Value) -> OpResult<Value> {
851    let mut p = crate::stats::Sum::default();
852    for (outcome, w) in &d.outcomes {
853        let counts = if b == Builtin::Pmf {
854            outcome == x
855        } else {
856            statistical_compare(outcome, x)?.is_le()
857        };
858        if counts {
859            p.add(*w);
860        }
861    }
862    ops::computed_prob(p.value(), b.name())
863}
864
865/// `pmf` or `cdf` of a count law from its formulas. As with its outcomes,
866/// `pmf` is of an int: a float is none of them.
867fn law_query(b: Builtin, law: &Counts, x: &Value) -> OpResult<Value> {
868    if b == Builtin::Pmf {
869        let p = match x {
870            Value::Int(_) => law.pmf(x.as_f64().unwrap_or(f64::INFINITY)),
871            _ => 0.0,
872        };
873        return ops::computed_prob(p, "pmf");
874    }
875    // The comparison an outcome would make with `x`, for its errors.
876    statistical_compare(&Value::Int(0.into()), x)?;
877    let t = x
878        .as_f64()
879        .ok_or_else(|| OpError::new(format!("`cdf` needs a number, found {}", article(&x.kind()))))?;
880    ops::computed_prob(law.cdf(t), "cdf")
881}
882
883/// `P(d)`: query an explicitly constructed boolean distribution.
884fn probability_of(v: &Value) -> OpResult<Value> {
885    match v {
886        Value::Dist(d) if d.truth().is_some() => ops::computed_prob(d.truth().unwrap().0 / d.total(), "P"),
887        _ => Err(OpError::new(format!(
888            "P needs a boolean distribution (dist[bool]), found {}",
889            article(&v.kind())
890        ))
891        .help("compare a distribution, like `P(d6 > 4)`; use `report event` for a fact across worlds")),
892    }
893}
894
895/// The parameters of `binomial`, `poisson` or `geometric`, checked; `None`
896/// for another built-in, or when an argument is a distribution.
897pub fn counts(b: Builtin, args: &[Value], budget: &mut Budget) -> OpResult<Option<Counts>> {
898    use Builtin as B;
899    if args.iter().any(Value::is_uncertain) {
900        return Ok(None);
901    }
902    Ok(Some(match b {
903        B::Binomial => {
904            let n = whole(&args[0], "binomial's number of trials", budget)?;
905            if n < 0 {
906                return Err(OpError::fault(
907                    Fault::DomainError,
908                    "binomial needs a number of trials of 0 or more",
909                ));
910            }
911            Counts::Binomial {
912                n: n as u64,
913                p: to_prob(&args[1])?,
914            }
915        }
916        B::Poisson => {
917            let rate = number(&args[0], "poisson")?;
918            if rate < 0.0 || !rate.is_finite() {
919                return Err(OpError::fault(Fault::DomainError, "poisson needs a rate of 0 or more"));
920            }
921            if rate > 1e15 {
922                return Err(OpError::new("poisson's rate is too large to count exactly")
923                    .help("above 10¹⁵, use a normal distribution with the same mean and variance"));
924            }
925            Counts::Poisson { rate }
926        }
927        B::Geometric => {
928            let p = to_prob(&args[0])?;
929            if p <= 0.0 {
930                return Err(OpError::fault(
931                    Fault::DomainError,
932                    "geometric needs a chance of success above 0%",
933                ));
934            }
935            Counts::Geometric { p }
936        }
937        _ => return Ok(None),
938    }))
939}
940
941fn continuous(family: OpResult<Family>) -> OpResult<Value> {
942    family.map(|f| Value::Continuous(Arc::new(f)))
943}
944
945/// Whether a value is, or mixes in, a continuous distribution.
946fn continuous_parts(v: &Value) -> bool {
947    match v {
948        Value::Continuous(_) | Value::Joint(_) => true,
949        Value::Dist(d) => d
950            .outcomes
951            .iter()
952            .any(|(x, _)| matches!(x, Value::Continuous(_) | Value::Joint(_))),
953        _ => false,
954    }
955}
956
957/// Questions about a continuous distribution, or a mixture that includes
958/// one, answered from the formulas (docs/semantics.md, section 13).
959fn continuous_query(b: Builtin, args: &[Value]) -> OpResult<Value> {
960    use Builtin as B;
961    // PMF is a typed-outcome query even when continuous components are mixed
962    // in. Converting atoms to f64 here would collapse distinct outcomes.
963    if b == B::Pmf {
964        return match &args[0] {
965            Value::Dist(d) => {
966                let at = crate::stats::sum(
967                    d.outcomes
968                        .iter()
969                        .filter(|(x, _)| !matches!(x, Value::Continuous(_)) && *x == args[1])
970                        .map(|(_, p)| *p),
971                );
972                ops::computed_prob(at / d.total(), "pmf")
973            }
974            Value::Continuous(_) => ops::computed_prob(0.0, "pmf"),
975            Value::Joint(_) => Err(OpError::new(
976                "`pmf` needs a distribution whose outcomes can be listed, not one of continuous outcomes",
977            )),
978            _ => unreachable!("checked continuous input"),
979        };
980    }
981    // The marginal of a `simulate` distribution of continuous outcomes.
982    let joint = |d: &Dist, scale: f64| match crate::joint::mixture(d) {
983        Some(m) => Ok(m.parts.into_iter().map(move |(part, p)| (part, p * scale))),
984        None => Err(OpError::unsupported(format!(
985            "`{}` of a distribution of continuous outcomes that aren't all numbers isn't supported",
986            b.name()
987        ))),
988    };
989    let parts = match &args[0] {
990        Value::Continuous(f) => vec![(Part::Continuous(**f), 1.0)],
991        Value::Joint(d) => joint(d, 1.0)?.collect(),
992        Value::Dist(d) => {
993            let mut parts = Vec::with_capacity(d.outcomes.len());
994            for (x, p) in &d.outcomes {
995                match x {
996                    Value::Continuous(f) => parts.push((Part::Continuous(**f), *p)),
997                    Value::Joint(j) => parts.extend(joint(j, *p)?),
998                    other => parts.push((Part::Point(number(other, b.name())?), *p)),
999                }
1000            }
1001            parts
1002        }
1003        _ => unreachable!("checked by `continuous_parts`"),
1004    };
1005    let m = Mixture { parts };
1006    match b {
1007        B::Mean => finite_float(m.mean(), "mean"),
1008        B::Variance => finite_float(m.variance(), "variance"),
1009        B::Sd => finite_float(m.sd(), "sd"),
1010        B::Median | B::MedianLow | B::MedianHigh => {
1011            let (lo, hi) = m.median_bounds();
1012            finite_float(
1013                match b {
1014                    B::MedianLow => lo,
1015                    B::MedianHigh => hi,
1016                    _ => crate::stats::midpoint(lo, hi),
1017                },
1018                b.name(),
1019            )
1020        }
1021        B::Quantile => finite_float(m.quantile(to_prob(&args[1])?), "quantile"),
1022        B::Cdf => ops::computed_prob(m.cdf(number(&args[1], "cdf")?), "cdf"),
1023        B::Pdf => {
1024            let x = number(&args[1], "pdf")?;
1025            let mut density = 0.0;
1026            for (part, p) in &m.parts {
1027                match part {
1028                    Part::Continuous(f) => density += p * f.pdf(x),
1029                    Part::Analytic(a) if !a.has_atoms() && a.piecewise_affine() => density += p * a.pdf(x),
1030                    Part::Analytic(a) if !a.has_atoms() => {
1031                        return Err(crate::analytic::unsupported("the density of this outcome"));
1032                    }
1033                    Part::Analytic(_) | Part::Point(_) => {
1034                        return Err(OpError::new(
1035                            "pdf needs a continuous distribution, without single values mixed in",
1036                        ));
1037                    }
1038                }
1039            }
1040            finite_float(density, "pdf")
1041        }
1042        _ => Err(OpError::new(format!(
1043            "`{}` needs a distribution whose outcomes can be listed, not a continuous one",
1044            b.name()
1045        ))),
1046    }
1047}
1048
1049/// Unit-step integer ranges are uniform populations; scalar queries do not
1050/// allocate their support. Rank arithmetic remains exact for arbitrary integers.
1051fn range_query(b: Builtin, lo: &Integer, hi: &Integer, arg: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1052    use Builtin as B;
1053    if hi < lo {
1054        return Err(OpError::fault(
1055            Fault::EmptyCollection,
1056            format!("`{}` needs a nonempty range", b.name()),
1057        ));
1058    }
1059    budget.integer_work(lo, hi, false)?;
1060    let n = range_len(lo, hi)?;
1061    budget.integer_allocation(n.bits(), 1)?;
1062    let two = Integer::from(2);
1063    match b {
1064        B::Support => Ok(Value::list(range_items(lo, hi, budget)?)),
1065        B::Mean => lo.add(hi)?.ratio(&two).map(Value::Float).ok_or_else(|| {
1066            OpError::fault(
1067                Fault::NumericOverflow,
1068                "`mean` gave a result that isn't a finite number",
1069            )
1070        }),
1071        B::Median | B::MedianLow | B::MedianHigh => {
1072            budget.integer_work(&n, &two, true)?;
1073            let (half, odd) = n.div_mod(&two)?;
1074            let high = Value::Int(lo.add(&half)?);
1075            let low = if odd.is_zero() {
1076                Value::Int(lo.add(&half)?.sub(&Integer::ONE)?)
1077            } else {
1078                high.clone()
1079            };
1080            match b {
1081                B::MedianLow => Ok(low),
1082                B::MedianHigh => Ok(high),
1083                _ => midpoint(&low, &high, budget),
1084            }
1085        }
1086        B::Sd | B::Variance => {
1087            let quarter = n.ratio(&Integer::from(4)).unwrap_or(f64::INFINITY);
1088            let inv = Integer::ONE.ratio(&n).unwrap_or(0.0);
1089            let sd = quarter * (4.0 / libm::sqrt(12.0)) * libm::sqrt((1.0 - inv) * (1.0 + inv));
1090            finite_float(if b == B::Sd { sd } else { sd * sd }, b.name())
1091        }
1092        B::Quantile => {
1093            let q = to_prob(arg.expect("quantile argument"))?;
1094            if q == 0.0 {
1095                return Ok(Value::Int(lo.clone()));
1096            }
1097            budget.work(n.bits().div_ceil(64).saturating_mul(17))?;
1098            let rank = if let Some(count) = n.to_u64().filter(|n| *n <= 1 << 53) {
1099                // Match the rounded cumulative weights of a finite uniform
1100                // population, including decimal boundaries such as 10% of 10.
1101                let mut weight = 1.0 / count as f64;
1102                let total = weight * count as f64;
1103                if total != 1.0 {
1104                    weight *= 1.0 / total;
1105                }
1106                let target = q * (weight * count as f64);
1107                let (mut left, mut right) = (1, count);
1108                while left < right {
1109                    budget.work(1)?;
1110                    let mid = left + (right - left) / 2;
1111                    if mid as f64 * weight >= target {
1112                        right = mid;
1113                    } else {
1114                        left = mid + 1;
1115                    }
1116                }
1117                Integer::from(left)
1118            } else {
1119                // Above exact f64 integer precision, keep the rank arithmetic
1120                // exact instead of rounding the population size to a float.
1121                n.probability_rank(q)?
1122            };
1123            Ok(Value::Int(lo.add(&rank.sub(&Integer::ONE)?)?))
1124        }
1125        B::Pmf => {
1126            let present = matches!(arg, Some(Value::Int(k)) if lo <= k && k <= hi);
1127            ops::computed_prob(
1128                if present {
1129                    Integer::ONE.ratio(&n).unwrap_or(0.0)
1130                } else {
1131                    0.0
1132                },
1133                "pmf",
1134            )
1135        }
1136        B::Cdf => {
1137            let x = arg.expect("cdf argument");
1138            if ops::compare(x, &Value::Int(lo.clone()))?.is_lt() {
1139                return ops::computed_prob(0.0, "cdf");
1140            }
1141            if ops::compare(x, &Value::Int(hi.clone()))?.is_ge() {
1142                return ops::computed_prob(1.0, "cdf");
1143            }
1144            let k = match x {
1145                Value::Int(k) => k.clone(),
1146                _ => Integer::from_f64(number(x, "cdf")?.floor()).expect("finite numeric bound"),
1147            };
1148            let count = k.sub(lo)?.add(&Integer::ONE)?;
1149            budget.integer_work(&count, &n, true)?;
1150            ops::computed_prob(count.ratio(&n).unwrap_or(0.0), "cdf")
1151        }
1152        _ => unreachable!("checked range query"),
1153    }
1154}
1155
1156/// A list is an empirical distribution: each element has equal weight,
1157/// including repetitions. Its elements remain values, never implicit draws.
1158fn stat_dist<'a>(v: &'a Value, what: &str, budget: &mut Budget) -> OpResult<Cow<'a, Dist>> {
1159    match v {
1160        Value::Dist(d) => {
1161            budget.work(d.outcomes.len() as u64)?;
1162            Ok(Cow::Borrowed(d))
1163        }
1164        Value::List(xs) if !xs.is_empty() => {
1165            budget.collection(xs.len() as u128)?;
1166            let comparisons = (xs.len().ilog2() + 1) as u64;
1167            budget.work((xs.len() as u64).saturating_mul(comparisons))?;
1168            for x in xs.iter() {
1169                match x {
1170                    Value::Int(n) => {
1171                        budget.integer_work(n, &Integer::ONE, false)?;
1172                    }
1173                    Value::Str(s) => {
1174                        budget.string_work(s)?;
1175                    }
1176                    x if x.is_uncertain() => {
1177                        return Err(OpError::new(format!(
1178                            "`{what}` needs list elements that are values, found {}",
1179                            article(&x.kind())
1180                        ))
1181                        .help("draw the elements first, or explicitly build a mixture with `one_of`"));
1182                    }
1183                    _ => {}
1184                }
1185            }
1186            Ok(Cow::Owned(Dist::uniform(xs.to_vec())))
1187        }
1188        _ => Err(OpError::new(format!(
1189            "`{what}` needs a distribution or a nonempty list"
1190        ))),
1191    }
1192}
1193
1194fn numeric_dist<'a>(v: &'a Value, what: &str, budget: &mut Budget) -> OpResult<Cow<'a, Dist>> {
1195    let d = stat_dist(v, what, budget)?;
1196    check_numeric(&d, what)?;
1197    Ok(d)
1198}
1199
1200fn check_numeric(d: &Dist, what: &str) -> OpResult<()> {
1201    for (x, _) in &d.outcomes {
1202        if !matches!(x, Value::Int(_) | Value::Float(_) | Value::Prob(_)) {
1203            return Err(OpError::new(format!(
1204                "`{what}` needs real numeric elements, found {}",
1205                article(&x.kind())
1206            )));
1207        }
1208        if x.as_f64().is_none_or(|x| !x.is_finite()) {
1209            return Err(OpError::new(format!(
1210                "`{what}` needs numbers representable as finite floats"
1211            )));
1212        }
1213    }
1214    if d.outcomes.is_empty() {
1215        return Err(OpError::new(format!("`{what}` needs at least one resolved outcome")));
1216    }
1217    Ok(())
1218}
1219
1220fn statistical_compare(a: &Value, b: &Value) -> OpResult<std::cmp::Ordering> {
1221    match (a, b) {
1222        (Value::Bool(a), Value::Bool(b)) => Ok(a.cmp(b)),
1223        _ => ops::compare(a, b),
1224    }
1225}
1226
1227/// Public ordering is separate from the typed total order used for storage.
1228/// Never put this reordered vector back into a Dist or a world/cache key.
1229fn ordered_outcomes(d: &Dist, budget: &mut Budget) -> OpResult<Vec<(Value, f64)>> {
1230    budget.collection(d.outcomes.len() as u128)?;
1231    for (x, _) in &d.outcomes {
1232        budget.work(1)?;
1233        statistical_compare(x, x)?;
1234    }
1235    crate::ordering::reserve_sort(d.outcomes.len(), budget)?;
1236    let mut outcomes = d.outcomes.clone();
1237    crate::ordering::try_sort_by(&mut outcomes, |(a, _), (b, _)| {
1238        budget.work(1)?;
1239        statistical_compare(a, b)
1240    })?;
1241    Ok(outcomes)
1242}
1243
1244fn quantile(v: &Value, q: f64, what: &str, budget: &mut Budget) -> OpResult<Value> {
1245    let d = stat_dist(v, what, budget)?;
1246    let outcomes = ordered_outcomes(&d, budget)?;
1247    crate::stats::quantile(&outcomes, q)
1248        .cloned()
1249        .ok_or_else(|| OpError::new(format!("`{what}` needs at least one resolved outcome")))
1250}
1251
1252fn median(v: &Value, b: Builtin, budget: &mut Budget) -> OpResult<Value> {
1253    let d = stat_dist(v, b.name(), budget)?;
1254    for (x, _) in &d.outcomes {
1255        budget.work(1)?;
1256        // Validate even singletons. Being sortable internally does not give a
1257        // record, complex number or recipe a mathematical ordering.
1258        statistical_compare(x, x)?;
1259        if b == Builtin::Median && !matches!(x, Value::Int(_) | Value::Float(_) | Value::Prob(_) | Value::Date(_)) {
1260            return Err(OpError::new(format!(
1261                "`median` needs real numeric or date elements, found {}",
1262                article(&x.kind())
1263            ))
1264            .help("use `median_low` or `median_high` for ordered values such as strings"));
1265        }
1266    }
1267    let outcomes = ordered_outcomes(&d, budget)?;
1268    let (lo, hi) = crate::stats::median_bounds(&outcomes)
1269        .ok_or_else(|| OpError::new(format!("`{}` needs at least one resolved outcome", b.name())))?;
1270    match b {
1271        Builtin::MedianLow => Ok(lo.clone()),
1272        Builtin::MedianHigh => Ok(hi.clone()),
1273        _ if lo == hi => Ok(lo.clone()),
1274        _ => midpoint(lo, hi, budget),
1275    }
1276}
1277
1278fn date_mean(v: &Value, d: &Dist) -> OpResult<Value> {
1279    // Lists permit exact integer arithmetic, so rounding cannot depend on
1280    // repeated values, list order, or a date's distance from the epoch.
1281    if let Value::List(xs) = v {
1282        let total = xs
1283            .iter()
1284            .try_fold(0i128, |sum, x| Ok::<_, OpError>(sum + date_value(x, "mean")? as i128))?;
1285        let n = xs.len() as i128;
1286        let day = total.div_euclid(n) + i128::from(2 * total.rem_euclid(n) > n);
1287        return Ok(Value::Date(day as i32));
1288    }
1289    let origin = date_value(&d.outcomes[0].0, "mean")?;
1290    let mut offsets = crate::stats::Sum::default();
1291    let mut weights = crate::stats::Sum::default();
1292    for (x, w) in &d.outcomes {
1293        offsets.add((date_value(x, "mean")? - origin) as f64 * w);
1294        weights.add(*w);
1295    }
1296    let days = offsets.value() / weights.value();
1297    let floor = days.floor();
1298    let tie_error = 4.0 * f64::EPSILON * days.abs().max(1.0);
1299    let rounded = floor + f64::from(days - floor > 0.5 + tie_error);
1300    date_result(
1301        i32::try_from(origin as i64 + rounded as i64)
1302            .ok()
1303            .filter(|d| dates::valid(*d)),
1304    )
1305}
1306
1307/// Preserve exact integral midpoints, including bigints beyond float range.
1308pub(crate) fn midpoint(a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
1309    if let (Value::Date(a), Value::Date(b)) = (a, b) {
1310        return Ok(Value::Date(((*a as i64 + *b as i64).div_euclid(2)) as i32));
1311    }
1312    if let (Value::Int(a), Value::Int(b)) = (a, b) {
1313        budget.integer_work(a, b, false)?;
1314        let sum = a.add(b)?;
1315        budget.integer_allocation(sum.bits(), 1)?;
1316        let two = Integer::from(2);
1317        budget.integer_work(&sum, &two, true)?;
1318        let (whole, remainder) = sum.div_mod(&two)?;
1319        if remainder.is_zero() {
1320            return Ok(Value::Int(whole));
1321        }
1322        return sum.ratio(&two).map(Value::Float).ok_or_else(|| {
1323            OpError::fault(
1324                Fault::NumericOverflow,
1325                "median's fractional midpoint is too large for a finite float",
1326            )
1327        });
1328    }
1329    let (a, b) = (number(a, "median")?, number(b, "median")?);
1330    // Same-sign subtraction and opposite-sign addition avoid overflow; this
1331    // also preserves equal subnormal values instead of halving both to zero.
1332    finite_float(crate::stats::midpoint(a, b), "median")
1333}
1334
1335fn iter_items(v: &Value, budget: &mut Budget) -> OpResult<Value> {
1336    match v {
1337        Value::List(_) | Value::Range(..) => Ok(v.clone()),
1338        Value::Map(m) => Ok(Value::list(
1339            m.iter().map(|(k, x)| Value::list(vec![k.clone(), x.clone()])).collect(),
1340        )),
1341        Value::Bag(b) => {
1342            let n: u128 = b.values().map(|n| *n as u128).sum();
1343            budget.collection(n)?;
1344            Ok(Value::list(
1345                b.iter()
1346                    .flat_map(|(k, n)| std::iter::repeat_n(k.clone(), *n as usize))
1347                    .collect(),
1348            ))
1349        }
1350        Value::Str(s) => crate::text::chars(s, budget).map(Value::list),
1351        v if v.is_uncertain() => Err(OpError::new(format!("can't loop over a {}", v.kind()))
1352            .help("draw a value first with `~`, or loop over `support(…)`")),
1353        other => Err(OpError::new(format!("can't loop over {}", article(&other.kind())))),
1354    }
1355}
1356
1357fn date_count(v: &Value, func: &str, budget: &mut Budget) -> OpResult<i64> {
1358    integer(v, func, budget)?
1359        .to_i64()
1360        .ok_or_else(|| OpError::fault(Fault::NumericOverflow, "date out of range"))
1361}
1362
1363fn date_value(v: &Value, func: &str) -> OpResult<i32> {
1364    match v {
1365        Value::Date(d) if dates::valid(*d) => Ok(*d),
1366        Value::Date(_) => Err(OpError::fault(Fault::NumericOverflow, "date out of range")),
1367        _ => Err(expected("a date", v, func)),
1368    }
1369}
1370
1371fn date_result(date: Option<i32>) -> OpResult<Value> {
1372    date.map(Value::Date)
1373        .ok_or_else(|| OpError::fault(Fault::NumericOverflow, "date out of range (0001-01-01..9999-12-31)"))
1374}
1375
1376fn holiday_calendar(v: Option<&Value>, func: &str, budget: &mut Budget) -> OpResult<Vec<i32>> {
1377    let Some(v) = v else {
1378        return Ok(Vec::new());
1379    };
1380    let values = list(v, func)?;
1381    budget.collection(values.len() as u128)?;
1382    let log = (values.len() as u64).checked_ilog2().unwrap_or(0) as u64 + 1;
1383    budget.work((values.len() as u64).saturating_mul(log))?;
1384    let mut holidays = Vec::new();
1385    for value in values.iter() {
1386        let d = date_value(value, func)?;
1387        if dates::weekday(d) < 5 {
1388            holidays.push(d);
1389        }
1390    }
1391    holidays.sort_unstable();
1392    holidays.dedup();
1393    Ok(holidays)
1394}
1395
1396fn expected(what: &str, v: &Value, func: &str) -> OpError {
1397    OpError::new(format!("`{func}` needs {what}, found {}", article(&v.kind())))
1398}
1399
1400fn number(v: &Value, func: &str) -> OpResult<f64> {
1401    match v {
1402        Value::Bool(_) => Err(expected("a number", v, func)),
1403        Value::Continuous(_) => {
1404            Err(expected("a number", v, func).help("draw a value first, like `let x ~ normal(0, 1)`"))
1405        }
1406        Value::Int(n) => n.to_f64().ok_or_else(|| {
1407            OpError::fault(
1408                Fault::NumericOverflow,
1409                format!("`{func}` needs an integer that fits in a finite float"),
1410            )
1411        }),
1412        _ => v.as_f64().ok_or_else(|| expected("a number", v, func)),
1413    }
1414}
1415
1416fn complex_number(v: &Value, func: &str) -> OpResult<Complex> {
1417    v.as_complex()
1418        .ok_or_else(|| expected("a finite real or complex number", v, func))
1419}
1420
1421fn finite_float(x: f64, func: &str) -> OpResult<Value> {
1422    if x.is_finite() {
1423        Ok(Value::Float(x))
1424    } else {
1425        Err(OpError::new(format!(
1426            "`{func}` gave a result that isn't a finite number"
1427        )))
1428    }
1429}
1430
1431fn whole(v: &Value, what: &str, budget: &mut Budget) -> OpResult<i64> {
1432    integer(v, what, budget)?
1433        .to_i64()
1434        .ok_or_else(|| OpError::new(format!("{what} is outside the supported count range")))
1435}
1436
1437fn nonnegative_int<'a>(v: &'a Value, func: &str, budget: &mut Budget) -> OpResult<Cow<'a, Integer>> {
1438    let n = integer(v, func, budget)?;
1439    if n.is_negative() {
1440        Err(OpError::fault(
1441            Fault::DomainError,
1442            format!("`{func}` needs nonnegative integers"),
1443        ))
1444    } else {
1445        Ok(n)
1446    }
1447}
1448
1449fn choose(n: &Integer, k: &Integer, budget: &mut Budget) -> OpResult<Value> {
1450    if k > n {
1451        return Ok(Value::Int(Integer::ZERO));
1452    }
1453    let k = k.min(&n.sub(k)?).clone();
1454    let steps = k
1455        .to_u64()
1456        .ok_or_else(|| OpError::limit("choose needs too many iterations"))?;
1457    budget.work(steps)?;
1458    let offset = n.sub(&k)?;
1459    let mut result = Integer::ONE;
1460    for i in 1..=steps {
1461        let numerator = offset.add(&i.into())?;
1462        let divisor = Integer::from(i);
1463        // Cancel first: an intermediate product must not exceed the integer
1464        // ceiling when the final binomial coefficient fits.
1465        let d = gcd(numerator.clone(), divisor.clone(), budget)?;
1466        let numerator = numerator.div_mod(&d)?.0;
1467        let divisor = divisor.div_mod(&d)?.0;
1468        budget.integer_work(&result, &divisor, true)?;
1469        result = result.div_mod(&divisor)?.0;
1470        budget.integer_work(&result, &numerator, true)?;
1471        result = result.mul(&numerator)?;
1472        budget.integer_bits(result.bits())?;
1473    }
1474    Ok(Value::Int(result))
1475}
1476
1477fn factorial(n: &Integer, budget: &mut Budget) -> OpResult<Value> {
1478    let n = n
1479        .to_u64()
1480        .ok_or_else(|| OpError::limit("factorial exceeds the integer size limit"))?;
1481    // n! contains at least n/2 factors of n/2 or more.
1482    let half = n / 2;
1483    let lower_bits = half.saturating_mul(63u64.saturating_sub(u64::from(half.leading_zeros())));
1484    budget.integer_bits(lower_bits)?;
1485    budget.work(n)?;
1486    let mut result = Integer::ONE;
1487    for i in 2..=n {
1488        let factor = Integer::from(i);
1489        budget.integer_work(&result, &factor, true)?;
1490        result = result.mul(&factor)?;
1491        budget.integer_bits(result.bits())?;
1492    }
1493    Ok(Value::Int(result))
1494}
1495
1496fn gcd(mut a: Integer, mut b: Integer, budget: &mut Budget) -> OpResult<Integer> {
1497    while !b.is_zero() {
1498        budget.integer_work(&a, &b, true)?;
1499        let r = a.div_mod(&b)?.1;
1500        a = b;
1501        b = r;
1502    }
1503    Ok(a)
1504}
1505
1506fn euler_phi(n: &Integer, budget: &mut Budget) -> OpResult<Value> {
1507    if n.is_zero() {
1508        return Err(OpError::fault(
1509            Fault::DomainError,
1510            "`euler_phi` needs a positive integer",
1511        ));
1512    }
1513    let mut result = n.clone();
1514    for (p, _) in primes::factorize(n, &mut |units| budget.work(units))? {
1515        budget.integer_work(&result, &p, true)?;
1516        result = result.sub(&result.div_mod(&p)?.0)?;
1517    }
1518    Ok(Value::Int(result))
1519}
1520
1521fn prime_factors(n: &Integer, budget: &mut Budget) -> OpResult<Value> {
1522    if n.is_zero() || n.is_negative() {
1523        return Err(OpError::fault(
1524            Fault::DomainError,
1525            "`prime_factors` needs a positive integer",
1526        ));
1527    }
1528    let factors = primes::factorize(n, &mut |units| budget.work(units))?;
1529    let count: u64 = factors.iter().map(|&(_, k)| u64::from(k)).sum();
1530    budget.collection(u128::from(count))?;
1531    budget.work(count)?;
1532    let mut items = Vec::with_capacity(count as usize);
1533    for (p, k) in factors {
1534        budget.integer_allocation(p.bits(), u64::from(k))?;
1535        items.extend(std::iter::repeat_n(Value::Int(p), k as usize));
1536    }
1537    Ok(Value::list(items))
1538}
1539
1540fn divisors(n: &Integer, budget: &mut Budget) -> OpResult<Value> {
1541    if n.is_zero() || n.is_negative() {
1542        return Err(OpError::fault(
1543            Fault::DomainError,
1544            "`divisors` needs a positive integer",
1545        ));
1546    }
1547    let factors = primes::factorize(n, &mut |units| budget.work(units))?;
1548    let count = factors
1549        .iter()
1550        .fold(1u128, |count, &(_, k)| count.saturating_mul(u128::from(k) + 1));
1551    budget.collection(count)?;
1552    // Each divisor takes a product, and sorting compares each about log₂ count times.
1553    let count = count as u64;
1554    let w = n.bits().div_ceil(64).max(1);
1555    budget.work(count.saturating_mul(w * w + w * u64::from(64 - count.leading_zeros())))?;
1556    budget.integer_allocation(n.bits(), count)?;
1557    Ok(Value::list(
1558        primes::divisors(&factors).into_iter().map(Value::Int).collect(),
1559    ))
1560}
1561
1562fn primes_between(lo: &Integer, hi: &Integer, budget: &mut Budget) -> OpResult<Value> {
1563    let mut primes = primes::Primes::new(lo, hi);
1564    let mut items = Vec::new();
1565    while let Some(p) = primes.next(&mut |units| budget.work(units))? {
1566        budget.collection(items.len() as u128 + 1)?;
1567        budget.integer_allocation(p.bits(), 1)?;
1568        items.push(Value::Int(p));
1569    }
1570    Ok(Value::list(items))
1571}
1572
1573/// The work of Euclid's algorithm modulo `m`: about twice its bit count in
1574/// steps, each a division of numbers below it.
1575fn euclid_work(m: &Integer) -> u64 {
1576    m.bits().max(1).saturating_mul(2 * m.bits().div_ceil(64).max(1))
1577}
1578
1579fn mod_pow(b: &Value, e: &Value, m: &Value, budget: &mut Budget) -> OpResult<Value> {
1580    let (b, e, m) = (
1581        integer(b, "mod_pow", budget)?,
1582        integer(e, "mod_pow", budget)?,
1583        integer(m, "mod_pow", budget)?,
1584    );
1585    if m.is_zero() {
1586        return Err(OpError::fault(
1587            Fault::DivisionByZero,
1588            "`mod_pow` needs a nonzero modulus",
1589        ));
1590    }
1591    budget.integer_work(&b, &m, true)?;
1592    if e.is_negative() {
1593        budget.work(euclid_work(&m))?;
1594    }
1595    // A squaring for each bit of the exponent, and a product for some.
1596    let w = m.bits().div_ceil(64).max(1);
1597    budget.work(e.bits().saturating_mul(2 * w * w))?;
1598    match b.mod_pow(&e, &m)? {
1599        Some(power) => Ok(Value::Int(power)),
1600        None => Err(OpError::fault(
1601            Fault::DomainError,
1602            "`mod_pow` with a negative exponent needs a base coprime to the modulus",
1603        )),
1604    }
1605}
1606
1607fn text(v: &Value, func: &str) -> OpResult<Arc<str>> {
1608    match v {
1609        Value::Str(s) => Ok(s.clone()),
1610        other => Err(expected("a string", other, func)),
1611    }
1612}
1613
1614fn list<'a>(v: &'a Value, func: &str) -> OpResult<&'a [Value]> {
1615    match v {
1616        Value::List(items) => Ok(items),
1617        other => Err(expected("a list", other, func)),
1618    }
1619}
1620
1621fn range_items(lo: &Integer, hi: &Integer, budget: &mut Budget) -> OpResult<Vec<Value>> {
1622    let n = range_count(lo, hi)?;
1623    budget.collection(n)?;
1624    let n = usize::try_from(n).map_err(|_| OpError::limit("range has too many elements"))?;
1625    budget.integer_allocation(lo.bits().max(hi.bits()), n as u64)?;
1626    budget.work((n as u64).saturating_mul(lo.bits().max(hi.bits()).div_ceil(64).max(1)))?;
1627    (0..n)
1628        .map(|i| lo.add(&Integer::from(i)).map(Value::Int).map_err(OpError::from))
1629        .collect()
1630}
1631
1632/// The elements of a list-like value, checking that a range isn't too long
1633/// to spell out.
1634pub fn items(v: &Value, func: &str, budget: &mut Budget) -> OpResult<Vec<Value>> {
1635    match v {
1636        Value::List(items) => {
1637            budget.collection(items.len() as u128)?;
1638            budget.work(items.len() as u64)?;
1639            Ok(items.to_vec())
1640        }
1641        Value::Range(lo, hi) => range_items(lo, hi, budget),
1642        Value::Str(s) => crate::text::chars(s, budget),
1643        other => Err(expected("a list, range or string", other, func)),
1644    }
1645}
1646
1647fn sort(items: &mut [Value], descending: bool, budget: &mut Budget) -> OpResult<()> {
1648    // Even a singleton must have a language ordering; internal storage order
1649    // doesn't make complex values, records or distributions sortable.
1650    for x in items.iter() {
1651        budget.work(1)?;
1652        ops::compare(x, x)?;
1653    }
1654    crate::ordering::reserve_sort(items.len(), budget)?;
1655    crate::ordering::try_sort_by(items, |x, y| {
1656        budget.work(1)?;
1657        ops::compare(x, y).map(|c| if descending { c.reverse() } else { c })
1658    })
1659}
1660
1661/// Real inputs keep their real domains; an explicit complex input requests the
1662/// principal complex extension. Distribution lifting happens in the interpreter.
1663fn elementary1(
1664    v: &Value,
1665    func: &str,
1666    complex: fn(Complex) -> OpResult<Complex>,
1667    real: impl Fn(f64) -> Option<f64>,
1668) -> OpResult<Value> {
1669    match v {
1670        Value::Complex(z) => complex(*z).map(Value::Complex).map_err(|_| {
1671            OpError::fault(
1672                Fault::DomainError,
1673                format!("`{func}` isn't defined for {v} or its result isn't finite"),
1674            )
1675        }),
1676        _ => float1(v, func, real),
1677    }
1678}
1679
1680fn float1(v: &Value, func: &str, f: impl Fn(f64) -> Option<f64>) -> OpResult<Value> {
1681    let x = number(v, func)?;
1682    f(x).filter(|y| x.is_finite() && y.is_finite())
1683        .map(Value::Float)
1684        .ok_or_else(|| {
1685            OpError::fault(
1686                Fault::DomainError,
1687                format!("`{func}` isn't defined for {}", fmt_float(x)),
1688            )
1689        })
1690}
1691
1692fn float2(a: &Value, b: &Value, func: &str, f: fn(f64, f64) -> f64) -> OpResult<Value> {
1693    let (a, b) = (number(a, func)?, number(b, func)?);
1694    if !a.is_finite() || !b.is_finite() {
1695        return Err(OpError::new(format!("`{func}` needs finite numbers")));
1696    }
1697    // Value equality merges signed zeros. In particular, atan2 must not
1698    // distinguish worlds (or memoized arguments) that the engine considers equal.
1699    let a = if a == 0.0 { 0.0 } else { a };
1700    let b = if b == 0.0 { 0.0 } else { b };
1701    let value = f(a, b);
1702    if !value.is_finite() {
1703        return Err(OpError::new(format!(
1704            "`{func}` gave a result that isn't a finite number"
1705        )));
1706    }
1707    Ok(Value::Float(value))
1708}
1709
1710/// `sqrt`, `exp`, `ln` and their kin of an outcome of a draw, through the
1711/// three transforms: `exp2(x)` is `exp(x ln 2)`, `log10(x)` is
1712/// `ln(x) / ln 10`, and so on.
1713fn transform(b: Builtin, v: &Value, budget: &mut Budget) -> OpResult<Value> {
1714    use crate::analytic::Kernel;
1715    use probl_syntax::ast::BinOp;
1716    use std::f64::consts::{LN_2, LN_10};
1717    let of = |k: Kernel, v: Value, concrete: Builtin, budget: &mut Budget| match &v {
1718        Value::Analytic(x) => crate::analytic::transformed(x, k, budget),
1719        _ => call_plain(concrete, &[v], budget),
1720    };
1721    let times = |v: &Value, c: f64, budget: &mut Budget| ops::binary(BinOp::Mul, v, &Value::Float(c), budget);
1722    match b {
1723        Builtin::Sqrt => of(Kernel::Sqrt, v.clone(), b, budget),
1724        Builtin::Exp => of(Kernel::Exp, v.clone(), b, budget),
1725        Builtin::Ln => of(Kernel::Ln, v.clone(), b, budget),
1726        Builtin::Exp2 => {
1727            let y = times(v, LN_2, budget)?;
1728            of(Kernel::Exp, y, Builtin::Exp, budget)
1729        }
1730        Builtin::Expm1 => {
1731            let y = of(Kernel::Exp, v.clone(), Builtin::Exp, budget)?;
1732            ops::binary(BinOp::Sub, &y, &Value::Float(1.0), budget)
1733        }
1734        Builtin::Log10 | Builtin::Log2 => {
1735            let y = of(Kernel::Ln, v.clone(), Builtin::Ln, budget)?;
1736            let base = if b == Builtin::Log10 { LN_10 } else { LN_2 };
1737            ops::binary(BinOp::Div, &y, &Value::Float(base), budget)
1738        }
1739        Builtin::Log1p => {
1740            let y = ops::binary(BinOp::Add, v, &Value::Float(1.0), budget)?;
1741            of(Kernel::Ln, y, Builtin::Ln, budget)
1742        }
1743        _ => unreachable!("a transform"),
1744    }
1745}
1746
1747fn to_int(v: &Value, f: fn(f64) -> f64) -> OpResult<Value> {
1748    match v {
1749        Value::Int(i) => Ok(Value::Int(i.clone())),
1750        other => {
1751            let x = number(other, "rounding")?;
1752            let r = f(x);
1753            Integer::from_f64(r).map(Value::Int).ok_or_else(|| {
1754                OpError::fault(
1755                    Fault::NumericOverflow,
1756                    format!("{} cannot be rounded to a finite int", fmt_float(x)),
1757                )
1758            })
1759        }
1760    }
1761}
1762
1763fn round(v: &Value, digits: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1764    let Some(digits) = digits else {
1765        return to_int(v, f64::round);
1766    };
1767    let digits = integer(digits, "round's digits", budget)?;
1768    if let Value::Int(n) = v {
1769        if *digits >= 0 {
1770            return Ok(v.clone());
1771        }
1772        let places = digits.abs();
1773        if places > Integer::from(probl_number::MAX_INTEGER_DIGITS) {
1774            return Ok(Value::Int(Integer::ZERO));
1775        }
1776        budget.integer_work(n, n, true)?;
1777        return Ok(Value::Int(n.round_decimal(places.to_u64().unwrap() as u32)?));
1778    }
1779    // Once outside this interval the exact magnitude of digits is irrelevant.
1780    let digits = if *digits > 323 {
1781        324
1782    } else if *digits < -308 {
1783        -309
1784    } else {
1785        digits.to_i64().unwrap()
1786    };
1787    float1(v, "round", |x| {
1788        // Beyond these bounds a decimal place cannot change a finite f64,
1789        // or every finite f64 rounds to zero. No unbounded powers or loops.
1790        if digits > 323 || x == 0.0 {
1791            return Some(x);
1792        }
1793        if digits < -308 {
1794            return Some(0.0);
1795        }
1796        let rounded = if digits >= 0 {
1797            // Splitting the scale supports subnormals (up to 323 places)
1798            // without forming an infinite power of ten.
1799            let high = libm::pow(10.0, digits.min(308) as f64);
1800            let low = libm::pow(10.0, (digits - 308).max(0) as f64);
1801            let scaled = (x * high) * low;
1802            // At this precision the decimal adjustment is smaller than half
1803            // an f64 step; scaling back could only introduce a new error.
1804            if scaled.abs() >= 1e16 || x.fract() == 0.0 {
1805                return Some(x);
1806            }
1807            (scaled.round() / low) / high
1808        } else {
1809            let scale = libm::pow(10.0, -digits as f64);
1810            let scaled = x / scale;
1811            if scaled.abs() >= 1e16 {
1812                return Some(x);
1813            }
1814            scaled.round() * scale
1815        };
1816        Some(rounded)
1817    })
1818}
1819
1820fn min_max(args: &[Value], want_max: bool, budget: &mut Budget) -> OpResult<Value> {
1821    let name = if want_max { "max" } else { "min" };
1822    let mut best: Option<Value> = None;
1823    for v in args {
1824        // Validate candidate outcomes, then lift each comparison. Reducing incrementally avoids a Cartesian product of
1825        // every candidate and preserves missing mass without drawing.
1826        let v = ops::lift1(v, budget, |v, budget| {
1827            budget.work(1)?;
1828            ops::compare(v, v)?;
1829            Ok(v.clone())
1830        })?;
1831        best = Some(match best {
1832            None => v,
1833            Some(b) => ops::lift2(&b, &v, budget, |b, v, _| {
1834                let ord = ops::compare(v, b)?;
1835                if (want_max && ord.is_gt()) || (!want_max && ord.is_lt()) {
1836                    Ok(v.clone())
1837                } else {
1838                    Ok(b.clone())
1839                }
1840            })?,
1841        });
1842    }
1843    best.ok_or_else(|| OpError::fault(Fault::EmptyCollection, format!("`{name}` of an empty list")))
1844}
1845
1846fn slice(v: &Value, start: &Value, end: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1847    let start = integer(start, "slice", budget)?;
1848    let end = end.map(|v| integer(v, "slice", budget)).transpose()?;
1849    let length: Integer = match v {
1850        Value::List(xs) => xs.len().into(),
1851        Value::Str(s) => s.chars().count().into(),
1852        Value::Range(lo, hi) => {
1853            budget.integer_work(lo, hi, false)?;
1854            range_len(lo, hi)?
1855        }
1856        other => return Err(expected("a list, range or string", other, "slice")),
1857    };
1858    let end = end.as_deref().unwrap_or(&length);
1859    let start = start.as_ref();
1860    if start.is_negative() || start > end || end > &length {
1861        return Err(OpError::fault(
1862            Fault::IndexOutOfBounds,
1863            "`slice` needs 0 <= start <= end <= length",
1864        ));
1865    }
1866    Ok(match v {
1867        Value::Range(lo, _) => {
1868            // Empty slices have a canonical empty range. Subtract one from the
1869            // offset before adding, so a slice ending at the largest int works.
1870            let (first, last) = if start == end {
1871                (Integer::ZERO, (-1).into())
1872            } else {
1873                (lo.add(start)?, lo.add(&end.sub(&Integer::ONE)?)?)
1874            };
1875            budget.integer_allocation(first.bits(), 1)?;
1876            budget.integer_allocation(last.bits(), 1)?;
1877            Value::Range(first, last)
1878        }
1879        Value::List(xs) => {
1880            let (start, end) = (start.to_u64().unwrap() as usize, end.to_u64().unwrap() as usize);
1881            budget.collection((end - start) as u128)?;
1882            budget.work((end - start) as u64)?;
1883            Value::list(xs[start..end].to_vec())
1884        }
1885        Value::Str(s) => {
1886            let (start, end) = (start.to_u64().unwrap() as usize, end.to_u64().unwrap() as usize);
1887            let boundary = |n| s.char_indices().nth(n).map_or(s.len(), |(i, _)| i);
1888            crate::text::value(&s[boundary(start)..boundary(end)], budget)?
1889        }
1890        _ => unreachable!(),
1891    })
1892}
1893
1894fn len(v: &Value) -> OpResult<Value> {
1895    let n: u128 = match v {
1896        Value::List(items) => items.len() as u128,
1897        Value::Str(s) => s.chars().count() as u128,
1898        Value::Map(m) => m.len() as u128,
1899        Value::Bag(b) => b.values().map(|n| *n as u128).sum(),
1900        Value::Range(lo, hi) => return Ok(Value::Int(range_len(lo, hi)?)),
1901        other => return Err(expected("a collection", other, "len")),
1902    };
1903    Ok(Value::Int(n.into()))
1904}
1905
1906fn sum(v: &Value, budget: &mut Budget) -> OpResult<Value> {
1907    let items = items(v, "sum", budget)?;
1908    let mut acc = Value::Int(0.into());
1909    for x in &items {
1910        acc = ops::binary(probl_syntax::ast::BinOp::Add, &acc, x, budget)?;
1911    }
1912    Ok(acc)
1913}
1914
1915fn get(coll: &Value, key: &Value, default: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1916    let found = match coll {
1917        Value::Map(m) => m.get(key).cloned(),
1918        Value::List(items) => integer(key, "index", budget)?
1919            .to_u64()
1920            .and_then(|n| usize::try_from(n).ok())
1921            .and_then(|n| items.get(n))
1922            .cloned(),
1923        Value::Str(s) => {
1924            let index = integer(key, "index", budget)?
1925                .to_u64()
1926                .and_then(|n| usize::try_from(n).ok());
1927            budget.string_work(s)?;
1928            index
1929                .and_then(|n| s.chars().nth(n))
1930                .map(|c| crate::text::value(c.encode_utf8(&mut [0; 4]), budget))
1931                .transpose()?
1932        }
1933        Value::Range(lo, hi) => {
1934            let index = integer(key, "index", budget)?;
1935            budget.integer_work(lo, hi, false)?;
1936            let len = range_len(lo, hi)?;
1937            if index.is_negative() || *index >= len {
1938                None
1939            } else {
1940                Some(Value::Int(lo.add(&index)?))
1941            }
1942        }
1943        Value::Bag(b) => Some(Value::Int(b.get(key).copied().unwrap_or(0).into())),
1944        other => return Err(expected("a map, bag or sequence", other, "get")),
1945    };
1946    match (found, default) {
1947        (Some(v), _) => Ok(v),
1948        (None, Some(d)) => Ok(d.clone()),
1949        (None, None) => Err(
1950            OpError::fault(Fault::MissingKey, format!("the key {key:?} isn't there"))
1951                .help("give a default: `get(key, default)`"),
1952        ),
1953    }
1954}
1955
1956/// Select an existing element; recipes inside collections are never lifted.
1957pub fn population_extreme(v: &Value, want_max: bool, default: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1958    let name = if want_max { "maximum" } else { "minimum" };
1959    match v {
1960        Value::Range(lo, hi) => {
1961            budget.integer_work(lo, hi, false)?;
1962            if hi < lo {
1963                return empty_extreme(name, default);
1964            }
1965            Ok(Value::Int(if want_max { hi.clone() } else { lo.clone() }))
1966        }
1967        Value::Continuous(f) => {
1968            let (lo, hi) = f.support();
1969            finite_bound(if want_max { hi } else { lo }, name)
1970        }
1971        Value::Dist(d) => {
1972            if d.missing > 0.0 {
1973                return Err(OpError::new(format!(
1974                    "`{name}` cannot determine a support bound while the distribution has unresolved mass"
1975                )));
1976            }
1977            let mut best = None;
1978            for (x, p) in &d.outcomes {
1979                if *p <= 0.0 {
1980                    continue;
1981                }
1982                let x = match x {
1983                    Value::Continuous(_) | Value::Dist(_) => population_extreme(x, want_max, None, budget)?,
1984                    x => x.clone(),
1985                };
1986                select_extreme(&mut best, x, want_max, budget)?;
1987            }
1988            best.ok_or_else(|| OpError::new(format!("`{name}` needs a nonempty distribution")))
1989        }
1990        Value::List(xs) if xs.iter().any(|x| matches!(x, Value::Analytic(_))) => {
1991            crate::analytic::extreme(xs, want_max, budget)
1992        }
1993        Value::List(_) | Value::Str(_) => {
1994            let mut best = None;
1995            for x in items(v, name, budget)? {
1996                select_extreme(&mut best, x, want_max, budget)?;
1997            }
1998            best.map_or_else(|| empty_extreme(name, default), Ok)
1999        }
2000        other => Err(expected(
2001            "a distribution or nonempty list, range or string",
2002            other,
2003            name,
2004        )),
2005    }
2006}
2007
2008pub fn empty_extreme(name: &str, default: Option<&Value>) -> OpResult<Value> {
2009    default.cloned().ok_or_else(|| {
2010        OpError::fault(
2011            Fault::EmptyCollection,
2012            format!("`{name}` needs a nonempty collection or a default"),
2013        )
2014        .help(format!("use `{name}(xs, default: value)` to handle empty collections"))
2015    })
2016}
2017
2018fn finite_bound(x: f64, name: &str) -> OpResult<Value> {
2019    if !x.is_finite() {
2020        return Err(OpError::new(format!("`{name}` has no finite support bound")));
2021    }
2022    Ok(Value::Float(x))
2023}
2024
2025fn select_extreme(best: &mut Option<Value>, x: Value, want_max: bool, budget: &mut Budget) -> OpResult<()> {
2026    budget.work(1)?;
2027    ops::compare(&x, &x)?;
2028    let replace = match best {
2029        None => true,
2030        Some(b) => {
2031            let ord = ops::compare(&x, b)?;
2032            if want_max { ord.is_gt() } else { ord.is_lt() }
2033        }
2034    };
2035    if replace {
2036        *best = Some(x);
2037    }
2038    Ok(())
2039}
2040
2041pub fn extreme_count(n: &Value, budget: &mut Budget) -> OpResult<usize> {
2042    let n = integer(n, "the count", budget)?;
2043    if n.is_negative() {
2044        return Err(OpError::fault(Fault::DomainError, "the count must be nonnegative"));
2045    }
2046    Ok(n.to_u64().and_then(|n| usize::try_from(n).ok()).unwrap_or(usize::MAX))
2047}
2048
2049fn extremes(v: &Value, n: &Value, highest: bool, budget: &mut Budget) -> OpResult<Value> {
2050    let n = extreme_count(n, budget)?;
2051    let mut items = items(v, if highest { "highest" } else { "lowest" }, budget)?;
2052    sort(&mut items, highest, budget)?;
2053    items.truncate(n);
2054    Ok(Value::list(items))
2055}
2056
2057fn insert(coll: &Value, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
2058    match coll {
2059        Value::List(items) => {
2060            budget.collection(items.len() as u128 + 1)?;
2061            budget.work(items.len() as u64)?;
2062            let mut items = items.to_vec();
2063            let i = as_index(a, items.len() as u128 + 1, budget)? as usize;
2064            items.insert(i, b.clone());
2065            Ok(Value::list(items))
2066        }
2067        Value::Map(m) => {
2068            budget.collection(m.len() as u128 + 1)?;
2069            budget.work(m.len() as u64)?;
2070            let mut m = BTreeMap::clone(m);
2071            m.insert(a.clone(), b.clone());
2072            Ok(Value::map(m))
2073        }
2074        other => Err(expected("a list or a map", other, "insert")),
2075    }
2076}
2077
2078fn remove(coll: &Value, key: &Value, budget: &mut Budget) -> OpResult<Value> {
2079    match coll {
2080        Value::List(items) => {
2081            budget.work(items.len() as u64)?;
2082            let mut items = items.to_vec();
2083            let i = as_index(key, items.len() as u128, budget)? as usize;
2084            items.remove(i);
2085            Ok(Value::list(items))
2086        }
2087        Value::Map(m) => {
2088            budget.work(m.len() as u64)?;
2089            let mut m = BTreeMap::clone(m);
2090            m.remove(key);
2091            Ok(Value::map(m))
2092        }
2093        Value::Bag(b) => match b.without(key) {
2094            Some(rest) => Ok(Value::multiset(rest)),
2095            None => Err(OpError::fault(Fault::MissingKey, format!("{key:?} isn't in the bag"))),
2096        },
2097        other => Err(expected("a list, map or bag", other, "remove")),
2098    }
2099}
2100
2101/// A choice among options. Options that are distributions are mixed in, so
2102/// that a value drawn from the result is settled (docs/semantics.md, section 2).
2103fn one_of(v: &Value, budget: &mut Budget) -> OpResult<Value> {
2104    match v {
2105        Value::List(items) if !items.is_empty() => {
2106            budget.outcomes(items.len() as u128)?;
2107            let p = 1.0 / items.len() as f64;
2108            ops::combine(items.iter().map(|x| (x.clone(), p)).collect(), 0.0, budget)
2109        }
2110        Value::Range(lo, hi) if hi >= lo => {
2111            budget.outcomes(range_count(lo, hi)?)?;
2112            Ok(Dist::uniform(range_items(lo, hi, budget)?).into_value())
2113        }
2114        Value::Map(m) if !m.is_empty() => {
2115            let all_probs = m.values().all(|w| matches!(w, Value::Prob(_)));
2116            if !all_probs && m.values().any(|w| matches!(w, Value::Prob(_))) {
2117                return Err(
2118                    OpError::new("one_of can't mix probabilities and relative numeric weights")
2119                        .help("use prob(...) for every absolute probability, or numbers for every relative weight"),
2120                );
2121            }
2122            let mut pairs = Vec::new();
2123            for (k, w) in m.iter() {
2124                let w = match w {
2125                    Value::Bool(_) => None,
2126                    _ => w.as_f64().filter(|w| *w >= 0.0 && w.is_finite()),
2127                }
2128                .ok_or_else(|| {
2129                    OpError::new(format!(
2130                        "one_of needs weights that are numbers of 0 or more, found {w:?}"
2131                    ))
2132                })?;
2133                pairs.push((k.clone(), w));
2134            }
2135            let largest = pairs.iter().map(|(_, w)| *w).fold(0.0, f64::max);
2136            if largest == 0.0 {
2137                return Err(OpError::new("one_of needs at least one positive weight"));
2138            }
2139            // Absolute probabilities are checked before normalization. Relative
2140            // weights may have an unrepresentable sum even though each is valid.
2141            let total = crate::stats::sum(pairs.iter().map(|(_, w)| *w));
2142            if all_probs && (total - 1.0).abs() > 1e-9 {
2143                return Err(OpError::new(format!(
2144                    "the chances add up to {}, not 100%",
2145                    crate::value::fmt_prob(total)
2146                ))
2147                .help("use plain numbers for relative weights, like [\"a\": 3, \"b\": 1]"));
2148            }
2149            for (_, w) in &mut pairs {
2150                *w /= largest;
2151            }
2152            let total = crate::stats::sum(pairs.iter().map(|(_, w)| *w));
2153            ops::combine(pairs.into_iter().map(|(k, w)| (k, w / total)).collect(), 0.0, budget)
2154        }
2155        Value::Bag(b) => {
2156            let total: u128 = b.values().map(|n| *n as u128).sum();
2157            if total == 0 {
2158                return Err(OpError::fault(Fault::EmptyCollection, "the bag is empty"));
2159            }
2160            let pairs = b.iter().map(|(k, n)| (k.clone(), *n as f64 / total as f64)).collect();
2161            ops::combine(pairs, 0.0, budget)
2162        }
2163        Value::List(_) | Value::Range(..) | Value::Map(_) => Err(OpError::fault(
2164            Fault::EmptyCollection,
2165            "one_of needs at least one option",
2166        )),
2167        other => Err(expected("a list, range, map or bag", other, "one_of")),
2168    }
2169}
2170
2171fn bag(v: &Value, budget: &mut Budget) -> OpResult<Value> {
2172    let mut counts = BTreeMap::new();
2173    match v {
2174        Value::Map(m) => {
2175            for (k, n) in m.iter() {
2176                let n = integer(n, "bag count", budget)?;
2177                if n.is_negative() {
2178                    return Err(OpError::new("bag counts must be whole numbers of 0 or more"));
2179                }
2180                let n = n
2181                    .to_u64()
2182                    .ok_or_else(|| OpError::new("bag count exceeds the supported count range"))?;
2183                if n > 0 {
2184                    counts.insert(k.clone(), n);
2185                }
2186            }
2187        }
2188        Value::List(items) => {
2189            for item in items.iter() {
2190                *counts.entry(item.clone()).or_insert(0) += 1;
2191            }
2192        }
2193        other => return Err(expected("a map of counts or a list", other, "bag")),
2194    }
2195    Ok(Value::bag(counts))
2196}