1use 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
16pub 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 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 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 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
571pub 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
637pub 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 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 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 _ => call_plain(b, args, budget),
782 }
783}
784
785fn 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 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 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
847fn 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
865fn 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 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
883fn 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
895pub 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
945fn 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
957fn continuous_query(b: Builtin, args: &[Value]) -> OpResult<Value> {
960 use Builtin as B;
961 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 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
1049fn 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 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 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
1156fn 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
1227fn 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 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 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
1307pub(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 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 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 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 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
1573fn 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 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
1632pub 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 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
1661fn 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 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
1710fn 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 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 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 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 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 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 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
1956pub 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
2101fn 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 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}