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;
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::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 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
496pub 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
562pub 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 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 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 _ => call_plain(b, args, budget),
707 }
708}
709
710fn 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 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 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
772fn 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
790fn 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 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
808fn 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
820pub 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
870fn 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
882fn continuous_query(b: Builtin, args: &[Value]) -> OpResult<Value> {
885 use Builtin as B;
886 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 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
974fn 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 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 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
1081fn 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
1152fn 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 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 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
1232pub(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 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 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 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
1494pub 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 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
1523fn 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 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
1572fn 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 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 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 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 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 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 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
1818pub 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
1959fn 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 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}