1use crate::dist::{Budget, Dist};
4use crate::error::{Fault, OpError, OpResult};
5use crate::value::{EnumValue, Record, Value};
6use probl_number::Integer;
7use probl_syntax::ast::{BinOp, UnOp};
8use std::borrow::Cow;
9use std::sync::Arc;
10
11pub fn integer<'a>(v: &'a Value, context: &str, budget: &mut Budget) -> OpResult<Cow<'a, Integer>> {
14 match v {
15 Value::Int(n) => Ok(Cow::Borrowed(n)),
16 Value::Float(f) if f.is_finite() && f.fract() == 0.0 => {
17 let n = Integer::from_f64(*f).expect("finite integral float fits the integer hard limit");
18 budget.integer_allocation(n.bits(), 1)?;
19 budget.work(n.bits().div_ceil(64).max(1))?;
20 Ok(Cow::Owned(n))
21 }
22 Value::Analytic(a) if a.int => Err(crate::analytic::unsupported(&format!(
23 "an int outcome of a continuous draw as `{context}`"
24 ))),
25 _ => {
26 let found = match v {
27 Value::Float(f) => format!("the float {f:?}"),
28 _ => article(&v.kind()),
29 };
30 Err(OpError::new(format!(
31 "`{context}` needs an int (or an exactly integral finite float), found {found}"
32 )))
33 }
34 }
35}
36
37#[cold]
41#[inline(never)]
42pub fn listed<'a>(v: &'a Value, budget: &mut Budget) -> OpResult<std::borrow::Cow<'a, Value>> {
43 match v {
44 Value::Counts(c) => Ok(std::borrow::Cow::Owned(c.list(budget)?.into_value())),
45 v => Ok(std::borrow::Cow::Borrowed(v)),
46 }
47}
48
49pub fn outcomes(v: &Value) -> std::borrow::Cow<'_, [(Value, f64)]> {
51 match v {
52 Value::Dist(d) => std::borrow::Cow::Borrowed(&d.outcomes),
53 other => std::borrow::Cow::Owned(vec![(other.clone(), 1.0)]),
54 }
55}
56
57fn missing(v: &Value) -> f64 {
58 match v {
59 Value::Dist(d) => d.missing,
60 _ => 0.0,
61 }
62}
63
64pub fn combine(results: Vec<(Value, f64)>, missing: f64, budget: &mut Budget) -> OpResult<Value> {
69 if results.is_empty() {
70 return Err(OpError::new("the result would be an empty distribution"));
71 }
72 let mut flat = Vec::with_capacity(results.len());
73 let mut missing = missing;
74 for (v, w) in results {
75 let v = if let Value::Counts(_) = v {
76 listed(&v, budget)?.into_owned()
77 } else {
78 v
79 };
80 match v {
81 Value::Dist(d) => {
82 missing += w * d.missing;
83 flat.extend(d.outcomes.iter().map(|(x, p)| (x.clone(), w * p)));
84 }
85 other => flat.push((other, w)),
86 }
87 }
88 budget.outcomes(flat.len() as u128)?;
89 budget.work(flat.len() as u64)?;
90 Ok(Dist::from_pairs(flat, missing.min(1.0)).into_value())
91}
92
93pub fn lift1(a: &Value, budget: &mut Budget, f: impl Fn(&Value, &mut Budget) -> OpResult<Value>) -> OpResult<Value> {
95 if !a.is_dist() && !matches!(a, Value::Counts(_)) {
96 return f(a, budget);
97 }
98 let a = listed(a, budget)?;
99 let a = a.as_ref();
100 let mut results = Vec::new();
101 for (v, w) in outcomes(a).iter() {
102 results.push((f(v, budget)?, *w));
103 }
104 combine(results, missing(a), budget)
105}
106
107pub fn lift2(
109 a: &Value,
110 b: &Value,
111 budget: &mut Budget,
112 f: impl Fn(&Value, &Value, &mut Budget) -> OpResult<Value>,
113) -> OpResult<Value> {
114 let plain = |v: &Value| !v.is_dist() && !matches!(v, Value::Counts(_));
115 if plain(a) && plain(b) {
116 return f(a, b, budget);
117 }
118 let (a, b) = (listed(a, budget)?, listed(b, budget)?);
119 let (a, b) = (a.as_ref(), b.as_ref());
120 let (oa, ob) = (outcomes(a), outcomes(b));
121 budget.outcomes(oa.len() as u128 * ob.len() as u128)?;
122 budget.work((oa.len() * ob.len()) as u64)?;
123 let mut results = Vec::with_capacity(oa.len() * ob.len());
124 for (x, p) in oa.iter() {
125 for (y, q) in ob.iter() {
126 results.push((f(x, y, budget)?, p * q));
127 }
128 }
129 let m = missing(a) + (1.0 - missing(a)) * missing(b);
130 combine(results, m, budget)
131}
132
133pub type NaryFn<'a> = &'a dyn Fn(&[Value], &mut Budget) -> OpResult<Value>;
136
137pub fn lift_n(args: &[Value], budget: &mut Budget, f: NaryFn) -> OpResult<Value> {
139 if !args.iter().any(|a| a.is_dist() || matches!(a, Value::Counts(_))) {
140 return f(args, budget);
141 }
142 let listed_args;
143 let args = if args.iter().any(|a| matches!(a, Value::Counts(_))) {
144 listed_args = args
145 .iter()
146 .map(|a| Ok(listed(a, budget)?.into_owned()))
147 .collect::<OpResult<Vec<_>>>()?;
148 &listed_args[..]
149 } else {
150 args
151 };
152 let size = args
153 .iter()
154 .fold(1u128, |acc, a| acc.saturating_mul(outcomes(a).len() as u128));
155 budget.outcomes(size)?;
156 budget.work(size as u64)?;
157 let missing = args.iter().fold(0.0, |m, a| m + (1.0 - m) * missing(a));
158 let mut results = Vec::new();
159 let mut current = Vec::with_capacity(args.len());
160 product(args, 0, &mut current, 1.0, f, budget, &mut results)?;
161 combine(results, missing, budget)
162}
163
164fn product(
165 args: &[Value],
166 i: usize,
167 current: &mut Vec<Value>,
168 weight: f64,
169 f: NaryFn,
170 budget: &mut Budget,
171 results: &mut Vec<(Value, f64)>,
172) -> OpResult<()> {
173 if i == args.len() {
174 results.push((f(current, budget)?, weight));
175 return Ok(());
176 }
177 for (v, w) in outcomes(&args[i]).iter() {
178 current.push(v.clone());
179 product(args, i + 1, current, weight * w, f, budget, results)?;
180 current.pop();
181 }
182 Ok(())
183}
184
185pub fn article(kind: &str) -> String {
188 let vowel = kind.starts_with(['a', 'e', 'i', 'o', 'u']);
189 format!("{} {kind}", if vowel { "an" } else { "a" })
190}
191
192pub fn to_prob(v: &Value) -> OpResult<f64> {
195 if crate::analytic::contains(v) {
196 return Err(crate::analytic::unsupported(
197 "using an analytic outcome as a probability parameter",
198 ));
199 }
200 match v {
201 Value::Prob(_) | Value::Float(_) | Value::Int(_) => {
203 match make_prob(v).map_err(|e| e.as_fault(Fault::DomainError))? {
204 Value::Prob(p) => Ok(p),
205 _ => unreachable!("make_prob returns a probability"),
206 }
207 }
208 Value::Bool(_) => Err(OpError::new("expected a probability, found a fact (true or false)")
209 .help("convert a boolean explicitly with `prob(fact)`")),
210 Value::Dist(_) => Err(OpError::new(format!("expected a probability, found a {}", v.kind()))
211 .help("`P(…)` gives the probability that a distribution of facts is true")),
212 other => Err(OpError::new(format!(
213 "expected a probability, found {}",
214 article(&other.kind())
215 ))),
216 }
217}
218
219pub fn make_prob(v: &Value) -> OpResult<Value> {
221 let p = match v {
222 Value::Bool(b) => {
223 if *b {
224 1.0
225 } else {
226 0.0
227 }
228 }
229 Value::Prob(p) | Value::Float(p) => *p,
230 Value::Int(n) if *n == 0 => 0.0,
231 Value::Int(n) if *n == 1 => 1.0,
232 Value::Int(_) => {
233 return Err(OpError::fault(
234 Fault::ConversionError,
235 "prob needs a finite number between 0 and 1",
236 ));
237 }
238 _ => {
239 return Err(
240 OpError::new(format!("prob needs a number or bool, found {}", article(&v.kind())))
241 .help("draw distribution outcomes explicitly; use `P(d)` to query a boolean distribution"),
242 );
243 }
244 };
245 if !p.is_finite() || !(0.0..=1.0).contains(&p) {
246 return Err(OpError::fault(
247 Fault::ConversionError,
248 "prob needs a finite number between 0 and 1",
249 ));
250 }
251 Ok(Value::Prob(p))
252}
253
254pub(crate) fn computed_prob(p: f64, what: &str) -> OpResult<Value> {
258 const ROUNDING: f64 = 8.0 * f64::EPSILON;
259 if !p.is_finite() || !(-ROUNDING..=1.0 + ROUNDING).contains(&p) {
260 return Err(OpError::new(format!(
261 "`{what}` could not compute a finite probability between 0 and 1"
262 )));
263 }
264 Ok(Value::Prob(if p <= 0.0 { 0.0 } else { p.min(1.0) }))
265}
266
267pub fn fact(v: &Value, context: &str) -> OpResult<bool> {
268 match v {
269 Value::Bool(b) => Ok(*b),
270 _ => Err(
271 OpError::new(format!("`{context}` needs a bool, found {}", article(&v.kind())))
272 .help("draw an outcome first with `~`; for a probability, draw from `bernoulli(p)`"),
273 ),
274 }
275}
276
277#[derive(Clone, Copy, Debug, PartialEq)]
279pub struct Condition {
280 pub yes: f64,
282 pub no: f64,
284 pub missing: f64,
287}
288
289pub enum Truth {
291 Fact(bool),
292 Analytic(Arc<crate::analytic::Event>),
293 Probability(f64),
294 Uncertain(Arc<Dist>),
295}
296
297pub fn truth(v: &Value, op: &str) -> OpResult<Truth> {
299 match v {
300 Value::Bool(b) => Ok(Truth::Fact(*b)),
301 Value::Event(e) => Ok(Truth::Analytic(e.clone())),
302 Value::Dist(d) if d.truth().is_some() => Ok(Truth::Uncertain(d.clone())),
303 Value::Prob(_) => Ok(Truth::Probability(to_prob(v)?)),
304 other => Err(OpError::new(format!(
305 "`{op}` needs bool, prob or dist[bool], found {}",
306 article(&other.kind())
307 ))),
308 }
309}
310
311pub fn not(v: &Value, budget: &mut Budget) -> OpResult<Value> {
312 match truth(v, "not")? {
313 Truth::Fact(b) => Ok(Value::Bool(!b)),
314 Truth::Analytic(e) => {
315 budget.work(e.size() as u64)?;
316 Ok(e.complement()?.value())
317 }
318 Truth::Probability(p) => computed_prob(1.0 - p, "not"),
319 Truth::Uncertain(d) => lift1(&Value::Dist(d), budget, |x, _| match x {
320 Value::Bool(b) => Ok(Value::Bool(!b)),
321 _ => unreachable!("checked by `truth`"),
322 }),
323 }
324}
325
326pub fn logic(and: bool, a: Truth, b: Truth, budget: &mut Budget) -> OpResult<Value> {
328 if matches!(a, Truth::Analytic(_)) || matches!(b, Truth::Analytic(_)) {
329 let value = |t| match t {
330 Truth::Fact(b) => Value::Bool(b),
331 Truth::Probability(p) => Value::Prob(p),
332 Truth::Uncertain(d) => Value::Dist(d),
333 Truth::Analytic(e) => Value::Event(e),
334 };
335 let (a, b) = (value(a), value(b));
336 for v in [&a, &b] {
337 if let Value::Event(e) = v {
338 budget.collection(e.size() as u128)?;
339 budget.work(e.size() as u64)?;
340 }
341 }
342 return crate::analytic::logic(and, &a, &b);
343 }
344 let op = |x: bool, y: bool| if and { x && y } else { x || y };
345 match (a, b) {
346 (Truth::Fact(x), Truth::Fact(y)) => Ok(Value::Bool(op(x, y))),
347 (Truth::Probability(p), Truth::Probability(q)) => computed_prob(
348 if and { p * q } else { p + (1.0 - p) * q },
349 if and { "and" } else { "or" },
350 ),
351 (Truth::Probability(p), Truth::Fact(b)) | (Truth::Fact(b), Truth::Probability(p)) => {
352 Ok(Value::Prob(if b == and {
353 p
354 } else if b {
355 1.0
356 } else {
357 0.0
358 }))
359 }
360 (Truth::Fact(x), Truth::Uncertain(d)) | (Truth::Uncertain(d), Truth::Fact(x)) => {
361 lift1(&Value::Dist(d), budget, |v, _| match v {
362 Value::Bool(y) => Ok(Value::Bool(op(x, *y))),
363 _ => unreachable!("checked by `truth`"),
364 })
365 }
366 (a, b) => {
367 let law = |v| match v {
368 Truth::Uncertain(d) => Value::Dist(d),
369 Truth::Probability(p) => Dist::bernoulli(p).into_value(),
370 Truth::Fact(b) => Value::Bool(b),
371 Truth::Analytic(_) => unreachable!("handled above"),
372 };
373 lift2(&law(a), &law(b), budget, |x, y, _| match (x, y) {
374 (Value::Bool(x), Value::Bool(y)) => Ok(Value::Bool(op(*x, *y))),
375 _ => unreachable!("checked by `truth`"),
376 })
377 }
378 }
379}
380
381pub fn boolean_law(v: &Value) -> OpResult<Value> {
383 match truth(v, "observe ~")? {
384 Truth::Analytic(e) => Ok(Value::Event(e)),
385 Truth::Fact(b) => Ok(Dist::bernoulli(if b { 1.0 } else { 0.0 }).into_value()),
386 Truth::Probability(p) => Ok(Dist::bernoulli(p).into_value()),
387 Truth::Uncertain(d) => Ok(Value::Dist(d)),
388 }
389}
390
391pub fn condition(v: &Value) -> OpResult<Condition> {
392 if matches!(v, Value::Analytic(_)) {
393 return Err(crate::analytic::unsupported(
394 "using a continuous outcome as a probability condition",
395 ));
396 }
397 if matches!(v, Value::Float(_) | Value::Int(_)) {
398 let p = to_prob(v)?;
399 return Ok(Condition {
400 yes: p,
401 no: 1.0 - p,
402 missing: 0.0,
403 });
404 }
405 match truth(v, "condition")? {
406 Truth::Analytic(e) => Ok(Condition {
407 yes: e.probability(),
408 no: 1.0 - e.probability(),
409 missing: 0.0,
410 }),
411 Truth::Fact(b) => Ok(Condition {
412 yes: if b { 1.0 } else { 0.0 },
413 no: if b { 0.0 } else { 1.0 },
414 missing: 0.0,
415 }),
416 Truth::Probability(p) => Ok(Condition {
417 yes: p,
418 no: 1.0 - p,
419 missing: 0.0,
420 }),
421 Truth::Uncertain(d) => {
422 let (yes, no) = d.truth().expect("checked by truth");
423 Ok(Condition {
424 yes,
425 no,
426 missing: d.missing,
427 })
428 }
429 }
430}
431
432pub fn unary(op: UnOp, v: &Value, budget: &mut Budget) -> OpResult<Value> {
435 match op {
436 UnOp::Neg => lift1(v, budget, |x, budget| match x {
437 Value::Int(i) => {
438 budget.integer_work(i, &Integer::ZERO, false)?;
439 {
440 let n = i.negated();
441 budget.integer_allocation(n.bits(), 1)?;
442 Ok(Value::Int(n))
443 }
444 }
445 Value::Float(f) | Value::Prob(f) => Ok(Value::Float(-f)),
446 Value::Complex(z) => Ok(Value::Complex(z.negated())),
447 Value::Analytic(a) => crate::analytic::negate(a),
448 other => Err(OpError::new(format!("can't negate {}", article(&other.kind())))),
449 }),
450 UnOp::Not => not(v, budget),
451 UnOp::Typeof => unreachable!("typeof is lowered to an interpreter intrinsic"),
452 }
453}
454
455pub fn binary(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
456 match op {
457 BinOp::And | BinOp::Or => unreachable!("`and` and `or` are evaluated lazily by the interpreter"),
458 BinOp::To => unreachable!("`to` is lowered to a built-in"),
459 BinOp::Range | BinOp::RangeExcl => range(op, a, b, budget),
460 BinOp::In => lift2(a, b, budget, |x, coll, budget| {
461 contains(coll, x, budget).map(Value::Bool)
462 }),
463 BinOp::NotIn => lift2(a, b, budget, |x, coll, budget| {
464 contains(coll, x, budget).map(|c| Value::Bool(!c))
465 }),
466 _ => lift2(a, b, budget, |x, y, budget| binary_plain(op, x, y, budget)),
467 }
468}
469
470fn binary_plain(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
471 let unusual = |v: &Value| matches!(v, Value::Analytic(_) | Value::Continuous(_) | Value::Joint(_));
474 if unusual(a) || unusual(b) {
475 return unusual_binary(op, a, b, budget);
476 }
477 if is_comparison(op) && (crate::analytic::contains(a) || crate::analytic::contains(b)) {
478 return analytic_comparison(op, a, b, budget);
479 }
480 if matches!(
481 op,
482 BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge
483 ) {
484 for v in [a, b] {
485 if let Value::Int(n) = v {
486 budget.integer_work(n, &Integer::ONE, false)?;
487 }
488 if let Value::Str(s) = v {
489 budget.string_work(s)?;
490 }
491 }
492 }
493 match op {
494 BinOp::Eq => Ok(Value::Bool(equals(a, b))),
495 BinOp::Ne => Ok(Value::Bool(!equals(a, b))),
496 BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge => {
497 let ord = compare(a, b)?;
498 Ok(Value::Bool(match op {
499 BinOp::Lt => ord.is_lt(),
500 BinOp::Le => ord.is_le(),
501 BinOp::Gt => ord.is_gt(),
502 _ => ord.is_ge(),
503 }))
504 }
505 BinOp::Add => add(a, b, budget),
506 BinOp::Sub => sub(a, b, budget),
507 _ => arith(op, a, b, budget),
508 }
509}
510
511fn add(a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
512 match (a, b) {
513 (Value::Str(x), Value::Str(y)) => {
514 let size = x
515 .len()
516 .checked_add(y.len())
517 .ok_or_else(|| OpError::limit("string size overflow"))?;
518 budget.string_size(size)?;
519 let mut out = String::new();
520 crate::text::push(&mut out, x, budget)?;
521 crate::text::push(&mut out, y, budget)?;
522 Ok(Value::str(&out))
523 }
524 (Value::List(x), Value::List(y)) => {
525 budget.collection(x.len() as u128 + y.len() as u128)?;
526 budget.work(x.len() as u64 + y.len() as u64)?;
527 let mut items = x.to_vec();
528 items.extend(y.iter().cloned());
529 Ok(Value::list(items))
530 }
531 (Value::Date(d), n @ (Value::Int(_) | Value::Float(_)))
532 | (n @ (Value::Int(_) | Value::Float(_)), Value::Date(d)) => {
533 let n = integer(n, "date offset", budget)?;
534 date_plus(
535 *d,
536 n.to_i64()
537 .ok_or_else(|| OpError::fault(Fault::NumericOverflow, "date out of range"))?,
538 )
539 }
540 (Value::Str(_), _) | (_, Value::Str(_)) => {
541 Err(
542 OpError::new(format!("can't add {} and {}", article(&a.kind()), article(&b.kind())))
543 .help("to build text, use interpolation: \"total: {x}\""),
544 )
545 }
546 _ => arith(BinOp::Add, a, b, budget),
547 }
548}
549
550fn sub(a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
551 match (a, b) {
552 (Value::Date(x), Value::Date(y)) => Ok(Value::Int((*x as i64 - *y as i64).into())),
553 (Value::Date(d), n @ (Value::Int(_) | Value::Float(_))) => {
554 match integer(n, "date offset", budget)?.negated().to_i64() {
555 Some(m) => date_plus(*d, m),
556 None => Err(OpError::fault(Fault::NumericOverflow, "date out of range")),
557 }
558 }
559 _ => arith(BinOp::Sub, a, b, budget),
560 }
561}
562
563fn date_plus(d: i32, n: i64) -> OpResult<Value> {
564 crate::dates::add_days(d, n)
565 .map(Value::Date)
566 .ok_or_else(|| OpError::fault(Fault::NumericOverflow, "date out of range"))
567}
568
569fn number(v: &Value) -> Option<f64> {
571 match v {
572 Value::Bool(_) => None,
573 _ => v.as_f64(),
574 }
575}
576
577fn arith(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
578 let bad = || {
579 OpError::new(format!(
580 "can't use `{}` with {} and {}",
581 op.symbol(),
582 article(&a.kind()),
583 article(&b.kind())
584 ))
585 };
586 if let (BinOp::Pow, Value::Complex(z), Value::Int(n)) = (op, a, b) {
587 budget.integer_work(n, &Integer::ONE, false)?;
588 budget.work(n.bits())?;
589 return Ok(Value::Complex(z.pow_integer(n)?));
590 }
591 if matches!(a, Value::Complex(_)) || matches!(b, Value::Complex(_)) {
592 let (Some(x), Some(y)) = (a.as_complex(), b.as_complex()) else {
593 return Err(bad());
594 };
595 let z = match op {
596 BinOp::Add => x.plus(y),
597 BinOp::Sub => x.minus(y),
598 BinOp::Mul => x.times(y),
599 BinOp::Div => x.divided_by(y),
600 BinOp::Pow => match b {
601 Value::Int(n) => {
602 budget.integer_work(n, &Integer::ONE, false)?;
603 budget.work(n.bits())?;
604 x.pow_integer(n)
605 }
606 _ => Err(OpError::new("a complex power needs an int exponent")),
607 },
608 _ => Err(bad()),
609 }?;
610 return Ok(Value::Complex(z));
611 }
612 if let (Value::Int(x), Value::Int(y)) = (a, b) {
613 budget.integer_work(x, y, matches!(op, BinOp::Mul | BinOp::Div | BinOp::IntDiv | BinOp::Mod))?;
614 let n = match op {
615 BinOp::Add => x.add(y)?,
616 BinOp::Sub => x.sub(y)?,
617 BinOp::Mul => {
618 if !x.is_zero() && !y.is_zero() {
619 budget.integer_bits((x.bits() + y.bits()).saturating_sub(1))?;
620 }
621 x.mul(y)?
622 }
623 BinOp::Div => {
624 if y.is_zero() {
625 return Err(division_by_zero());
626 }
627 return x.ratio(y).map(Value::Float).ok_or_else(|| {
628 OpError::fault(
629 Fault::NumericOverflow,
630 "division result is too large for a finite float",
631 )
632 });
633 }
634 BinOp::IntDiv => x.div_mod(y)?.0,
635 BinOp::Mod => x.div_mod(y)?.1,
636 BinOp::Pow => return int_power(x, y, budget),
637 _ => return Err(bad()),
638 };
639 budget.integer_allocation(n.bits(), 1)?;
640 return Ok(Value::Int(n));
641 }
642 let (Some(x), Some(y)) = (number(a), number(b)) else {
643 return Err(bad());
644 };
645 let v = match op {
646 BinOp::Add => x + y,
647 BinOp::Sub => x - y,
648 BinOp::Mul => x * y,
649 BinOp::Div => {
650 if y == 0.0 {
651 return Err(division_by_zero());
652 }
653 x / y
654 }
655 BinOp::IntDiv => {
656 if y == 0.0 {
657 return Err(division_by_zero());
658 }
659 let q = (x / y).floor();
660 let n = Integer::from_f64(q)
661 .ok_or_else(|| OpError::fault(Fault::NumericOverflow, "integer division result is not finite"))?;
662 budget.integer_allocation(n.bits(), 1)?;
663 return Ok(Value::Int(n));
664 }
665 BinOp::Mod => {
666 if y == 0.0 {
667 return Err(division_by_zero());
668 }
669 x - (x / y).floor() * y
670 }
671 BinOp::Pow => libm::pow(x, y),
672 _ => return Err(bad()),
673 };
674 if !v.is_finite() {
675 return Err(OpError::fault(
676 Fault::NumericOverflow,
677 format!("`{}` gave a result that isn't a finite number", op.symbol()),
678 ));
679 }
680 Ok(Value::Float(v))
681}
682
683fn is_comparison(op: BinOp) -> bool {
687 matches!(
688 op,
689 BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge
690 )
691}
692
693#[inline(never)]
696fn unusual_binary(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
697 if matches!(a, Value::Analytic(_)) || matches!(b, Value::Analytic(_)) {
698 return crate::analytic::binary(op, a, b, budget);
699 }
700 if is_comparison(op) && (crate::analytic::contains(a) || crate::analytic::contains(b)) {
701 return analytic_comparison(op, a, b, budget);
702 }
703 match continuous_binary(op, a, b, budget)? {
704 Some(v) => Ok(v),
705 None => unreachable!("a continuous operand"),
706 }
707}
708
709#[inline(never)]
712fn analytic_comparison(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
713 if matches!(op, BinOp::Eq | BinOp::Ne) {
714 if let (Value::Event(x), Value::Bool(b)) | (Value::Bool(b), Value::Event(x)) = (a, b) {
715 return if *b == (op == BinOp::Eq) {
716 Ok(Value::Event(x.clone()))
717 } else {
718 not(&Value::Event(x.clone()), budget)
719 };
720 }
721 }
722 if let (Value::Event(x), Value::Event(y)) = (a, b) {
723 if matches!(op, BinOp::Eq | BinOp::Ne) {
724 budget.work((x.size() + y.size()) as u64)?;
725 let equal = x.combine(y, BinOp::Eq)?;
726 return Ok(if op == BinOp::Eq { equal } else { equal.complement()? }.value());
727 }
728 }
729 Err(crate::analytic::unsupported(
730 "comparing aggregate or boolean analytic outcomes",
731 ))
732}
733
734fn continuous_binary(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Option<Value>> {
735 let (family, other, flipped) = match (a, b) {
736 (Value::Continuous(f), other) => (f, other, false),
737 (other, Value::Continuous(f)) => (f, other, true),
738 (Value::Joint(_), _) | (_, Value::Joint(_)) => return crate::joint::binary(op, a, b, budget).map(Some),
739 _ => return Ok(None),
740 };
741 let needs_value = || {
742 OpError::new(format!(
743 "`{}` needs a value, not a {} distribution",
744 op.symbol(),
745 family.name()
746 ))
747 .help("draw a value first, like `let x ~ normal(0, 1)`; comparing a distribution with a number works too")
748 };
749 let x = match number(other) {
750 Some(x) if !x.is_nan() => x,
751 _ => return Err(needs_value()),
752 };
753 let below = family.cdf(x);
755 let yes = match (op, flipped) {
756 (BinOp::Lt | BinOp::Le, false) | (BinOp::Gt | BinOp::Ge, true) => below,
757 (BinOp::Gt | BinOp::Ge, false) | (BinOp::Lt | BinOp::Le, true) => 1.0 - below,
758 (BinOp::Eq, _) => 0.0,
759 (BinOp::Ne, _) => 1.0,
760 _ => {
762 return Err(OpError::unsupported(format!(
763 "`{}` on an undrawn {} distribution isn't supported yet",
764 op.symbol(),
765 family.name()
766 ))
767 .help("draw a value first, like `let x ~ normal(0, 1)`, and compute with it"));
768 }
769 };
770 Ok(Some(Dist::bernoulli(yes).into_value()))
771}
772
773fn division_by_zero() -> OpError {
774 OpError::fault(Fault::DivisionByZero, "division by zero")
775}
776
777fn int_power(base: &Integer, exponent: &Integer, budget: &mut Budget) -> OpResult<Value> {
778 if exponent.is_zero() {
779 return Ok(Value::Int(Integer::ONE));
780 }
781 if base.is_zero() {
782 return if exponent.is_negative() {
783 Err(division_by_zero())
784 } else {
785 Ok(Value::Int(Integer::ZERO))
786 };
787 }
788 if exponent.is_negative() {
789 let reciprocal = Integer::ONE.ratio(&base.abs()).unwrap_or(0.0);
790 let power = exponent.abs().to_f64().unwrap_or(f64::INFINITY);
791 let result = libm::pow(reciprocal, power);
792 return Ok(Value::Float(if base.is_negative() && exponent.is_odd() {
793 -result
794 } else {
795 result
796 }));
797 }
798 if *base == 1 || *base == -1 {
799 return Ok(Value::Int(if *base == -1 && exponent.is_odd() {
800 (-1).into()
801 } else {
802 Integer::ONE
803 }));
804 }
805 let mut n = exponent
806 .to_u64()
807 .ok_or_else(|| OpError::limit("integer power exceeds the integer size limit"))?;
808 budget.integer_bits(base.bits().saturating_sub(1).saturating_mul(n).saturating_add(1))?;
809 let mut result = Integer::ONE;
810 let mut b = base.clone();
811 while n != 0 {
812 if n & 1 != 0 {
813 budget.integer_work(&result, &b, true)?;
814 result = result.mul(&b)?;
815 budget.integer_bits(result.bits())?;
816 }
817 n >>= 1;
818 if n != 0 {
819 budget.integer_work(&b, &b, true)?;
820 b = b.mul(&b)?;
821 budget.integer_bits(b.bits())?;
822 }
823 }
824 budget.integer_allocation(result.bits(), 1)?;
825 Ok(Value::Int(result))
826}
827
828fn numeric_compare(a: &Value, b: &Value) -> Option<std::cmp::Ordering> {
829 match (a, b) {
830 (Value::Int(a), Value::Int(b)) => Some(a.cmp(b)),
831 (Value::Int(a), Value::Float(b) | Value::Prob(b)) => a.cmp_f64(*b),
832 (Value::Float(a) | Value::Prob(a), Value::Int(b)) => b.cmp_f64(*a).map(std::cmp::Ordering::reverse),
833 (Value::Float(a) | Value::Prob(a), Value::Float(b) | Value::Prob(b)) => a.partial_cmp(b),
834 _ => None,
835 }
836}
837
838pub fn equals(a: &Value, b: &Value) -> bool {
840 if let (Value::Complex(x), Value::Complex(y)) = (a, b) {
841 return x == y;
842 }
843 if let (Value::Complex(z), real) | (real, Value::Complex(z)) = (a, b) {
844 if z.im() != 0.0 {
845 return false;
846 }
847 return match real {
848 Value::Int(n) => {
850 let x = z.re();
851 n.cmp_f64(x).is_some_and(|c| c.is_eq())
852 }
853 Value::Float(x) | Value::Prob(x) => z.re() == *x,
854 _ => false,
855 };
856 }
857 if let Some(c) = numeric_compare(a, b) {
858 return c.is_eq();
859 }
860 if matches!(a, Value::Float(x) | Value::Prob(x) if x.is_nan())
861 || matches!(b, Value::Float(x) | Value::Prob(x) if x.is_nan())
862 {
863 return false;
864 }
865 a == b
866}
867
868pub fn compare(a: &Value, b: &Value) -> OpResult<std::cmp::Ordering> {
870 if matches!(a, Value::Complex(_)) || matches!(b, Value::Complex(_)) {
871 return Err(OpError::new("complex values have no ordering").help("compare `abs(z)`, `real(z)` or `imag(z)`"));
872 }
873 if let Some(c) = numeric_compare(a, b) {
874 return Ok(c);
875 }
876 if matches!(a, Value::Float(x) | Value::Prob(x) if x.is_nan())
877 || matches!(b, Value::Float(x) | Value::Prob(x) if x.is_nan())
878 {
879 return Err(OpError::new("can't order NaN"));
880 }
881 match (a, b) {
882 (Value::Str(x), Value::Str(y)) => Ok(x.cmp(y)),
883 (Value::Date(x), Value::Date(y)) => Ok(x.cmp(y)),
884 (Value::Enum(x), Value::Enum(y)) if x.ty == y.ty => Ok(x.variant.cmp(&y.variant)),
885 (Value::List(x), Value::List(y)) => {
886 for (p, q) in x.iter().zip(y.iter()) {
887 let c = compare(p, q)?;
888 if c.is_ne() {
889 return Ok(c);
890 }
891 }
892 Ok(x.len().cmp(&y.len()))
893 }
894 _ => Err(OpError::new(format!(
895 "can't compare {} with {}",
896 article(&a.kind()),
897 article(&b.kind())
898 ))),
899 }
900}
901
902pub fn contains(coll: &Value, item: &Value, budget: &mut Budget) -> OpResult<bool> {
904 if crate::analytic::contains(coll) || crate::analytic::contains(item) {
905 return Err(crate::analytic::unsupported("membership involving analytic outcomes"));
906 }
907 if let Value::Str(s) = coll {
908 budget.string_work(s)?;
909 if let Value::Str(s) = item {
910 budget.string_work(s)?;
911 }
912 }
913 Ok(match coll {
914 Value::List(items) => items.iter().any(|x| equals(x, item)),
915 Value::Map(m) => m.contains_key(item),
916 Value::Bag(b) => b.get(item).is_some_and(|n| *n > 0),
917 Value::Range(lo, hi) => match item {
918 Value::Int(n) => n >= lo && n <= hi,
919 Value::Float(x) | Value::Prob(x) if x.is_finite() && x.fract() == 0.0 => {
920 lo.cmp_f64(*x).is_some_and(|c| !c.is_gt()) && hi.cmp_f64(*x).is_some_and(|c| !c.is_lt())
921 }
922 _ => false,
923 },
924 Value::Str(s) => match item {
925 Value::Str(sub) => s.contains(&**sub),
926 _ => {
927 return Err(OpError::new(format!(
928 "can't look for {} in a string",
929 article(&item.kind())
930 )));
931 }
932 },
933 other => return Err(OpError::new(format!("can't look inside {}", article(&other.kind())))),
934 })
935}
936
937fn range(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
938 if a.is_uncertain() || b.is_uncertain() {
939 return Err(OpError::new("a range needs plain whole numbers, not distributions")
940 .help("draw a value first, like `let n ~ d6`"));
941 }
942 let lo = integer(a, "range bound", budget)?;
943 let hi = integer(b, "range bound", budget)?;
944 let hi = if op == BinOp::RangeExcl {
945 hi.sub(&Integer::ONE)?
946 } else {
947 hi.into_owned()
948 };
949 budget.integer_work(&lo, &hi, false)?;
950 if op == BinOp::RangeExcl {
951 budget.integer_allocation(hi.bits(), 1)?;
952 }
953 Ok(Value::Range(lo.into_owned(), hi))
954}
955
956pub fn range_len(lo: &Integer, hi: &Integer) -> OpResult<Integer> {
958 if hi < lo {
959 Ok(Integer::ZERO)
960 } else {
961 Ok(hi.sub(lo)?.add(&Integer::ONE)?)
962 }
963}
964
965pub fn range_count(lo: &Integer, hi: &Integer) -> OpResult<u128> {
967 range_len(lo, hi)?
968 .to_u128()
969 .ok_or_else(|| OpError::limit("range has too many elements to materialize"))
970}
971
972pub fn field(v: &Value, name: &str, budget: &mut Budget) -> OpResult<Value> {
975 lift1(v, budget, |x, _| match x {
976 Value::Record(r) => r.get(name).cloned().ok_or_else(|| {
977 let known: Vec<String> = r.fields.iter().map(|(n, _)| format!("`{n}`")).collect();
978 OpError::new(format!("{} has no field `{name}`", article(&x.kind())))
979 .help(format!("its fields are {}", known.join(", ")))
980 }),
981 other => Err(OpError::new(format!(
982 "can't read the field `{name}` of {}",
983 article(&other.kind())
984 ))),
985 })
986}
987
988pub fn index(coll: &Value, i: &Value, budget: &mut Budget) -> OpResult<Value> {
989 lift2(coll, i, budget, |c, i, budget| {
990 let v = index_plain(c, i, budget)?;
991 if let (Value::Range(lo, hi), Value::Int(n)) = (c, &v) {
992 budget.integer_work(lo, hi, false)?;
993 budget.integer_allocation(n.bits(), 1)?;
994 }
995 Ok(v)
996 })
997}
998
999pub fn index_plain(coll: &Value, i: &Value, budget: &mut Budget) -> OpResult<Value> {
1000 match coll {
1001 Value::List(items) => {
1002 let k = as_index(i, items.len() as u128, budget)?;
1003 Ok(items[k as usize].clone())
1004 }
1005 Value::Range(lo, hi) => {
1006 let k = integer(i, "index", budget)?;
1007 if k.is_negative() {
1008 return Err(OpError::fault(Fault::IndexOutOfBounds, "index out of range"));
1009 }
1010 let n = lo.add(&k)?;
1011 if &n > hi {
1012 return Err(OpError::fault(Fault::IndexOutOfBounds, "index out of range"));
1013 }
1014 Ok(Value::Int(n))
1015 }
1016 Value::Str(s) => {
1017 budget.string_work(s)?;
1018 let k = as_index(i, s.chars().count() as u128, budget)?;
1019 let c = s.chars().nth(k as usize).expect("checked scalar index");
1020 crate::text::value(c.encode_utf8(&mut [0; 4]), budget)
1021 }
1022 Value::Map(m) => m.get(i).cloned().ok_or_else(|| {
1023 OpError::fault(Fault::MissingKey, format!("the key {i:?} isn't in the map"))
1024 .help("use `get(key, default)` for keys that may be missing")
1025 }),
1026 other => Err(OpError::new(format!("can't index {}", article(&other.kind())))),
1027 }
1028}
1029
1030pub fn as_index(i: &Value, len: u128, budget: &mut Budget) -> OpResult<u128> {
1031 let k = integer(i, "index", budget)?.to_u128();
1032 k.filter(|k| *k < len).ok_or_else(|| {
1033 OpError::fault(
1034 Fault::IndexOutOfBounds,
1035 format!("index {i} is out of range for a length of {len}"),
1036 )
1037 .help("indices start at 0")
1038 })
1039}
1040
1041pub fn make_record(ty: Option<Arc<str>>, mut fields: Vec<(Arc<str>, Value)>) -> Value {
1042 fields.sort_by(|a, b| a.0.cmp(&b.0));
1043 Value::record(Record { ty, fields })
1044}
1045
1046pub fn with_fields(base: &Value, updates: &[(Arc<str>, Value)]) -> OpResult<Value> {
1047 let Value::Record(r) = base else {
1048 return Err(OpError::new(format!(
1049 "`with` needs a record, found {}",
1050 article(&base.kind())
1051 )));
1052 };
1053 let mut r = Record::clone(r);
1054 for (name, v) in updates {
1055 match r.get_mut(name) {
1056 Some(slot) => *slot = v.clone(),
1057 None => {
1058 return Err(OpError::new(format!("{} has no field `{name}`", article(&base.kind()))));
1059 }
1060 }
1061 }
1062 Ok(Value::record(r))
1063}
1064
1065pub fn enum_value(ty: u32, variant: u32, name: &str) -> Value {
1066 Value::Enum(Arc::new(EnumValue {
1067 ty,
1068 variant,
1069 name: Arc::from(name),
1070 }))
1071}
1072
1073pub fn is_certain(v: &Value, truth: bool) -> bool {
1075 matches!(v, Value::Bool(b) if *b == truth)
1076}
1077
1078#[cfg(test)]
1079mod probability_tests {
1080 use super::*;
1081
1082 #[test]
1083 fn computed_probabilities_only_correct_boundary_roundoff() {
1084 assert_eq!(computed_prob(1.0 + f64::EPSILON, "test").unwrap(), Value::Prob(1.0));
1085 assert_eq!(computed_prob(-f64::EPSILON, "test").unwrap(), Value::Prob(0.0));
1086 assert_eq!(computed_prob(1e-300, "test").unwrap(), Value::Prob(1e-300));
1087 for p in [1.001, -0.001, f64::INFINITY, f64::NEG_INFINITY, f64::NAN] {
1088 assert!(computed_prob(p, "test").is_err());
1089 }
1090 }
1091}