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