Skip to main content

probl_engine/
builtins.rs

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