1use crate::complex::Complex;
4use crate::continuous::{Family, Mixture, Part};
5use crate::dates;
6use crate::dist::{Budget, Counts, Dist};
7use crate::error::{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 => min_max(args, b == B::Max, budget),
44 B::Abs => match a(0) {
45 Value::Complex(z) => finite_float(z.abs(), "abs"),
46 Value::Int(n) => Ok(Value::Int(n.abs())),
47 Value::Prob(p) => Ok(Value::Prob(p.abs())),
48 v => float1(v, "abs", |x| Some(x.abs())),
49 },
50 B::Floor => to_int(a(0), f64::floor),
51 B::Ceil => to_int(a(0), f64::ceil),
52 B::Trunc => to_int(a(0), libm::trunc),
53 B::Round => round(a(0), args.get(1), budget),
54 B::Sqrt => elementary1(a(0), "sqrt", Complex::sqrt, |x| (x >= 0.0).then(|| x.sqrt())),
55 B::Cbrt => elementary1(a(0), "cbrt", Complex::cbrt, |x| Some(libm::cbrt(x))),
56 B::Exp => elementary1(a(0), "exp", Complex::exp, |x| {
57 Some(crate::math::exp(x)).filter(|y| y.is_finite())
58 }),
59 B::Exp2 => elementary1(a(0), "exp2", Complex::exp2, |x| Some(libm::exp2(x))),
60 B::Ln => elementary1(a(0), "ln", Complex::ln, |x| (x > 0.0).then(|| libm::log(x))),
61 B::Log10 => elementary1(a(0), "log10", Complex::log10, |x| (x > 0.0).then(|| libm::log10(x))),
62 B::Log2 => elementary1(a(0), "log2", Complex::log2, |x| (x > 0.0).then(|| libm::log2(x))),
63 B::Log1p => elementary1(a(0), "log1p", Complex::log1p, |x| (x > -1.0).then(|| libm::log1p(x))),
64 B::Expm1 => elementary1(a(0), "expm1", Complex::expm1, |x| {
65 Some(libm::expm1(x)).filter(|y| y.is_finite())
66 }),
67 B::Sin => elementary1(a(0), "sin", Complex::sin, |x| Some(libm::sin(x))),
68 B::Cos => elementary1(a(0), "cos", Complex::cos, |x| Some(libm::cos(x))),
69 B::Tan => elementary1(a(0), "tan", Complex::tan, |x| Some(libm::tan(x))),
70 B::Asin => elementary1(a(0), "asin", Complex::asin, |x| {
71 (-1.0..=1.0).contains(&x).then(|| libm::asin(x))
72 }),
73 B::Acos => elementary1(a(0), "acos", Complex::acos, |x| {
74 (-1.0..=1.0).contains(&x).then(|| libm::acos(x))
75 }),
76 B::Atan => elementary1(a(0), "atan", Complex::atan, |x| Some(libm::atan(x))),
77 B::Atan2 => float2(a(0), a(1), "atan2", libm::atan2),
78 B::Hypot => float2(a(0), a(1), "hypot", libm::hypot),
79 B::Sinh => elementary1(a(0), "sinh", Complex::sinh, |x| {
80 Some(libm::sinh(x)).filter(|y| y.is_finite())
81 }),
82 B::Cosh => elementary1(a(0), "cosh", Complex::cosh, |x| {
83 Some(libm::cosh(x)).filter(|y| y.is_finite())
84 }),
85 B::Tanh => elementary1(a(0), "tanh", Complex::tanh, |x| Some(libm::tanh(x))),
86 B::Asinh => elementary1(a(0), "asinh", Complex::asinh, |x| Some(libm::asinh(x))),
87 B::Acosh => elementary1(a(0), "acosh", Complex::acosh, |x| (x >= 1.0).then(|| libm::acosh(x))),
88 B::Atanh => elementary1(a(0), "atanh", Complex::atanh, |x| {
89 (x.abs() < 1.0).then(|| libm::atanh(x))
90 }),
91 B::BitLength => Ok(Value::Int(integer(a(0), "bit_length", budget)?.bits().into())),
92 B::BitAnd => Ok(Value::Int(
93 integer(a(0), b.name(), budget)?.bit_and(integer(a(1), b.name(), budget)?.as_ref())?,
94 )),
95 B::BitOr => Ok(Value::Int(
96 integer(a(0), b.name(), budget)?.bit_or(integer(a(1), b.name(), budget)?.as_ref())?,
97 )),
98 B::BitXor => Ok(Value::Int(
99 integer(a(0), b.name(), budget)?.bit_xor(integer(a(1), b.name(), budget)?.as_ref())?,
100 )),
101 B::BitNot => Ok(Value::Int(integer(a(0), b.name(), budget)?.bit_not()?)),
102 B::BitCount => Ok(Value::Int(integer(a(0), b.name(), budget)?.bit_count().into())),
103 B::ILog2 => {
104 let n = integer(a(0), "ilog2", budget)?;
105 if n.is_zero() || n.is_negative() {
106 return Err(OpError::new("`ilog2` needs a positive integer"));
107 }
108 Ok(Value::Int((n.bits() - 1).into()))
109 }
110 B::Choose => choose(
111 nonnegative_int(a(0), "choose", budget)?.as_ref(),
112 nonnegative_int(a(1), "choose", budget)?.as_ref(),
113 budget,
114 ),
115 B::Factorial => factorial(nonnegative_int(a(0), "factorial", budget)?.as_ref(), budget),
116 B::Gcd | B::Lcm => {
117 let (x, y) = (
118 integer(a(0), b.name(), budget)?.abs(),
119 integer(a(1), b.name(), budget)?.abs(),
120 );
121 let d = gcd(x.clone(), y.clone(), budget)?;
122 let result = if b == B::Gcd {
123 d
124 } else if x.is_zero() || y.is_zero() {
125 Integer::ZERO
126 } else {
127 let q = x.div_mod(&d)?.0;
128 budget.integer_work(&q, &y, true)?;
129 q.mul(&y)?
130 };
131 Ok(Value::Int(result))
132 }
133 B::EulerPhi => euler_phi(nonnegative_int(a(0), "euler_phi", budget)?.as_ref(), budget),
134 B::LnGamma => float1(a(0), "ln_gamma", |x| (x > 0.0).then(|| libm::lgamma(x))),
135 B::Erf => float1(a(0), "erf", |x| Some(libm::erf(x))),
136 B::Erfc => float1(a(0), "erfc", |x| Some(libm::erfc(x))),
137 B::Complex => {
138 let z = if args.len() == 1 {
139 complex_number(a(0), "complex")?
140 } else {
141 Complex::new(number(a(0), "complex")?, number(a(1), "complex")?)?
142 };
143 Ok(Value::Complex(z))
144 }
145 B::Real => Ok(Value::Float(complex_number(a(0), "real")?.re())),
146 B::Imag => Ok(Value::Float(complex_number(a(0), "imag")?.im())),
147 B::Conj => Ok(Value::Complex(complex_number(a(0), "conj")?.conjugate())),
148 B::Abs2 => finite_float(complex_number(a(0), "abs2")?.abs2(), "abs2"),
149 B::Arg => Ok(Value::Float(complex_number(a(0), "arg")?.arg())),
150 B::Cis => {
151 let theta = number(a(0), "cis")?;
152 Complex::new(libm::cos(theta), libm::sin(theta)).map(Value::Complex)
153 }
154 B::Clamp => {
155 let (lo, hi) = (a(1), a(2));
156 if ops::compare(lo, hi)?.is_gt() {
157 return Err(OpError::new("clamp's lower bound is above its upper bound"));
158 }
159 if ops::compare(a(0), lo)?.is_lt() {
160 Ok(lo.clone())
161 } else if ops::compare(a(0), hi)?.is_gt() {
162 Ok(hi.clone())
163 } else {
164 Ok(a(0).clone())
165 }
166 }
167 B::Str => crate::text::formatted(a(0), budget),
168 B::Upper | B::Lower => crate::text::case(&text(a(0), b.name())?, b == B::Upper, budget),
169 B::Trim | B::TrimStart | B::TrimEnd => {
170 let s = text(a(0), b.name())?;
171 let set = args.get(1).map(|v| text(v, b.name())).transpose()?;
172 let set = set
173 .as_ref()
174 .map(|s| {
175 let mut set = rustc_hash::FxHashSet::default();
176 for c in s.chars() {
177 if !set.contains(&c) {
178 budget.collection(set.len() as u128 + 1)?;
179 set.insert(c);
180 }
181 }
182 Ok::<_, OpError>(set)
183 })
184 .transpose()?;
185 let matches = |c: char| set.as_ref().map_or_else(|| c.is_whitespace(), |set| set.contains(&c));
186 let trimmed = match b {
187 B::TrimStart => s.trim_start_matches(matches),
188 B::TrimEnd => s.trim_end_matches(matches),
189 _ => s.trim_matches(matches),
190 };
191 crate::text::value(trimmed, budget)
192 }
193 B::StartsWith | B::EndsWith => {
194 let (s, part) = (text(a(0), b.name())?, text(a(1), b.name())?);
195 Ok(Value::Bool(if b == B::StartsWith {
196 s.starts_with(&*part)
197 } else {
198 s.ends_with(&*part)
199 }))
200 }
201 B::Chars => crate::text::chars(&text(a(0), "chars")?, budget).map(Value::list),
202 B::Split => {
203 let (s, sep) = (text(a(0), "split")?, text(a(1), "split")?);
204 if sep.is_empty() {
205 return crate::text::chars(&s, budget).map(Value::list);
206 }
207 let mut parts = Vec::new();
208 for part in s.split(&*sep) {
209 budget.collection(parts.len() as u128 + 1)?;
210 parts.push(crate::text::value(part, budget)?);
211 }
212 Ok(Value::list(parts))
213 }
214 B::Join => {
215 let items = list(a(0), "join")?;
216 let sep = text(a(1), "join")?;
217 budget.collection(items.len() as u128)?;
218 budget.work(items.len() as u64)?;
219 let mut out = String::new();
220 for (i, item) in items.iter().enumerate() {
221 if i != 0 {
222 crate::text::push(&mut out, &sep, budget)?;
223 }
224 crate::text::push_value(&mut out, item, budget)?;
225 }
226 Ok(Value::str(&out))
227 }
228 B::Len => len(a(0)),
229 B::Slice => slice(a(0), a(1), args.get(2), budget),
230 B::Sum => sum(a(0), budget),
231 B::Count if args.len() == 1 => len(a(0)),
232 B::Sort | B::SortDesc => {
233 assert_eq!(args.len(), 1, "comparator sorting is handled by the interpreter");
234 let mut items = items(a(0), b.name(), budget)?;
235 sort(&mut items, b == B::SortDesc, budget)?;
236 Ok(Value::list(items))
237 }
238 B::Reverse => match a(0) {
239 Value::Str(s) => {
240 budget.string_allocation(s.len())?;
241 Ok(Value::str(&s.chars().rev().collect::<String>()))
242 }
243 v => {
244 let mut items = items(v, "reverse", budget)?;
245 items.reverse();
246 Ok(Value::list(items))
247 }
248 },
249 B::Keys => match a(0) {
250 Value::Map(m) => Ok(Value::list(m.keys().cloned().collect())),
251 Value::Bag(m) => Ok(Value::list(m.keys().cloned().collect())),
252 v => Err(expected("a map", v, "keys")),
253 },
254 B::Values => match a(0) {
255 Value::Map(m) => Ok(Value::list(m.values().cloned().collect())),
256 v => Err(expected("a map", v, "values")),
257 },
258 B::Get => get(a(0), a(1), args.get(2), budget),
259 B::Contains => ops::contains(a(0), a(1), budget).map(Value::Bool),
260 B::Highest | B::Lowest => extremes(a(0), a(1), b == B::Highest, budget),
261 B::Enumerate => {
262 let items = items(a(0), "enumerate", budget)?;
263 Ok(Value::list(
264 items
265 .into_iter()
266 .enumerate()
267 .map(|(i, x)| Value::list(vec![Value::Int((i as i64).into()), x]))
268 .collect(),
269 ))
270 }
271 B::Zip => {
272 let (xs, ys) = (items(a(0), "zip", budget)?, items(a(1), "zip", budget)?);
273 Ok(Value::list(
274 xs.into_iter().zip(ys).map(|(x, y)| Value::list(vec![x, y])).collect(),
275 ))
276 }
277 B::Push => {
278 let items = list(a(0), "push")?;
279 budget.collection(items.len() as u128 + 1)?;
280 let mut items = items.to_vec();
281 items.push(a(1).clone());
282 Ok(Value::list(items))
283 }
284 B::Insert => insert(a(0), a(1), a(2), budget),
285 B::Remove => remove(a(0), a(1), budget),
286 B::Pop => Err(
287 OpError::new("`pop` changes a list, so call it as a method on a variable")
288 .help("write `let top = xs.pop()`"),
289 ),
290 B::Last => {
291 let items = list(a(0), "pop")?;
292 items
293 .last()
294 .cloned()
295 .ok_or_else(|| OpError::new("can't pop from an empty list"))
296 }
297 B::DropLast => {
298 let items = list(a(0), "pop")?;
299 let n = items.len().saturating_sub(1);
300 Ok(Value::list(items[..n].to_vec()))
301 }
302 B::Prob => ops::make_prob(a(0)),
303 B::BooleanLaw => ops::boolean_law(a(0)),
304 B::Bernoulli | B::ScoreLaw => Ok(Dist::bernoulli(to_prob(a(0))?).into_value()),
305 B::OneOf => one_of(a(0), budget),
306 B::Binomial | B::Poisson | B::Geometric => {
307 let counts = counts(b, args, budget)?.expect("called with plain arguments");
308 Ok(counts.list(budget)?.into_value())
309 }
310 B::Bag => bag(a(0), budget),
311 B::Normal => continuous(Family::normal(number(a(0), "normal")?, number(a(1), "normal")?)),
312 B::Lognormal => continuous(Family::lognormal(
313 number(a(0), "lognormal")?,
314 number(a(1), "lognormal")?,
315 )),
316 B::Uniform => continuous(Family::uniform(number(a(0), "uniform")?, number(a(1), "uniform")?)),
317 B::Beta => continuous(Family::beta(number(a(0), "beta")?, number(a(1), "beta")?)),
318 B::Gamma => continuous(Family::gamma(number(a(0), "gamma")?, number(a(1), "gamma")?)),
319 B::Exponential => continuous(Family::exponential(number(a(0), "exponential")?)),
320 B::Triangular => continuous(Family::triangular(
321 number(a(0), "triangular")?,
322 number(a(1), "triangular")?,
323 number(a(2), "triangular")?,
324 )),
325 B::Pert => continuous(Family::pert(
326 number(a(0), "pert")?,
327 number(a(1), "pert")?,
328 number(a(2), "pert")?,
329 )),
330 B::NormalRange => continuous(Family::normal_range(
331 number(a(0), "normal_range")?,
332 number(a(1), "normal_range")?,
333 )),
334 B::To => continuous(Family::estimate(number(a(0), "to")?, number(a(1), "to")?)),
335 B::Mixture | B::Truncate | B::Bins => {
336 Err(OpError::unsupported(format!("`{}` isn't implemented yet", b.name())))
337 }
338 B::Odds => {
339 let p = to_prob(a(0))?;
340 if p >= 1.0 {
341 return Err(OpError::new("the odds of a certain event are infinite"));
342 }
343 Ok(Value::Float(p / (1.0 - p)))
344 }
345 B::Logit => {
346 let p = to_prob(a(0))?;
347 if p <= 0.0 || p >= 1.0 {
348 return Err(OpError::new("logit needs a probability strictly between 0% and 100%"));
349 }
350 Ok(Value::Float(libm::log(p / (1.0 - p))))
351 }
352 B::InvLogit => ops::computed_prob(1.0 / (1.0 + crate::math::exp(-number(a(0), "inv_logit")?)), "inv_logit"),
353 B::Date => {
354 let date = match args.len() {
355 1 => dates::parse(&text(a(0), "date")?),
356 3 => dates::from_parts(
357 date_count(a(0), "date", budget)?,
358 date_count(a(1), "date", budget)?,
359 date_count(a(2), "date", budget)?,
360 ),
361 _ => {
362 return Err(OpError::new(
363 "`date` takes one ISO string or three integers (year, month, day)",
364 ));
365 }
366 };
367 date.map(Value::Date)
368 .ok_or_else(|| OpError::new("invalid date; use YYYY-MM-DD within 0001-01-01..9999-12-31"))
369 }
370 B::Days => to_int(a(0), f64::round),
371 B::Weeks if matches!(a(0), Value::Int(_)) => {
372 ops::binary(probl_syntax::ast::BinOp::Mul, a(0), &Value::Int(7.into()), budget)
373 }
374 B::Weeks => {
375 let n = number(a(0), "weeks")?;
376 to_int(&Value::Float(n * 7.0), f64::round)
377 }
378 B::AddWorkdays => {
379 let d = date_value(a(0), b.name())?;
380 let n = date_count(a(1), b.name(), budget)?;
381 let holidays = holiday_calendar(args.get(2), b.name(), budget)?;
382 date_result(dates::add_workdays_with_holidays(d, n, &holidays))
383 }
384 B::IsWorkday => {
385 let d = date_value(a(0), b.name())?;
386 let holidays = holiday_calendar(args.get(1), b.name(), budget)?;
387 Ok(Value::Bool(
388 dates::weekday(d) < 5 && holidays.binary_search(&d).is_err(),
389 ))
390 }
391 B::AddMonths | B::AddYears => {
392 let d = date_value(a(0), b.name())?;
393 let n = date_count(a(1), b.name(), budget)?;
394 date_result(if b == B::AddMonths {
395 dates::add_months(d, n)
396 } else {
397 dates::add_years(d, n)
398 })
399 }
400 B::StartOfMonth | B::EndOfMonth => {
401 let d = date_value(a(0), b.name())?;
402 date_result(if b == B::StartOfMonth {
403 dates::start_of_month(d)
404 } else {
405 dates::end_of_month(d)
406 })
407 }
408 B::Year | B::Month | B::Day => {
409 let (y, m, d) = dates::to_civil(date_value(a(0), b.name())? as i64);
410 Ok(Value::Int(match b {
411 B::Year => y.into(),
412 B::Month => m.into(),
413 _ => d.into(),
414 }))
415 }
416 B::Weekday => crate::text::value(
417 dates::WEEKDAYS[dates::weekday(date_value(a(0), b.name())?) as usize],
418 budget,
419 ),
420 B::IsFalse => Ok(Value::Bool(ops::is_certain(a(0), false))),
421 B::IsTrue => Ok(Value::Bool(ops::is_certain(a(0), true))),
422 B::IsListOfLen => Ok(Value::Bool(
423 matches!((a(0), a(1)), (Value::List(items), Value::Int(n)) if *n == items.len() as i64),
424 )),
425 B::Minimum | B::Maximum => unreachable!("population extrema receive distributions whole"),
426 B::Typeof | B::RunDate | B::Count | B::Map | B::Filter | B::Reduce | B::Print | B::Roll | B::Take => {
427 unreachable!("`{}` is handled by the interpreter", b.name())
428 }
429 B::P
430 | B::Mean
431 | B::Sd
432 | B::Variance
433 | B::Median
434 | B::MedianLow
435 | B::MedianHigh
436 | B::Quantile
437 | B::Support
438 | B::Cdf
439 | B::Pmf
440 | B::IterItems
441 | B::RepeatCount
442 | B::Pdf
443 | B::Settled => unreachable!("`{}` takes distributions as they are", b.name()),
444 }
445}
446
447pub fn check_query_input(b: Builtin, args: &[Value]) -> OpResult<()> {
451 use Builtin as B;
452 let expected = match b {
453 B::Minimum | B::Maximum => "a distribution or nonempty list, range or string",
454 B::P => "a boolean distribution (dist[bool])",
455 B::Pdf => "a continuous distribution",
456 B::Mean
457 | B::Variance
458 | B::Sd
459 | B::Median
460 | B::MedianLow
461 | B::MedianHigh
462 | B::Quantile
463 | B::Cdf
464 | B::Pmf
465 | B::Support => "a distribution, nonempty list or nonempty range",
466 _ => return Ok(()),
467 };
468 let v = &args[0];
469 let valid = match b {
470 B::Minimum | B::Maximum => matches!(
471 v,
472 Value::Dist(_) | Value::Continuous(_) | Value::List(_) | Value::Range(..) | Value::Str(_)
473 ),
474 B::P => matches!(v, Value::Dist(_)),
475 B::Pdf => matches!(v, Value::Dist(_) | Value::Continuous(_)),
476 _ => matches!(
477 v,
478 Value::Dist(_) | Value::Continuous(_) | Value::List(_) | Value::Range(..)
479 ),
480 };
481 if !valid {
482 let help = if b == B::P {
483 "use `report event` to measure a fact across worlds, or `prob(event)` to convert a bool to 0 or 1"
484 } else {
485 "use `report x` to summarize values across worlds, or put the model inside `simulate { ... }` to obtain a distribution"
486 };
487 return Err(OpError::new(format!(
488 "`{}` expects {expected}, found {}",
489 b.name(),
490 article(&v.kind())
491 ))
492 .help(help));
493 }
494 if !matches!(b, B::Minimum | B::Maximum) && matches!(v, Value::List(xs) if xs.is_empty()) {
495 return Err(OpError::new(format!("`{}` needs a nonempty list", b.name())));
496 }
497 Ok(())
498}
499
500pub fn call_raw(b: Builtin, args: &[Value], budget: &mut Budget) -> OpResult<Value> {
502 use Builtin as B;
503 check_query_input(b, args)?;
504 let v = &args[0];
505 if let Value::Range(lo, hi) = v {
506 if matches!(
507 b,
508 B::Mean
509 | B::Variance
510 | B::Sd
511 | B::Median
512 | B::MedianLow
513 | B::MedianHigh
514 | B::Quantile
515 | B::Cdf
516 | B::Pmf
517 | B::Support
518 ) {
519 return range_query(b, lo, hi, args.get(1), budget);
520 }
521 }
522 if matches!(b, B::P | B::Cdf | B::Pmf | B::Pdf) {
523 if let Value::Dist(d) = v {
524 if d.missing > 0.0 {
525 return Err(OpError::new(format!("`{}` cannot return an exact scalar probability while the distribution has unresolved mass", b.name()))
526 .help("report a distribution comparison to retain probability bounds; scalar queries require a fully resolved distribution"));
527 }
528 }
529 }
530 if matches!(
531 b,
532 B::Mean
533 | B::Variance
534 | B::Sd
535 | B::Median
536 | B::MedianLow
537 | B::MedianHigh
538 | B::Quantile
539 | B::Cdf
540 | B::Pdf
541 | B::Pmf
542 | B::Support
543 ) && continuous_parts(v)
544 {
545 if matches!(b, B::Median | B::MedianLow | B::MedianHigh) {
546 let n = match v {
547 Value::Dist(d) => d.outcomes.len(),
548 _ => 1,
549 };
550 budget.work((n as u64).saturating_mul(n.max(1).ilog2() as u64 + 1))?;
552 }
553 return continuous_query(b, args);
554 }
555 match b {
556 B::Minimum | B::Maximum => population_extreme(v, b == B::Maximum, None, budget),
557 B::P => probability_of(v),
558 B::Pdf => Err(OpError::new("pdf needs a continuous distribution")
559 .help("for a distribution whose outcomes can be listed, use `pmf`")),
560 B::Mean => {
561 let d = stat_dist(v, b.name(), budget)?;
562 if d.outcomes.iter().any(|(x, _)| matches!(x, Value::Date(_))) {
563 return date_mean(v, &d);
564 }
565 if d.outcomes.iter().any(|(x, _)| matches!(x, Value::Complex(_))) {
566 budget.work(d.outcomes.len() as u64)?;
567 let mut sum = Complex::new(0.0, 0.0)?;
568 let total = d.total();
569 for (x, p) in &d.outcomes {
570 let x = complex_number(x, "mean")?;
571 sum = sum.plus(x.times(Complex::new(p / total, 0.0)?)?)?;
572 }
573 Ok(Value::Complex(sum))
574 } else {
575 check_numeric(&d, "mean")?;
576 finite_float(d.mean().unwrap(), "mean")
577 }
578 }
579 B::Variance => {
580 numeric_dist(v, "variance", budget).and_then(|d| finite_float(d.variance().unwrap(), "variance"))
581 }
582 B::Sd => numeric_dist(v, "sd", budget).and_then(|d| finite_float(d.sd().unwrap(), "sd")),
583 B::Median | B::MedianLow | B::MedianHigh => median(v, b, budget),
584 B::Quantile => {
585 let q = to_prob(&args[1])?;
586 quantile(v, q, "quantile", budget)
587 }
588 B::Support => {
589 let d = stat_dist(v, b.name(), budget)?;
590 budget.collection(d.outcomes.len() as u128)?;
591 Ok(Value::list(d.outcomes.iter().map(|(x, _)| x.clone()).collect()))
592 }
593 B::Cdf => {
594 let d = stat_dist(v, b.name(), budget)?;
595 let mut p = crate::stats::Sum::default();
596 for (x, w) in &d.outcomes {
597 if statistical_compare(x, &args[1])?.is_le() {
598 p.add(*w);
599 }
600 }
601 ops::computed_prob(p.value() / d.total(), "cdf")
602 }
603 B::Pmf => {
604 let d = stat_dist(v, b.name(), budget)?;
605 let p = crate::stats::sum(d.outcomes.iter().filter(|(x, _)| *x == args[1]).map(|(_, w)| *w));
606 ops::computed_prob(p / d.total(), "pmf")
607 }
608 B::IterItems => iter_items(v, budget),
609 B::RepeatCount => {
610 if v.is_uncertain() {
611 return Err(OpError::new(format!("`repeat` needs a number, not a {}", v.kind()))
612 .help("draw a value first, like `let n ~ d6`, then `repeat n { … }`"));
613 }
614 let n = integer(v, "repeat count", budget)?;
615 if n.is_negative() {
616 return Err(OpError::new("`repeat` needs a count of 0 or more"));
617 }
618 Ok(Value::Int(n.into_owned()))
619 }
620 B::Settled => match v {
621 v if v.is_uncertain() => Err(
622 OpError::new(format!("`match` needs a settled value, not a {}", v.kind()))
623 .help("draw a value first, like `let x ~ d6`, and match on `x`"),
624 ),
625 other => Ok(other.clone()),
626 },
627 _ => call_plain(b, args, budget),
629 }
630}
631
632fn probability_of(v: &Value) -> OpResult<Value> {
634 match v {
635 Value::Dist(d) if d.truth().is_some() => ops::computed_prob(d.truth().unwrap().0 / d.total(), "P"),
636 _ => Err(OpError::new(format!(
637 "P needs a boolean distribution (dist[bool]), found {}",
638 article(&v.kind())
639 ))
640 .help("compare a distribution, like `P(d6 > 4)`; use `report event` for a fact across worlds")),
641 }
642}
643
644pub fn counts(b: Builtin, args: &[Value], budget: &mut Budget) -> OpResult<Option<Counts>> {
647 use Builtin as B;
648 if args.iter().any(Value::is_uncertain) {
649 return Ok(None);
650 }
651 Ok(Some(match b {
652 B::Binomial => {
653 let n = whole(&args[0], "binomial's number of trials", budget)?;
654 if n < 0 {
655 return Err(OpError::new("binomial needs a number of trials of 0 or more"));
656 }
657 Counts::Binomial {
658 n: n as u64,
659 p: to_prob(&args[1])?,
660 }
661 }
662 B::Poisson => {
663 let rate = number(&args[0], "poisson")?;
664 if rate < 0.0 || !rate.is_finite() {
665 return Err(OpError::new("poisson needs a rate of 0 or more"));
666 }
667 if rate > 1e15 {
668 return Err(OpError::new("poisson's rate is too large to count exactly")
669 .help("above 10¹⁵, use a normal distribution with the same mean and variance"));
670 }
671 Counts::Poisson { rate }
672 }
673 B::Geometric => {
674 let p = to_prob(&args[0])?;
675 if p <= 0.0 {
676 return Err(OpError::new("geometric needs a chance of success above 0%"));
677 }
678 Counts::Geometric { p }
679 }
680 _ => return Ok(None),
681 }))
682}
683
684fn continuous(family: OpResult<Family>) -> OpResult<Value> {
685 family.map(|f| Value::Continuous(Arc::new(f)))
686}
687
688fn continuous_parts(v: &Value) -> bool {
690 match v {
691 Value::Continuous(_) => true,
692 Value::Dist(d) => d.outcomes.iter().any(|(x, _)| matches!(x, Value::Continuous(_))),
693 _ => false,
694 }
695}
696
697fn continuous_query(b: Builtin, args: &[Value]) -> OpResult<Value> {
700 use Builtin as B;
701 if b == B::Pmf {
704 return match &args[0] {
705 Value::Dist(d) => {
706 let at = crate::stats::sum(
707 d.outcomes
708 .iter()
709 .filter(|(x, _)| !matches!(x, Value::Continuous(_)) && *x == args[1])
710 .map(|(_, p)| *p),
711 );
712 ops::computed_prob(at / d.total(), "pmf")
713 }
714 Value::Continuous(_) => ops::computed_prob(0.0, "pmf"),
715 _ => unreachable!("checked continuous input"),
716 };
717 }
718 let parts = match &args[0] {
719 Value::Continuous(f) => vec![(Part::Continuous(**f), 1.0)],
720 Value::Dist(d) => d
721 .outcomes
722 .iter()
723 .map(|(x, p)| {
724 let part = match x {
725 Value::Continuous(f) => Part::Continuous(**f),
726 other => Part::Point(number(other, b.name())?),
727 };
728 Ok((part, *p))
729 })
730 .collect::<OpResult<Vec<_>>>()?,
731 _ => unreachable!("checked by `continuous_parts`"),
732 };
733 let m = Mixture { parts };
734 match b {
735 B::Mean => finite_float(m.mean(), "mean"),
736 B::Variance => finite_float(m.variance(), "variance"),
737 B::Sd => finite_float(m.sd(), "sd"),
738 B::Median | B::MedianLow | B::MedianHigh => {
739 let (lo, hi) = m.median_bounds();
740 finite_float(
741 match b {
742 B::MedianLow => lo,
743 B::MedianHigh => hi,
744 _ => crate::stats::midpoint(lo, hi),
745 },
746 b.name(),
747 )
748 }
749 B::Quantile => finite_float(m.quantile(to_prob(&args[1])?), "quantile"),
750 B::Cdf => ops::computed_prob(m.cdf(number(&args[1], "cdf")?), "cdf"),
751 B::Pdf => {
752 let x = number(&args[1], "pdf")?;
753 let mut density = 0.0;
754 for (part, p) in &m.parts {
755 match part {
756 Part::Continuous(f) => density += p * f.pdf(x),
757 Part::Analytic(a) => density += p * a.pdf(x),
758 Part::Point(_) => {
759 return Err(OpError::new(
760 "pdf needs a continuous distribution, without single values mixed in",
761 ));
762 }
763 }
764 }
765 finite_float(density, "pdf")
766 }
767 _ => Err(OpError::new(format!(
768 "`{}` needs a distribution whose outcomes can be listed, not a continuous one",
769 b.name()
770 ))),
771 }
772}
773
774fn range_query(b: Builtin, lo: &Integer, hi: &Integer, arg: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
777 use Builtin as B;
778 if hi < lo {
779 return Err(OpError::new(format!("`{}` needs a nonempty range", b.name())));
780 }
781 budget.integer_work(lo, hi, false)?;
782 let n = range_len(lo, hi)?;
783 budget.integer_allocation(n.bits(), 1)?;
784 let two = Integer::from(2);
785 match b {
786 B::Support => Ok(Value::list(range_items(lo, hi, budget)?)),
787 B::Mean => lo
788 .add(hi)?
789 .ratio(&two)
790 .map(Value::Float)
791 .ok_or_else(|| OpError::new("`mean` gave a result that isn't a finite number")),
792 B::Median | B::MedianLow | B::MedianHigh => {
793 budget.integer_work(&n, &two, true)?;
794 let (half, odd) = n.div_mod(&two)?;
795 let high = Value::Int(lo.add(&half)?);
796 let low = if odd.is_zero() {
797 Value::Int(lo.add(&half)?.sub(&Integer::ONE)?)
798 } else {
799 high.clone()
800 };
801 match b {
802 B::MedianLow => Ok(low),
803 B::MedianHigh => Ok(high),
804 _ => midpoint(&low, &high, budget),
805 }
806 }
807 B::Sd | B::Variance => {
808 let quarter = n.ratio(&Integer::from(4)).unwrap_or(f64::INFINITY);
809 let inv = Integer::ONE.ratio(&n).unwrap_or(0.0);
810 let sd = quarter * (4.0 / libm::sqrt(12.0)) * libm::sqrt((1.0 - inv) * (1.0 + inv));
811 finite_float(if b == B::Sd { sd } else { sd * sd }, b.name())
812 }
813 B::Quantile => {
814 let q = to_prob(arg.expect("quantile argument"))?;
815 if q == 0.0 {
816 return Ok(Value::Int(lo.clone()));
817 }
818 budget.work(n.bits().div_ceil(64).saturating_mul(17))?;
819 let rank = if let Some(count) = n.to_u64().filter(|n| *n <= 1 << 53) {
820 let mut weight = 1.0 / count as f64;
823 let total = weight * count as f64;
824 if total != 1.0 {
825 weight *= 1.0 / total;
826 }
827 let target = q * (weight * count as f64);
828 let (mut left, mut right) = (1, count);
829 while left < right {
830 budget.work(1)?;
831 let mid = left + (right - left) / 2;
832 if mid as f64 * weight >= target {
833 right = mid;
834 } else {
835 left = mid + 1;
836 }
837 }
838 Integer::from(left)
839 } else {
840 n.probability_rank(q)?
843 };
844 Ok(Value::Int(lo.add(&rank.sub(&Integer::ONE)?)?))
845 }
846 B::Pmf => {
847 let present = matches!(arg, Some(Value::Int(k)) if lo <= k && k <= hi);
848 ops::computed_prob(
849 if present {
850 Integer::ONE.ratio(&n).unwrap_or(0.0)
851 } else {
852 0.0
853 },
854 "pmf",
855 )
856 }
857 B::Cdf => {
858 let x = arg.expect("cdf argument");
859 if ops::compare(x, &Value::Int(lo.clone()))?.is_lt() {
860 return ops::computed_prob(0.0, "cdf");
861 }
862 if ops::compare(x, &Value::Int(hi.clone()))?.is_ge() {
863 return ops::computed_prob(1.0, "cdf");
864 }
865 let k = match x {
866 Value::Int(k) => k.clone(),
867 _ => Integer::from_f64(number(x, "cdf")?.floor()).expect("finite numeric bound"),
868 };
869 let count = k.sub(lo)?.add(&Integer::ONE)?;
870 budget.integer_work(&count, &n, true)?;
871 ops::computed_prob(count.ratio(&n).unwrap_or(0.0), "cdf")
872 }
873 _ => unreachable!("checked range query"),
874 }
875}
876
877fn stat_dist<'a>(v: &'a Value, what: &str, budget: &mut Budget) -> OpResult<Cow<'a, Dist>> {
880 match v {
881 Value::Dist(d) => {
882 budget.work(d.outcomes.len() as u64)?;
883 Ok(Cow::Borrowed(d))
884 }
885 Value::List(xs) if !xs.is_empty() => {
886 budget.collection(xs.len() as u128)?;
887 let comparisons = (xs.len().ilog2() + 1) as u64;
888 budget.work((xs.len() as u64).saturating_mul(comparisons))?;
889 for x in xs.iter() {
890 match x {
891 Value::Int(n) => {
892 budget.integer_work(n, &Integer::ONE, false)?;
893 }
894 Value::Str(s) => {
895 budget.string_work(s)?;
896 }
897 x if x.is_uncertain() => {
898 return Err(OpError::new(format!(
899 "`{what}` needs list elements that are values, found {}",
900 article(&x.kind())
901 ))
902 .help("draw the elements first, or explicitly build a mixture with `one_of`"));
903 }
904 _ => {}
905 }
906 }
907 Ok(Cow::Owned(Dist::uniform(xs.to_vec())))
908 }
909 _ => Err(OpError::new(format!(
910 "`{what}` needs a distribution or a nonempty list"
911 ))),
912 }
913}
914
915fn numeric_dist<'a>(v: &'a Value, what: &str, budget: &mut Budget) -> OpResult<Cow<'a, Dist>> {
916 let d = stat_dist(v, what, budget)?;
917 check_numeric(&d, what)?;
918 Ok(d)
919}
920
921fn check_numeric(d: &Dist, what: &str) -> OpResult<()> {
922 for (x, _) in &d.outcomes {
923 if !matches!(x, Value::Int(_) | Value::Float(_) | Value::Prob(_)) {
924 return Err(OpError::new(format!(
925 "`{what}` needs real numeric elements, found {}",
926 article(&x.kind())
927 )));
928 }
929 if x.as_f64().is_none_or(|x| !x.is_finite()) {
930 return Err(OpError::new(format!(
931 "`{what}` needs numbers representable as finite floats"
932 )));
933 }
934 }
935 if d.outcomes.is_empty() {
936 return Err(OpError::new(format!("`{what}` needs at least one resolved outcome")));
937 }
938 Ok(())
939}
940
941fn statistical_compare(a: &Value, b: &Value) -> OpResult<std::cmp::Ordering> {
942 match (a, b) {
943 (Value::Bool(a), Value::Bool(b)) => Ok(a.cmp(b)),
944 _ => ops::compare(a, b),
945 }
946}
947
948fn ordered_outcomes(d: &Dist, budget: &mut Budget) -> OpResult<Vec<(Value, f64)>> {
951 budget.collection(d.outcomes.len() as u128)?;
952 for (x, _) in &d.outcomes {
953 budget.work(1)?;
954 statistical_compare(x, x)?;
955 }
956 crate::ordering::reserve_sort(d.outcomes.len(), budget)?;
957 let mut outcomes = d.outcomes.clone();
958 crate::ordering::try_sort_by(&mut outcomes, |(a, _), (b, _)| {
959 budget.work(1)?;
960 statistical_compare(a, b)
961 })?;
962 Ok(outcomes)
963}
964
965fn quantile(v: &Value, q: f64, what: &str, budget: &mut Budget) -> OpResult<Value> {
966 let d = stat_dist(v, what, budget)?;
967 let outcomes = ordered_outcomes(&d, budget)?;
968 crate::stats::quantile(&outcomes, q)
969 .cloned()
970 .ok_or_else(|| OpError::new(format!("`{what}` needs at least one resolved outcome")))
971}
972
973fn median(v: &Value, b: Builtin, budget: &mut Budget) -> OpResult<Value> {
974 let d = stat_dist(v, b.name(), budget)?;
975 for (x, _) in &d.outcomes {
976 budget.work(1)?;
977 statistical_compare(x, x)?;
980 if b == Builtin::Median && !matches!(x, Value::Int(_) | Value::Float(_) | Value::Prob(_) | Value::Date(_)) {
981 return Err(OpError::new(format!(
982 "`median` needs real numeric or date elements, found {}",
983 article(&x.kind())
984 ))
985 .help("use `median_low` or `median_high` for ordered values such as strings"));
986 }
987 }
988 let outcomes = ordered_outcomes(&d, budget)?;
989 let (lo, hi) = crate::stats::median_bounds(&outcomes)
990 .ok_or_else(|| OpError::new(format!("`{}` needs at least one resolved outcome", b.name())))?;
991 match b {
992 Builtin::MedianLow => Ok(lo.clone()),
993 Builtin::MedianHigh => Ok(hi.clone()),
994 _ if lo == hi => Ok(lo.clone()),
995 _ => midpoint(lo, hi, budget),
996 }
997}
998
999fn date_mean(v: &Value, d: &Dist) -> OpResult<Value> {
1000 if let Value::List(xs) = v {
1003 let total = xs
1004 .iter()
1005 .try_fold(0i128, |sum, x| Ok::<_, OpError>(sum + date_value(x, "mean")? as i128))?;
1006 let n = xs.len() as i128;
1007 let day = total.div_euclid(n) + i128::from(2 * total.rem_euclid(n) > n);
1008 return Ok(Value::Date(day as i32));
1009 }
1010 let origin = date_value(&d.outcomes[0].0, "mean")?;
1011 let mut offsets = crate::stats::Sum::default();
1012 let mut weights = crate::stats::Sum::default();
1013 for (x, w) in &d.outcomes {
1014 offsets.add((date_value(x, "mean")? - origin) as f64 * w);
1015 weights.add(*w);
1016 }
1017 let days = offsets.value() / weights.value();
1018 let floor = days.floor();
1019 let tie_error = 4.0 * f64::EPSILON * days.abs().max(1.0);
1020 let rounded = floor + f64::from(days - floor > 0.5 + tie_error);
1021 date_result(
1022 i32::try_from(origin as i64 + rounded as i64)
1023 .ok()
1024 .filter(|d| dates::valid(*d)),
1025 )
1026}
1027
1028pub(crate) fn midpoint(a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
1030 if let (Value::Date(a), Value::Date(b)) = (a, b) {
1031 return Ok(Value::Date(((*a as i64 + *b as i64).div_euclid(2)) as i32));
1032 }
1033 if let (Value::Int(a), Value::Int(b)) = (a, b) {
1034 budget.integer_work(a, b, false)?;
1035 let sum = a.add(b)?;
1036 budget.integer_allocation(sum.bits(), 1)?;
1037 let two = Integer::from(2);
1038 budget.integer_work(&sum, &two, true)?;
1039 let (whole, remainder) = sum.div_mod(&two)?;
1040 if remainder.is_zero() {
1041 return Ok(Value::Int(whole));
1042 }
1043 return sum
1044 .ratio(&two)
1045 .map(Value::Float)
1046 .ok_or_else(|| OpError::new("median's fractional midpoint is too large for a finite float"));
1047 }
1048 let (a, b) = (number(a, "median")?, number(b, "median")?);
1049 finite_float(crate::stats::midpoint(a, b), "median")
1052}
1053
1054fn iter_items(v: &Value, budget: &mut Budget) -> OpResult<Value> {
1055 match v {
1056 Value::List(_) | Value::Range(..) => Ok(v.clone()),
1057 Value::Map(m) => Ok(Value::list(
1058 m.iter().map(|(k, x)| Value::list(vec![k.clone(), x.clone()])).collect(),
1059 )),
1060 Value::Bag(b) => {
1061 let n: u128 = b.values().map(|n| *n as u128).sum();
1062 budget.collection(n)?;
1063 Ok(Value::list(
1064 b.iter()
1065 .flat_map(|(k, n)| std::iter::repeat_n(k.clone(), *n as usize))
1066 .collect(),
1067 ))
1068 }
1069 Value::Str(s) => crate::text::chars(s, budget).map(Value::list),
1070 v if v.is_uncertain() => Err(OpError::new(format!("can't loop over a {}", v.kind()))
1071 .help("draw a value first with `~`, or loop over `support(…)`")),
1072 other => Err(OpError::new(format!("can't loop over {}", article(&other.kind())))),
1073 }
1074}
1075
1076fn date_count(v: &Value, func: &str, budget: &mut Budget) -> OpResult<i64> {
1077 integer(v, func, budget)?
1078 .to_i64()
1079 .ok_or_else(|| OpError::new("date out of range"))
1080}
1081
1082fn date_value(v: &Value, func: &str) -> OpResult<i32> {
1083 match v {
1084 Value::Date(d) if dates::valid(*d) => Ok(*d),
1085 Value::Date(_) => Err(OpError::new("date out of range")),
1086 _ => Err(expected("a date", v, func)),
1087 }
1088}
1089
1090fn date_result(date: Option<i32>) -> OpResult<Value> {
1091 date.map(Value::Date)
1092 .ok_or_else(|| OpError::new("date out of range (0001-01-01..9999-12-31)"))
1093}
1094
1095fn holiday_calendar(v: Option<&Value>, func: &str, budget: &mut Budget) -> OpResult<Vec<i32>> {
1096 let Some(v) = v else {
1097 return Ok(Vec::new());
1098 };
1099 let values = list(v, func)?;
1100 budget.collection(values.len() as u128)?;
1101 let log = (values.len() as u64).checked_ilog2().unwrap_or(0) as u64 + 1;
1102 budget.work((values.len() as u64).saturating_mul(log))?;
1103 let mut holidays = Vec::new();
1104 for value in values.iter() {
1105 let d = date_value(value, func)?;
1106 if dates::weekday(d) < 5 {
1107 holidays.push(d);
1108 }
1109 }
1110 holidays.sort_unstable();
1111 holidays.dedup();
1112 Ok(holidays)
1113}
1114
1115fn expected(what: &str, v: &Value, func: &str) -> OpError {
1116 OpError::new(format!("`{func}` needs {what}, found {}", article(&v.kind())))
1117}
1118
1119fn number(v: &Value, func: &str) -> OpResult<f64> {
1120 match v {
1121 Value::Bool(_) => Err(expected("a number", v, func)),
1122 Value::Continuous(_) => {
1123 Err(expected("a number", v, func).help("draw a value first, like `let x ~ normal(0, 1)`"))
1124 }
1125 Value::Int(n) => n
1126 .to_f64()
1127 .ok_or_else(|| OpError::new(format!("`{func}` needs an integer that fits in a finite float"))),
1128 _ => v.as_f64().ok_or_else(|| expected("a number", v, func)),
1129 }
1130}
1131
1132fn complex_number(v: &Value, func: &str) -> OpResult<Complex> {
1133 v.as_complex()
1134 .ok_or_else(|| expected("a finite real or complex number", v, func))
1135}
1136
1137fn finite_float(x: f64, func: &str) -> OpResult<Value> {
1138 if x.is_finite() {
1139 Ok(Value::Float(x))
1140 } else {
1141 Err(OpError::new(format!(
1142 "`{func}` gave a result that isn't a finite number"
1143 )))
1144 }
1145}
1146
1147fn whole(v: &Value, what: &str, budget: &mut Budget) -> OpResult<i64> {
1148 integer(v, what, budget)?
1149 .to_i64()
1150 .ok_or_else(|| OpError::new(format!("{what} is outside the supported count range")))
1151}
1152
1153fn nonnegative_int<'a>(v: &'a Value, func: &str, budget: &mut Budget) -> OpResult<Cow<'a, Integer>> {
1154 let n = integer(v, func, budget)?;
1155 if n.is_negative() {
1156 Err(OpError::new(format!("`{func}` needs nonnegative integers")))
1157 } else {
1158 Ok(n)
1159 }
1160}
1161
1162fn choose(n: &Integer, k: &Integer, budget: &mut Budget) -> OpResult<Value> {
1163 if k > n {
1164 return Ok(Value::Int(Integer::ZERO));
1165 }
1166 let k = k.min(&n.sub(k)?).clone();
1167 let steps = k
1168 .to_u64()
1169 .ok_or_else(|| OpError::limit("choose needs too many iterations"))?;
1170 budget.work(steps)?;
1171 let offset = n.sub(&k)?;
1172 let mut result = Integer::ONE;
1173 for i in 1..=steps {
1174 let numerator = offset.add(&i.into())?;
1175 let divisor = Integer::from(i);
1176 let d = gcd(numerator.clone(), divisor.clone(), budget)?;
1179 let numerator = numerator.div_mod(&d)?.0;
1180 let divisor = divisor.div_mod(&d)?.0;
1181 budget.integer_work(&result, &divisor, true)?;
1182 result = result.div_mod(&divisor)?.0;
1183 budget.integer_work(&result, &numerator, true)?;
1184 result = result.mul(&numerator)?;
1185 budget.integer_bits(result.bits())?;
1186 }
1187 Ok(Value::Int(result))
1188}
1189
1190fn factorial(n: &Integer, budget: &mut Budget) -> OpResult<Value> {
1191 let n = n
1192 .to_u64()
1193 .ok_or_else(|| OpError::limit("factorial exceeds the integer size limit"))?;
1194 let half = n / 2;
1196 let lower_bits = half.saturating_mul(63u64.saturating_sub(u64::from(half.leading_zeros())));
1197 budget.integer_bits(lower_bits)?;
1198 budget.work(n)?;
1199 let mut result = Integer::ONE;
1200 for i in 2..=n {
1201 let factor = Integer::from(i);
1202 budget.integer_work(&result, &factor, true)?;
1203 result = result.mul(&factor)?;
1204 budget.integer_bits(result.bits())?;
1205 }
1206 Ok(Value::Int(result))
1207}
1208
1209fn gcd(mut a: Integer, mut b: Integer, budget: &mut Budget) -> OpResult<Integer> {
1210 while !b.is_zero() {
1211 budget.integer_work(&a, &b, true)?;
1212 let r = a.div_mod(&b)?.1;
1213 a = b;
1214 b = r;
1215 }
1216 Ok(a)
1217}
1218
1219fn euler_phi(n: &Integer, budget: &mut Budget) -> OpResult<Value> {
1220 if n.is_zero() {
1221 return Err(OpError::new("`euler_phi` needs a positive integer"));
1222 }
1223 let mut n = n.clone();
1224 let mut result = n.clone();
1225 let mut divisor = Integer::from(2);
1226 loop {
1227 budget.integer_work(&n, &divisor, true)?;
1228 let (q, rem) = n.div_mod(&divisor)?;
1229 if divisor > q {
1230 break;
1231 }
1232 if rem.is_zero() {
1233 budget.integer_work(&result, &divisor, true)?;
1234 result = result.sub(&result.div_mod(&divisor)?.0)?;
1235 n = q;
1236 loop {
1237 budget.integer_work(&n, &divisor, true)?;
1238 let (q, rem) = n.div_mod(&divisor)?;
1239 if !rem.is_zero() {
1240 break;
1241 }
1242 n = q;
1243 }
1244 }
1245 divisor = divisor.add(&if divisor == 2 { Integer::ONE } else { 2.into() })?;
1246 }
1247 if n > 1 {
1248 budget.integer_work(&result, &n, true)?;
1249 result = result.sub(&result.div_mod(&n)?.0)?;
1250 }
1251 Ok(Value::Int(result))
1252}
1253
1254fn text(v: &Value, func: &str) -> OpResult<Arc<str>> {
1255 match v {
1256 Value::Str(s) => Ok(s.clone()),
1257 other => Err(expected("a string", other, func)),
1258 }
1259}
1260
1261fn list<'a>(v: &'a Value, func: &str) -> OpResult<&'a [Value]> {
1262 match v {
1263 Value::List(items) => Ok(items),
1264 other => Err(expected("a list", other, func)),
1265 }
1266}
1267
1268fn range_items(lo: &Integer, hi: &Integer, budget: &mut Budget) -> OpResult<Vec<Value>> {
1269 let n = range_count(lo, hi)?;
1270 budget.collection(n)?;
1271 let n = usize::try_from(n).map_err(|_| OpError::limit("range has too many elements"))?;
1272 budget.integer_allocation(lo.bits().max(hi.bits()), n as u64)?;
1273 budget.work((n as u64).saturating_mul(lo.bits().max(hi.bits()).div_ceil(64).max(1)))?;
1274 (0..n)
1275 .map(|i| lo.add(&Integer::from(i)).map(Value::Int).map_err(OpError::from))
1276 .collect()
1277}
1278
1279pub fn items(v: &Value, func: &str, budget: &mut Budget) -> OpResult<Vec<Value>> {
1282 match v {
1283 Value::List(items) => {
1284 budget.collection(items.len() as u128)?;
1285 budget.work(items.len() as u64)?;
1286 Ok(items.to_vec())
1287 }
1288 Value::Range(lo, hi) => range_items(lo, hi, budget),
1289 Value::Str(s) => crate::text::chars(s, budget),
1290 other => Err(expected("a list, range or string", other, func)),
1291 }
1292}
1293
1294fn sort(items: &mut [Value], descending: bool, budget: &mut Budget) -> OpResult<()> {
1295 for x in items.iter() {
1298 budget.work(1)?;
1299 ops::compare(x, x)?;
1300 }
1301 crate::ordering::reserve_sort(items.len(), budget)?;
1302 crate::ordering::try_sort_by(items, |x, y| {
1303 budget.work(1)?;
1304 ops::compare(x, y).map(|c| if descending { c.reverse() } else { c })
1305 })
1306}
1307
1308fn elementary1(
1311 v: &Value,
1312 func: &str,
1313 complex: fn(Complex) -> OpResult<Complex>,
1314 real: impl Fn(f64) -> Option<f64>,
1315) -> OpResult<Value> {
1316 match v {
1317 Value::Complex(z) => complex(*z)
1318 .map(Value::Complex)
1319 .map_err(|_| OpError::new(format!("`{func}` isn't defined for {v} or its result isn't finite"))),
1320 _ => float1(v, func, real),
1321 }
1322}
1323
1324fn float1(v: &Value, func: &str, f: impl Fn(f64) -> Option<f64>) -> OpResult<Value> {
1325 let x = number(v, func)?;
1326 f(x).filter(|y| x.is_finite() && y.is_finite())
1327 .map(Value::Float)
1328 .ok_or_else(|| OpError::new(format!("`{func}` isn't defined for {}", fmt_float(x))))
1329}
1330
1331fn float2(a: &Value, b: &Value, func: &str, f: fn(f64, f64) -> f64) -> OpResult<Value> {
1332 let (a, b) = (number(a, func)?, number(b, func)?);
1333 if !a.is_finite() || !b.is_finite() {
1334 return Err(OpError::new(format!("`{func}` needs finite numbers")));
1335 }
1336 let a = if a == 0.0 { 0.0 } else { a };
1339 let b = if b == 0.0 { 0.0 } else { b };
1340 let value = f(a, b);
1341 if !value.is_finite() {
1342 return Err(OpError::new(format!(
1343 "`{func}` gave a result that isn't a finite number"
1344 )));
1345 }
1346 Ok(Value::Float(value))
1347}
1348
1349fn to_int(v: &Value, f: fn(f64) -> f64) -> OpResult<Value> {
1350 match v {
1351 Value::Int(i) => Ok(Value::Int(i.clone())),
1352 other => {
1353 let x = number(other, "rounding")?;
1354 let r = f(x);
1355 Integer::from_f64(r)
1356 .map(Value::Int)
1357 .ok_or_else(|| OpError::new(format!("{} cannot be rounded to a finite int", fmt_float(x))))
1358 }
1359 }
1360}
1361
1362fn round(v: &Value, digits: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1363 let Some(digits) = digits else {
1364 return to_int(v, f64::round);
1365 };
1366 let digits = integer(digits, "round's digits", budget)?;
1367 if let Value::Int(n) = v {
1368 if *digits >= 0 {
1369 return Ok(v.clone());
1370 }
1371 let places = digits.abs();
1372 if places > Integer::from(probl_number::MAX_INTEGER_DIGITS) {
1373 return Ok(Value::Int(Integer::ZERO));
1374 }
1375 budget.integer_work(n, n, true)?;
1376 return Ok(Value::Int(n.round_decimal(places.to_u64().unwrap() as u32)?));
1377 }
1378 let digits = if *digits > 323 {
1380 324
1381 } else if *digits < -308 {
1382 -309
1383 } else {
1384 digits.to_i64().unwrap()
1385 };
1386 float1(v, "round", |x| {
1387 if digits > 323 || x == 0.0 {
1390 return Some(x);
1391 }
1392 if digits < -308 {
1393 return Some(0.0);
1394 }
1395 let rounded = if digits >= 0 {
1396 let high = libm::pow(10.0, digits.min(308) as f64);
1399 let low = libm::pow(10.0, (digits - 308).max(0) as f64);
1400 let scaled = (x * high) * low;
1401 if scaled.abs() >= 1e16 || x.fract() == 0.0 {
1404 return Some(x);
1405 }
1406 (scaled.round() / low) / high
1407 } else {
1408 let scale = libm::pow(10.0, -digits as f64);
1409 let scaled = x / scale;
1410 if scaled.abs() >= 1e16 {
1411 return Some(x);
1412 }
1413 scaled.round() * scale
1414 };
1415 Some(rounded)
1416 })
1417}
1418
1419fn min_max(args: &[Value], want_max: bool, budget: &mut Budget) -> OpResult<Value> {
1420 let name = if want_max { "max" } else { "min" };
1421 let mut best: Option<Value> = None;
1422 for v in args {
1423 let v = ops::lift1(v, budget, |v, budget| {
1426 budget.work(1)?;
1427 ops::compare(v, v)?;
1428 Ok(v.clone())
1429 })?;
1430 best = Some(match best {
1431 None => v,
1432 Some(b) => ops::lift2(&b, &v, budget, |b, v, _| {
1433 let ord = ops::compare(v, b)?;
1434 if (want_max && ord.is_gt()) || (!want_max && ord.is_lt()) {
1435 Ok(v.clone())
1436 } else {
1437 Ok(b.clone())
1438 }
1439 })?,
1440 });
1441 }
1442 best.ok_or_else(|| OpError::new(format!("`{name}` of an empty list")))
1443}
1444
1445fn slice(v: &Value, start: &Value, end: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1446 let start = integer(start, "slice", budget)?;
1447 let end = end.map(|v| integer(v, "slice", budget)).transpose()?;
1448 let length: Integer = match v {
1449 Value::List(xs) => xs.len().into(),
1450 Value::Str(s) => s.chars().count().into(),
1451 Value::Range(lo, hi) => {
1452 budget.integer_work(lo, hi, false)?;
1453 range_len(lo, hi)?
1454 }
1455 other => return Err(expected("a list, range or string", other, "slice")),
1456 };
1457 let end = end.as_deref().unwrap_or(&length);
1458 let start = start.as_ref();
1459 if start.is_negative() || start > end || end > &length {
1460 return Err(OpError::new("`slice` needs 0 <= start <= end <= length"));
1461 }
1462 Ok(match v {
1463 Value::Range(lo, _) => {
1464 let (first, last) = if start == end {
1467 (Integer::ZERO, (-1).into())
1468 } else {
1469 (lo.add(start)?, lo.add(&end.sub(&Integer::ONE)?)?)
1470 };
1471 budget.integer_allocation(first.bits(), 1)?;
1472 budget.integer_allocation(last.bits(), 1)?;
1473 Value::Range(first, last)
1474 }
1475 Value::List(xs) => {
1476 let (start, end) = (start.to_u64().unwrap() as usize, end.to_u64().unwrap() as usize);
1477 budget.collection((end - start) as u128)?;
1478 budget.work((end - start) as u64)?;
1479 Value::list(xs[start..end].to_vec())
1480 }
1481 Value::Str(s) => {
1482 let (start, end) = (start.to_u64().unwrap() as usize, end.to_u64().unwrap() as usize);
1483 let boundary = |n| s.char_indices().nth(n).map_or(s.len(), |(i, _)| i);
1484 crate::text::value(&s[boundary(start)..boundary(end)], budget)?
1485 }
1486 _ => unreachable!(),
1487 })
1488}
1489
1490fn len(v: &Value) -> OpResult<Value> {
1491 let n: u128 = match v {
1492 Value::List(items) => items.len() as u128,
1493 Value::Str(s) => s.chars().count() as u128,
1494 Value::Map(m) => m.len() as u128,
1495 Value::Bag(b) => b.values().map(|n| *n as u128).sum(),
1496 Value::Range(lo, hi) => return Ok(Value::Int(range_len(lo, hi)?)),
1497 other => return Err(expected("a collection", other, "len")),
1498 };
1499 Ok(Value::Int(n.into()))
1500}
1501
1502fn sum(v: &Value, budget: &mut Budget) -> OpResult<Value> {
1503 let items = items(v, "sum", budget)?;
1504 let mut acc = Value::Int(0.into());
1505 for x in &items {
1506 acc = ops::binary(probl_syntax::ast::BinOp::Add, &acc, x, budget)?;
1507 }
1508 Ok(acc)
1509}
1510
1511fn get(coll: &Value, key: &Value, default: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1512 let found = match coll {
1513 Value::Map(m) => m.get(key).cloned(),
1514 Value::List(items) => integer(key, "index", budget)?
1515 .to_u64()
1516 .and_then(|n| usize::try_from(n).ok())
1517 .and_then(|n| items.get(n))
1518 .cloned(),
1519 Value::Str(s) => {
1520 let index = integer(key, "index", budget)?
1521 .to_u64()
1522 .and_then(|n| usize::try_from(n).ok());
1523 budget.string_work(s)?;
1524 index
1525 .and_then(|n| s.chars().nth(n))
1526 .map(|c| crate::text::value(c.encode_utf8(&mut [0; 4]), budget))
1527 .transpose()?
1528 }
1529 Value::Range(lo, hi) => {
1530 let index = integer(key, "index", budget)?;
1531 budget.integer_work(lo, hi, false)?;
1532 let len = range_len(lo, hi)?;
1533 if index.is_negative() || *index >= len {
1534 None
1535 } else {
1536 Some(Value::Int(lo.add(&index)?))
1537 }
1538 }
1539 Value::Bag(b) => Some(Value::Int(b.get(key).copied().unwrap_or(0).into())),
1540 other => return Err(expected("a map, bag or sequence", other, "get")),
1541 };
1542 match (found, default) {
1543 (Some(v), _) => Ok(v),
1544 (None, Some(d)) => Ok(d.clone()),
1545 (None, None) => {
1546 Err(OpError::new(format!("the key {key:?} isn't there")).help("give a default: `get(key, default)`"))
1547 }
1548 }
1549}
1550
1551pub fn population_extreme(v: &Value, want_max: bool, default: Option<&Value>, budget: &mut Budget) -> OpResult<Value> {
1553 let name = if want_max { "maximum" } else { "minimum" };
1554 match v {
1555 Value::Range(lo, hi) => {
1556 budget.integer_work(lo, hi, false)?;
1557 if hi < lo {
1558 return empty_extreme(name, default);
1559 }
1560 Ok(Value::Int(if want_max { hi.clone() } else { lo.clone() }))
1561 }
1562 Value::Continuous(f) => {
1563 let (lo, hi) = f.support();
1564 finite_bound(if want_max { hi } else { lo }, name)
1565 }
1566 Value::Dist(d) => {
1567 if d.missing > 0.0 {
1568 return Err(OpError::new(format!(
1569 "`{name}` cannot determine a support bound while the distribution has unresolved mass"
1570 )));
1571 }
1572 let mut best = None;
1573 for (x, p) in &d.outcomes {
1574 if *p <= 0.0 {
1575 continue;
1576 }
1577 let x = match x {
1578 Value::Continuous(_) | Value::Dist(_) => population_extreme(x, want_max, None, budget)?,
1579 x => x.clone(),
1580 };
1581 select_extreme(&mut best, x, want_max, budget)?;
1582 }
1583 best.ok_or_else(|| OpError::new(format!("`{name}` needs a nonempty distribution")))
1584 }
1585 Value::List(_) | Value::Str(_) => {
1586 let mut best = None;
1587 for x in items(v, name, budget)? {
1588 select_extreme(&mut best, x, want_max, budget)?;
1589 }
1590 best.map_or_else(|| empty_extreme(name, default), Ok)
1591 }
1592 other => Err(expected(
1593 "a distribution or nonempty list, range or string",
1594 other,
1595 name,
1596 )),
1597 }
1598}
1599
1600pub fn empty_extreme(name: &str, default: Option<&Value>) -> OpResult<Value> {
1601 default.cloned().ok_or_else(|| {
1602 OpError::new(format!("`{name}` needs a nonempty collection or a default"))
1603 .help(format!("use `{name}(xs, default: value)` to handle empty collections"))
1604 })
1605}
1606
1607fn finite_bound(x: f64, name: &str) -> OpResult<Value> {
1608 if !x.is_finite() {
1609 return Err(OpError::new(format!("`{name}` has no finite support bound")));
1610 }
1611 Ok(Value::Float(x))
1612}
1613
1614fn select_extreme(best: &mut Option<Value>, x: Value, want_max: bool, budget: &mut Budget) -> OpResult<()> {
1615 budget.work(1)?;
1616 ops::compare(&x, &x)?;
1617 let replace = match best {
1618 None => true,
1619 Some(b) => {
1620 let ord = ops::compare(&x, b)?;
1621 if want_max { ord.is_gt() } else { ord.is_lt() }
1622 }
1623 };
1624 if replace {
1625 *best = Some(x);
1626 }
1627 Ok(())
1628}
1629
1630pub fn extreme_count(n: &Value, budget: &mut Budget) -> OpResult<usize> {
1631 let n = integer(n, "the count", budget)?;
1632 if n.is_negative() {
1633 return Err(OpError::new("the count must be nonnegative"));
1634 }
1635 Ok(n.to_u64().and_then(|n| usize::try_from(n).ok()).unwrap_or(usize::MAX))
1636}
1637
1638fn extremes(v: &Value, n: &Value, highest: bool, budget: &mut Budget) -> OpResult<Value> {
1639 let n = extreme_count(n, budget)?;
1640 let mut items = items(v, if highest { "highest" } else { "lowest" }, budget)?;
1641 sort(&mut items, highest, budget)?;
1642 items.truncate(n);
1643 Ok(Value::list(items))
1644}
1645
1646fn insert(coll: &Value, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
1647 match coll {
1648 Value::List(items) => {
1649 budget.collection(items.len() as u128 + 1)?;
1650 let mut items = items.to_vec();
1651 let i = as_index(a, items.len() as u128 + 1, budget)? as usize;
1652 items.insert(i, b.clone());
1653 Ok(Value::list(items))
1654 }
1655 Value::Map(m) => {
1656 budget.collection(m.len() as u128 + 1)?;
1657 let mut m = BTreeMap::clone(m);
1658 m.insert(a.clone(), b.clone());
1659 Ok(Value::map(m))
1660 }
1661 other => Err(expected("a list or a map", other, "insert")),
1662 }
1663}
1664
1665fn remove(coll: &Value, key: &Value, budget: &mut Budget) -> OpResult<Value> {
1666 match coll {
1667 Value::List(items) => {
1668 let mut items = items.to_vec();
1669 let i = as_index(key, items.len() as u128, budget)? as usize;
1670 items.remove(i);
1671 Ok(Value::list(items))
1672 }
1673 Value::Map(m) => {
1674 let mut m = BTreeMap::clone(m);
1675 m.remove(key);
1676 Ok(Value::map(m))
1677 }
1678 Value::Bag(b) => match b.without(key) {
1679 Some(rest) => Ok(Value::multiset(rest)),
1680 None => Err(OpError::new(format!("{key:?} isn't in the bag"))),
1681 },
1682 other => Err(expected("a list, map or bag", other, "remove")),
1683 }
1684}
1685
1686fn one_of(v: &Value, budget: &mut Budget) -> OpResult<Value> {
1689 match v {
1690 Value::List(items) if !items.is_empty() => {
1691 budget.outcomes(items.len() as u128)?;
1692 let p = 1.0 / items.len() as f64;
1693 ops::combine(items.iter().map(|x| (x.clone(), p)).collect(), 0.0, budget)
1694 }
1695 Value::Range(lo, hi) if hi >= lo => {
1696 budget.outcomes(range_count(lo, hi)?)?;
1697 Ok(Dist::uniform(range_items(lo, hi, budget)?).into_value())
1698 }
1699 Value::Map(m) if !m.is_empty() => {
1700 let all_probs = m.values().all(|w| matches!(w, Value::Prob(_)));
1701 if !all_probs && m.values().any(|w| matches!(w, Value::Prob(_))) {
1702 return Err(
1703 OpError::new("one_of can't mix probabilities and relative numeric weights")
1704 .help("use prob(...) for every absolute probability, or numbers for every relative weight"),
1705 );
1706 }
1707 let mut pairs = Vec::new();
1708 for (k, w) in m.iter() {
1709 let w = match w {
1710 Value::Bool(_) => None,
1711 _ => w.as_f64().filter(|w| *w >= 0.0 && w.is_finite()),
1712 }
1713 .ok_or_else(|| {
1714 OpError::new(format!(
1715 "one_of needs weights that are numbers of 0 or more, found {w:?}"
1716 ))
1717 })?;
1718 pairs.push((k.clone(), w));
1719 }
1720 let largest = pairs.iter().map(|(_, w)| *w).fold(0.0, f64::max);
1721 if largest == 0.0 {
1722 return Err(OpError::new("one_of needs at least one positive weight"));
1723 }
1724 let total = crate::stats::sum(pairs.iter().map(|(_, w)| *w));
1727 if all_probs && (total - 1.0).abs() > 1e-9 {
1728 return Err(OpError::new(format!(
1729 "the chances add up to {}, not 100%",
1730 crate::value::fmt_prob(total)
1731 ))
1732 .help("use plain numbers for relative weights, like [\"a\": 3, \"b\": 1]"));
1733 }
1734 for (_, w) in &mut pairs {
1735 *w /= largest;
1736 }
1737 let total = crate::stats::sum(pairs.iter().map(|(_, w)| *w));
1738 ops::combine(pairs.into_iter().map(|(k, w)| (k, w / total)).collect(), 0.0, budget)
1739 }
1740 Value::Bag(b) => {
1741 let total: u128 = b.values().map(|n| *n as u128).sum();
1742 if total == 0 {
1743 return Err(OpError::new("the bag is empty"));
1744 }
1745 let pairs = b.iter().map(|(k, n)| (k.clone(), *n as f64 / total as f64)).collect();
1746 ops::combine(pairs, 0.0, budget)
1747 }
1748 Value::List(_) | Value::Range(..) | Value::Map(_) => Err(OpError::new("one_of needs at least one option")),
1749 other => Err(expected("a list, range, map or bag", other, "one_of")),
1750 }
1751}
1752
1753fn bag(v: &Value, budget: &mut Budget) -> OpResult<Value> {
1754 let mut counts = BTreeMap::new();
1755 match v {
1756 Value::Map(m) => {
1757 for (k, n) in m.iter() {
1758 let n = integer(n, "bag count", budget)?;
1759 if n.is_negative() {
1760 return Err(OpError::new("bag counts must be whole numbers of 0 or more"));
1761 }
1762 let n = n
1763 .to_u64()
1764 .ok_or_else(|| OpError::new("bag count exceeds the supported count range"))?;
1765 if n > 0 {
1766 counts.insert(k.clone(), n);
1767 }
1768 }
1769 }
1770 Value::List(items) => {
1771 for item in items.iter() {
1772 *counts.entry(item.clone()).or_insert(0) += 1;
1773 }
1774 }
1775 other => return Err(expected("a map of counts or a list", other, "bag")),
1776 }
1777 Ok(Value::bag(counts))
1778}