1use super::ir::{Assignment, ConstraintEvalError, Expr, IrError, Ty, VarId, eval};
45use super::solve::Weight;
46use core::fmt;
47use num_bigint::BigInt;
48use num_bigint::BigUint;
49use num_traits::{One, Zero};
50
51pub const SELECTOR_BASE: VarId = 1 << 30;
55
56#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
59pub enum Span {
60 Point(Expr),
62 Range {
64 lo: Option<Expr>,
66 hi: Option<Expr>,
68 },
69}
70
71#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
73pub struct Item {
74 pub member: Expr,
77 pub ty: Ty,
80 pub span: Span,
82 pub weight: Expr,
84 pub per_range: bool,
86}
87
88#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
90pub struct DefaultWeight {
91 pub weight: Expr,
93 pub per_range: bool,
95}
96
97#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
99pub struct Dist {
100 pub subject_ty: Ty,
102 pub plain_var: Option<VarId>,
105 pub vars: Vec<VarId>,
107 pub items: Vec<Item>,
109 pub default: Option<DefaultWeight>,
111 pub subject: Option<Expr>,
120}
121
122#[derive(Clone, Debug, PartialEq, Eq)]
124pub enum DistError {
125 Eval(ConstraintEvalError),
128 NegativeWeight,
131}
132
133impl fmt::Display for DistError {
134 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
135 match self {
136 DistError::Eval(e) => write!(f, "in a `dist` weight or bound: {e}"),
137 DistError::NegativeWeight => f.write_str("a `dist` weight is negative"),
138 }
139 }
140}
141
142impl std::error::Error for DistError {}
143
144impl From<ConstraintEvalError> for DistError {
145 fn from(e: ConstraintEvalError) -> Self {
146 DistError::Eval(e)
147 }
148}
149
150fn numeric(v: &BigUint, t: Ty) -> BigInt {
151 if t.signed && v.bit(u64::from(t.width) - 1) {
152 BigInt::from(v.clone()) - (BigInt::one() << t.width)
153 } else {
154 BigInt::from(v.clone())
155 }
156}
157
158fn subject_range_in(t: Ty, s: Ty) -> (BigInt, BigInt) {
163 if t.signed && s.signed {
164 let half = BigInt::one() << (s.width - 1);
165 (-half.clone(), half - BigInt::one())
166 } else {
167 (BigInt::zero(), (BigInt::one() << s.width) - BigInt::one())
168 }
169}
170
171fn mixed_bound_cover(
177 s: Ty,
178 (lo, lo_ty): (&BigInt, Ty),
179 (hi, hi_ty): (&BigInt, Ty),
180) -> Vec<(BigInt, BigInt)> {
181 let w = s.width;
182 let (nmin, nmax) = if s.signed {
183 let half = BigInt::one() << (w - 1);
184 (-half.clone(), half - BigInt::one())
185 } else {
186 (BigInt::zero(), (BigInt::one() << w) - BigInt::one())
187 };
188 let side = |v: &BigInt, t: Ty, is_lo: bool| -> Vec<(BigInt, BigInt)> {
190 if t.signed || !s.signed {
191 let (a, b) = if is_lo {
194 (v.clone().max(nmin.clone()), nmax.clone())
195 } else {
196 (nmin.clone(), v.clone().min(nmax.clone()))
197 };
198 return if a <= b { vec![(a, b)] } else { Vec::new() };
199 }
200 let pmax = (BigInt::one() << w) - BigInt::one();
203 let (a, b) = if is_lo {
204 (v.clone().max(BigInt::zero()), pmax.clone())
205 } else {
206 (BigInt::zero(), v.clone().min(pmax.clone()))
207 };
208 if a > b {
209 return Vec::new();
210 }
211 let half = BigInt::one() << (w - 1);
212 let full = BigInt::one() << w;
213 let mut out = Vec::new();
214 if a < half {
215 out.push((a.clone(), (&half - BigInt::one()).min(b.clone())));
216 }
217 if b >= half {
218 out.push((a.max(half) - &full, b - &full));
219 }
220 out
221 };
222 let (ls, hs) = (side(lo, lo_ty, true), side(hi, hi_ty, false));
223 let mut out = Vec::new();
224 for (a, b) in &ls {
225 for (c, d) in &hs {
226 let (x, y) = (a.max(c).clone(), b.min(d).clone());
227 if x <= y {
228 out.push((x, y));
229 }
230 }
231 }
232 out
233}
234
235fn union_len(mut iv: Vec<(BigInt, BigInt)>) -> BigInt {
237 iv.sort();
238 let mut total = BigInt::zero();
239 let mut cur: Option<(BigInt, BigInt)> = None;
240 for (a, b) in iv {
241 cur = match cur {
242 Some((ca, cb)) if a <= &cb + BigInt::one() => Some((ca, if b > cb { b } else { cb })),
243 Some((ca, cb)) => {
244 total += cb - ca + BigInt::one();
245 Some((a, b))
246 }
247 None => Some((a, b)),
248 };
249 }
250 if let Some((a, b)) = cur {
251 total += b - a + BigInt::one();
252 }
253 total
254}
255
256impl Dist {
257 pub fn validate(&self) -> Result<(), IrError> {
260 let state_only = |e: &Expr| -> Result<(), IrError> {
261 e.validate()?;
262 let mut has_var = false;
263 e.for_each_var(&mut |_, _| has_var = true);
264 if has_var {
265 return Err(IrError::IllTyped {
266 what: "a `dist` weight or bound reads a random variable".into(),
267 });
268 }
269 Ok(())
270 };
271 for it in &self.items {
272 it.member.validate()?;
273 state_only(&it.weight)?;
274 match &it.span {
275 Span::Point(v) => state_only(v)?,
276 Span::Range { lo, hi } => {
277 for b in [lo, hi].into_iter().flatten() {
278 state_only(b)?;
279 }
280 }
281 }
282 }
283 if let Some(d) = &self.default {
284 state_only(&d.weight)?;
285 }
286 Ok(())
287 }
288
289 pub fn codes(&self) -> usize {
291 self.items.len() + usize::from(self.default.is_some())
292 }
293
294 pub fn selector_ty(&self) -> Ty {
296 let top = self.codes().saturating_sub(1) as u64;
297 Ty::unsigned((64 - top.leading_zeros()).max(1))
298 }
299
300 pub fn conjunct(&self, sel: VarId) -> Expr {
303 let kt = self.selector_ty();
304 let bin = |op, a: Expr, b: Expr| Expr::Binary(op, Box::new(a), Box::new(b), Ty::BIT);
305 let code = |i: usize| {
306 bin(
307 super::ir::BinOp::Eq,
308 Expr::Var(sel, kt),
309 Expr::constant(kt, i as u64),
310 )
311 };
312 let mut terms: Vec<Expr> = self
313 .items
314 .iter()
315 .enumerate()
316 .map(|(i, it)| bin(super::ir::BinOp::LogAnd, code(i), it.member.clone()))
317 .collect();
318 if self.default.is_some() {
319 let any = self
320 .items
321 .iter()
322 .map(|it| it.member.clone())
323 .reduce(|a, b| bin(super::ir::BinOp::LogOr, a, b));
324 let none = match any {
325 Some(a) => Expr::Unary(super::ir::UnOp::LogNot, Box::new(a), Ty::BIT),
326 None => Expr::Const(Ty::BIT, BigUint::one()),
327 };
328 terms.push(bin(super::ir::BinOp::LogAnd, code(self.items.len()), none));
329 }
330 terms
331 .into_iter()
332 .reduce(|a, b| bin(super::ir::BinOp::LogOr, a, b))
333 .unwrap_or(Expr::Const(Ty::BIT, BigUint::zero()))
334 }
335
336 pub fn weights(&self, states: &[Option<BigUint>]) -> Result<Vec<Weight>, DistError> {
341 let a = Assignment {
342 vars: Vec::new(),
343 states: states.to_vec(),
344 };
345 let weight_of = |e: &Expr| -> Result<BigUint, DistError> {
346 let v = numeric(&eval(e, &a)?, e.ty());
347 v.to_biguint().ok_or(DistError::NegativeWeight)
348 };
349 let per = |w: BigUint, per_range: bool, n: &BigInt| -> Weight {
350 if !per_range {
351 return Weight::int(w);
352 }
353 match n.to_biguint() {
354 Some(n) if !n.is_zero() => Weight { num: w, den: n },
355 _ => Weight::int(0u32),
356 }
357 };
358 let mut out = Vec::with_capacity(self.codes());
359 let mut covered: Vec<(BigInt, BigInt)> = Vec::new();
360 for it in &self.items {
361 let (smin, smax) = subject_range_in(it.ty, self.subject_ty);
362 let bound = |e: &Option<Expr>, dflt: &BigInt| -> Result<BigInt, DistError> {
363 match e {
364 Some(e) => Ok(numeric(&eval(e, &a)?, e.ty())),
365 None => Ok(dflt.clone()),
366 }
367 };
368 let (lo, hi) = match &it.span {
369 Span::Point(v) => {
370 let n = numeric(&eval(v, &a)?, v.ty());
371 (n.clone(), n)
372 }
373 Span::Range { lo, hi } => (bound(lo, &smin)?, bound(hi, &smax)?),
374 };
375 let n = if hi >= lo {
376 &hi - &lo + 1
377 } else {
378 BigInt::zero()
379 };
380 out.push(per(weight_of(&it.weight)?, it.per_range, &n));
381 if let Span::Range {
384 lo: Some(l),
385 hi: Some(h),
386 } = &it.span
387 && l.ty() != h.ty()
388 {
389 covered.extend(mixed_bound_cover(
390 self.subject_ty,
391 (&lo, l.ty()),
392 (&hi, h.ty()),
393 ));
394 continue;
395 }
396 let (a_, b_) = (lo.max(smin), hi.min(smax));
397 if a_ > b_ {
398 continue;
399 }
400 if self.subject_ty.signed && !it.ty.signed {
401 let half = BigInt::one() << (self.subject_ty.width - 1);
404 let full = BigInt::one() << self.subject_ty.width;
405 if a_ < half {
406 covered.push((a_.clone(), (&half - BigInt::one()).min(b_.clone())));
407 }
408 if b_ >= half {
409 covered.push((a_.max(half) - &full, b_ - &full));
410 }
411 } else {
412 covered.push((a_, b_));
413 }
414 }
415 if let Some(d) = &self.default {
416 let width = self.subject_ty.width;
417 let n = (BigInt::one() << width) - union_len(covered);
418 out.push(per(weight_of(&d.weight)?, d.per_range, &n));
419 }
420 Ok(out)
421 }
422}
423
424#[cfg(test)]
425mod tests {
426 use super::*;
427 use crate::constraints::ir::BinOp;
428
429 fn u(w: u32) -> Ty {
430 Ty::unsigned(w)
431 }
432 fn k(t: Ty, v: u64) -> Expr {
433 Expr::constant(t, v)
434 }
435 fn bin(op: BinOp, a: Expr, b: Expr) -> Expr {
436 Expr::Binary(op, Box::new(a), Box::new(b), Ty::BIT)
437 }
438 fn item(x: &Expr, lo: u64, hi: u64, w: u64, per_range: bool) -> Item {
439 let t = x.ty();
440 Item {
441 member: bin(
442 BinOp::LogAnd,
443 bin(BinOp::Ge, x.clone(), k(t, lo)),
444 bin(BinOp::Le, x.clone(), k(t, hi)),
445 ),
446 ty: t,
447 span: Span::Range {
448 lo: Some(k(t, lo)),
449 hi: Some(k(t, hi)),
450 },
451 weight: k(u(32), w),
452 per_range,
453 }
454 }
455 fn dist(subject_ty: Ty, items: Vec<Item>, default: Option<DefaultWeight>) -> Dist {
456 Dist {
457 subject_ty,
458 plain_var: Some(0),
459 vars: vec![0],
460 items,
461 default,
462 subject: None,
463 }
464 }
465 fn w(num: u32, den: u32) -> Weight {
466 Weight {
467 num: num.into(),
468 den: den.into(),
469 }
470 }
471
472 #[test]
473 fn the_selector_has_one_code_per_item_and_the_default() {
474 let x = Expr::Var(0, u(8));
475 let d = dist(u(8), vec![item(&x, 1, 1, 1, false)], None);
476 assert_eq!((d.codes(), d.selector_ty()), (1, u(1)));
477 let d = dist(
478 u(8),
479 vec![item(&x, 1, 1, 1, false), item(&x, 2, 2, 1, false)],
480 None,
481 );
482 assert_eq!((d.codes(), d.selector_ty()), (2, u(1)));
483 let d = dist(
484 u(8),
485 vec![item(&x, 1, 1, 1, false), item(&x, 2, 2, 1, false)],
486 Some(DefaultWeight {
487 weight: k(u(32), 1),
488 per_range: true,
489 }),
490 );
491 assert_eq!((d.codes(), d.selector_ty()), (3, u(2)));
492 }
493
494 #[test]
495 fn per_value_and_per_range_weights() {
496 let x = Expr::Var(0, u(8));
498 let d = dist(
499 u(8),
500 vec![item(&x, 0, 3, 5, false), item(&x, 4, 11, 8, true)],
501 None,
502 );
503 assert_eq!(d.weights(&[]).unwrap(), vec![w(5, 1), w(8, 8)]);
504 }
505
506 #[test]
507 fn the_default_weight_spreads_over_the_values_no_item_names() {
508 let x = Expr::Var(0, u(4));
511 let d = dist(
512 u(4),
513 vec![item(&x, 0, 3, 1, false), item(&x, 2, 5, 1, false)],
514 Some(DefaultWeight {
515 weight: k(u(32), 20),
516 per_range: true,
517 }),
518 );
519 assert_eq!(d.weights(&[]).unwrap()[2], w(20, 10));
520 }
521
522 #[test]
523 fn an_item_beyond_the_subjects_range_does_not_shrink_the_default() {
524 let wide = Item {
528 member: Expr::Const(Ty::BIT, BigUint::one()),
529 ty: u(32),
530 span: Span::Range {
531 lo: Some(k(u(32), 10)),
532 hi: Some(k(u(32), 40)),
533 },
534 weight: k(u(32), 31),
535 per_range: true,
536 };
537 let d = dist(
538 u(4),
539 vec![wide],
540 Some(DefaultWeight {
541 weight: k(u(32), 10),
542 per_range: true,
543 }),
544 );
545 assert_eq!(d.weights(&[]).unwrap(), vec![w(31, 31), w(10, 10)]);
546 }
547
548 #[test]
549 fn an_empty_range_and_a_full_default_have_weight_zero_not_a_zero_divisor() {
550 let x = Expr::Var(0, u(2));
551 let d = dist(
554 u(2),
555 vec![item(&x, 3, 1, 4, true), item(&x, 0, 3, 1, false)],
556 Some(DefaultWeight {
557 weight: k(u(32), 9),
558 per_range: true,
559 }),
560 );
561 assert_eq!(d.weights(&[]).unwrap(), vec![w(0, 1), w(1, 1), w(0, 1)]);
562 }
563
564 #[test]
565 fn weights_and_bounds_read_state_variables() {
566 let st = |i: u32| Expr::State(i, u(8));
568 let d = dist(
569 u(8),
570 vec![Item {
571 member: Expr::Const(Ty::BIT, BigUint::one()),
572 ty: u(8),
573 span: Span::Range {
574 lo: Some(st(0)),
575 hi: Some(st(1)),
576 },
577 weight: st(2),
578 per_range: true,
579 }],
580 None,
581 );
582 let states = [Some(0u32.into()), Some(9u32.into()), Some(20u32.into())];
583 assert_eq!(d.weights(&states).unwrap(), vec![w(20, 10)]);
584 let unknown = [Some(0u32.into()), None, Some(20u32.into())];
586 assert_eq!(
587 d.weights(&unknown),
588 Err(DistError::Eval(ConstraintEvalError::StateUnknown(1)))
589 );
590 }
591
592 #[test]
593 fn a_negative_weight_is_an_error() {
594 let x = Expr::Var(0, u(8));
595 let mut it = item(&x, 1, 1, 1, false);
596 it.weight = Expr::constant(Ty::signed(8), 0xFFu32); let d = dist(u(8), vec![it], None);
598 assert_eq!(d.weights(&[]), Err(DistError::NegativeWeight));
599 }
600
601 #[test]
602 fn a_signed_subject_counts_negative_values_in_the_default() {
603 let t = Ty::signed(4);
606 let neg2 = Expr::constant(t, 0xEu32);
607 let one = Expr::constant(t, 1u32);
608 let d = dist(
609 t,
610 vec![Item {
611 member: Expr::Const(Ty::BIT, BigUint::one()),
612 ty: t,
613 span: Span::Range {
614 lo: Some(neg2),
615 hi: Some(one),
616 },
617 weight: k(u(32), 1),
618 per_range: false,
619 }],
620 Some(DefaultWeight {
621 weight: k(u(32), 12),
622 per_range: true,
623 }),
624 );
625 assert_eq!(d.weights(&[]).unwrap(), vec![w(1, 1), w(12, 12)]);
626 }
627
628 #[test]
629 fn range_bounds_of_different_types_cover_the_intersection() {
630 let t = Ty::signed(4);
635 let d = dist(
636 t,
637 vec![Item {
638 member: Expr::Const(Ty::BIT, BigUint::one()),
639 ty: t,
640 span: Span::Range {
641 lo: Some(Expr::constant(t, 0xEu32)),
642 hi: Some(Expr::constant(u(4), 12u32)),
643 },
644 weight: k(u(32), 1),
645 per_range: false,
646 }],
647 Some(DefaultWeight {
648 weight: k(u(32), 8),
649 per_range: true,
650 }),
651 );
652 assert_eq!(d.weights(&[]).unwrap(), vec![w(1, 1), w(8, 8)]);
653 }
654
655 #[test]
656 fn the_conjunct_is_membership_of_exactly_one_item() {
657 let x = Expr::Var(0, u(2));
661 let d = dist(
662 u(2),
663 vec![item(&x, 1, 1, 1, false), item(&x, 2, 2, 1, false)],
664 Some(DefaultWeight {
665 weight: k(u(32), 1),
666 per_range: true,
667 }),
668 );
669 let c = d.conjunct(1);
670 let kt = d.selector_ty();
671 let ok = |xv: u64, kv: u64| {
672 let a = Assignment {
673 vars: vec![xv.into(), kv.into()],
674 states: vec![],
675 };
676 eval(&c, &a).unwrap() == BigUint::one()
677 };
678 for xv in 0..4 {
679 for kv in 0..(1u64 << kt.width) {
680 let want = match kv {
681 0 => xv == 1,
682 1 => xv == 2,
683 2 => xv == 0 || xv == 3,
684 _ => false,
685 };
686 assert_eq!(ok(xv, kv), want, "x={xv} K={kv}");
687 }
688 }
689 }
690}