use super::ir::{Assignment, ConstraintEvalError, Expr, IrError, Ty, VarId, eval};
use super::solve::Weight;
use core::fmt;
use num_bigint::BigInt;
use num_bigint::BigUint;
use num_traits::{One, Zero};
pub const SELECTOR_BASE: VarId = 1 << 30;
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum Span {
Point(Expr),
Range {
lo: Option<Expr>,
hi: Option<Expr>,
},
}
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Item {
pub member: Expr,
pub ty: Ty,
pub span: Span,
pub weight: Expr,
pub per_range: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DefaultWeight {
pub weight: Expr,
pub per_range: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Dist {
pub subject_ty: Ty,
pub plain_var: Option<VarId>,
pub vars: Vec<VarId>,
pub items: Vec<Item>,
pub default: Option<DefaultWeight>,
pub subject: Option<Expr>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum DistError {
Eval(ConstraintEvalError),
NegativeWeight,
}
impl fmt::Display for DistError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
DistError::Eval(e) => write!(f, "in a `dist` weight or bound: {e}"),
DistError::NegativeWeight => f.write_str("a `dist` weight is negative"),
}
}
}
impl std::error::Error for DistError {}
impl From<ConstraintEvalError> for DistError {
fn from(e: ConstraintEvalError) -> Self {
DistError::Eval(e)
}
}
fn numeric(v: &BigUint, t: Ty) -> BigInt {
if t.signed && v.bit(u64::from(t.width) - 1) {
BigInt::from(v.clone()) - (BigInt::one() << t.width)
} else {
BigInt::from(v.clone())
}
}
fn subject_range_in(t: Ty, s: Ty) -> (BigInt, BigInt) {
if t.signed && s.signed {
let half = BigInt::one() << (s.width - 1);
(-half.clone(), half - BigInt::one())
} else {
(BigInt::zero(), (BigInt::one() << s.width) - BigInt::one())
}
}
fn mixed_bound_cover(
s: Ty,
(lo, lo_ty): (&BigInt, Ty),
(hi, hi_ty): (&BigInt, Ty),
) -> Vec<(BigInt, BigInt)> {
let w = s.width;
let (nmin, nmax) = if s.signed {
let half = BigInt::one() << (w - 1);
(-half.clone(), half - BigInt::one())
} else {
(BigInt::zero(), (BigInt::one() << w) - BigInt::one())
};
let side = |v: &BigInt, t: Ty, is_lo: bool| -> Vec<(BigInt, BigInt)> {
if t.signed || !s.signed {
let (a, b) = if is_lo {
(v.clone().max(nmin.clone()), nmax.clone())
} else {
(nmin.clone(), v.clone().min(nmax.clone()))
};
return if a <= b { vec![(a, b)] } else { Vec::new() };
}
let pmax = (BigInt::one() << w) - BigInt::one();
let (a, b) = if is_lo {
(v.clone().max(BigInt::zero()), pmax.clone())
} else {
(BigInt::zero(), v.clone().min(pmax.clone()))
};
if a > b {
return Vec::new();
}
let half = BigInt::one() << (w - 1);
let full = BigInt::one() << w;
let mut out = Vec::new();
if a < half {
out.push((a.clone(), (&half - BigInt::one()).min(b.clone())));
}
if b >= half {
out.push((a.max(half) - &full, b - &full));
}
out
};
let (ls, hs) = (side(lo, lo_ty, true), side(hi, hi_ty, false));
let mut out = Vec::new();
for (a, b) in &ls {
for (c, d) in &hs {
let (x, y) = (a.max(c).clone(), b.min(d).clone());
if x <= y {
out.push((x, y));
}
}
}
out
}
fn union_len(mut iv: Vec<(BigInt, BigInt)>) -> BigInt {
iv.sort();
let mut total = BigInt::zero();
let mut cur: Option<(BigInt, BigInt)> = None;
for (a, b) in iv {
cur = match cur {
Some((ca, cb)) if a <= &cb + BigInt::one() => Some((ca, if b > cb { b } else { cb })),
Some((ca, cb)) => {
total += cb - ca + BigInt::one();
Some((a, b))
}
None => Some((a, b)),
};
}
if let Some((a, b)) = cur {
total += b - a + BigInt::one();
}
total
}
impl Dist {
pub fn validate(&self) -> Result<(), IrError> {
let state_only = |e: &Expr| -> Result<(), IrError> {
e.validate()?;
let mut has_var = false;
e.for_each_var(&mut |_, _| has_var = true);
if has_var {
return Err(IrError::IllTyped {
what: "a `dist` weight or bound reads a random variable".into(),
});
}
Ok(())
};
for it in &self.items {
it.member.validate()?;
state_only(&it.weight)?;
match &it.span {
Span::Point(v) => state_only(v)?,
Span::Range { lo, hi } => {
for b in [lo, hi].into_iter().flatten() {
state_only(b)?;
}
}
}
}
if let Some(d) = &self.default {
state_only(&d.weight)?;
}
Ok(())
}
pub fn codes(&self) -> usize {
self.items.len() + usize::from(self.default.is_some())
}
pub fn selector_ty(&self) -> Ty {
let top = self.codes().saturating_sub(1) as u64;
Ty::unsigned((64 - top.leading_zeros()).max(1))
}
pub fn conjunct(&self, sel: VarId) -> Expr {
let kt = self.selector_ty();
let bin = |op, a: Expr, b: Expr| Expr::Binary(op, Box::new(a), Box::new(b), Ty::BIT);
let code = |i: usize| {
bin(
super::ir::BinOp::Eq,
Expr::Var(sel, kt),
Expr::constant(kt, i as u64),
)
};
let mut terms: Vec<Expr> = self
.items
.iter()
.enumerate()
.map(|(i, it)| bin(super::ir::BinOp::LogAnd, code(i), it.member.clone()))
.collect();
if self.default.is_some() {
let any = self
.items
.iter()
.map(|it| it.member.clone())
.reduce(|a, b| bin(super::ir::BinOp::LogOr, a, b));
let none = match any {
Some(a) => Expr::Unary(super::ir::UnOp::LogNot, Box::new(a), Ty::BIT),
None => Expr::Const(Ty::BIT, BigUint::one()),
};
terms.push(bin(super::ir::BinOp::LogAnd, code(self.items.len()), none));
}
terms
.into_iter()
.reduce(|a, b| bin(super::ir::BinOp::LogOr, a, b))
.unwrap_or(Expr::Const(Ty::BIT, BigUint::zero()))
}
pub fn weights(&self, states: &[Option<BigUint>]) -> Result<Vec<Weight>, DistError> {
let a = Assignment {
vars: Vec::new(),
states: states.to_vec(),
};
let weight_of = |e: &Expr| -> Result<BigUint, DistError> {
let v = numeric(&eval(e, &a)?, e.ty());
v.to_biguint().ok_or(DistError::NegativeWeight)
};
let per = |w: BigUint, per_range: bool, n: &BigInt| -> Weight {
if !per_range {
return Weight::int(w);
}
match n.to_biguint() {
Some(n) if !n.is_zero() => Weight { num: w, den: n },
_ => Weight::int(0u32),
}
};
let mut out = Vec::with_capacity(self.codes());
let mut covered: Vec<(BigInt, BigInt)> = Vec::new();
for it in &self.items {
let (smin, smax) = subject_range_in(it.ty, self.subject_ty);
let bound = |e: &Option<Expr>, dflt: &BigInt| -> Result<BigInt, DistError> {
match e {
Some(e) => Ok(numeric(&eval(e, &a)?, e.ty())),
None => Ok(dflt.clone()),
}
};
let (lo, hi) = match &it.span {
Span::Point(v) => {
let n = numeric(&eval(v, &a)?, v.ty());
(n.clone(), n)
}
Span::Range { lo, hi } => (bound(lo, &smin)?, bound(hi, &smax)?),
};
let n = if hi >= lo {
&hi - &lo + 1
} else {
BigInt::zero()
};
out.push(per(weight_of(&it.weight)?, it.per_range, &n));
if let Span::Range {
lo: Some(l),
hi: Some(h),
} = &it.span
&& l.ty() != h.ty()
{
covered.extend(mixed_bound_cover(
self.subject_ty,
(&lo, l.ty()),
(&hi, h.ty()),
));
continue;
}
let (a_, b_) = (lo.max(smin), hi.min(smax));
if a_ > b_ {
continue;
}
if self.subject_ty.signed && !it.ty.signed {
let half = BigInt::one() << (self.subject_ty.width - 1);
let full = BigInt::one() << self.subject_ty.width;
if a_ < half {
covered.push((a_.clone(), (&half - BigInt::one()).min(b_.clone())));
}
if b_ >= half {
covered.push((a_.max(half) - &full, b_ - &full));
}
} else {
covered.push((a_, b_));
}
}
if let Some(d) = &self.default {
let width = self.subject_ty.width;
let n = (BigInt::one() << width) - union_len(covered);
out.push(per(weight_of(&d.weight)?, d.per_range, &n));
}
Ok(out)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::constraints::ir::BinOp;
fn u(w: u32) -> Ty {
Ty::unsigned(w)
}
fn k(t: Ty, v: u64) -> Expr {
Expr::constant(t, v)
}
fn bin(op: BinOp, a: Expr, b: Expr) -> Expr {
Expr::Binary(op, Box::new(a), Box::new(b), Ty::BIT)
}
fn item(x: &Expr, lo: u64, hi: u64, w: u64, per_range: bool) -> Item {
let t = x.ty();
Item {
member: bin(
BinOp::LogAnd,
bin(BinOp::Ge, x.clone(), k(t, lo)),
bin(BinOp::Le, x.clone(), k(t, hi)),
),
ty: t,
span: Span::Range {
lo: Some(k(t, lo)),
hi: Some(k(t, hi)),
},
weight: k(u(32), w),
per_range,
}
}
fn dist(subject_ty: Ty, items: Vec<Item>, default: Option<DefaultWeight>) -> Dist {
Dist {
subject_ty,
plain_var: Some(0),
vars: vec![0],
items,
default,
subject: None,
}
}
fn w(num: u32, den: u32) -> Weight {
Weight {
num: num.into(),
den: den.into(),
}
}
#[test]
fn the_selector_has_one_code_per_item_and_the_default() {
let x = Expr::Var(0, u(8));
let d = dist(u(8), vec![item(&x, 1, 1, 1, false)], None);
assert_eq!((d.codes(), d.selector_ty()), (1, u(1)));
let d = dist(
u(8),
vec![item(&x, 1, 1, 1, false), item(&x, 2, 2, 1, false)],
None,
);
assert_eq!((d.codes(), d.selector_ty()), (2, u(1)));
let d = dist(
u(8),
vec![item(&x, 1, 1, 1, false), item(&x, 2, 2, 1, false)],
Some(DefaultWeight {
weight: k(u(32), 1),
per_range: true,
}),
);
assert_eq!((d.codes(), d.selector_ty()), (3, u(2)));
}
#[test]
fn per_value_and_per_range_weights() {
let x = Expr::Var(0, u(8));
let d = dist(
u(8),
vec![item(&x, 0, 3, 5, false), item(&x, 4, 11, 8, true)],
None,
);
assert_eq!(d.weights(&[]).unwrap(), vec![w(5, 1), w(8, 8)]);
}
#[test]
fn the_default_weight_spreads_over_the_values_no_item_names() {
let x = Expr::Var(0, u(4));
let d = dist(
u(4),
vec![item(&x, 0, 3, 1, false), item(&x, 2, 5, 1, false)],
Some(DefaultWeight {
weight: k(u(32), 20),
per_range: true,
}),
);
assert_eq!(d.weights(&[]).unwrap()[2], w(20, 10));
}
#[test]
fn an_item_beyond_the_subjects_range_does_not_shrink_the_default() {
let wide = Item {
member: Expr::Const(Ty::BIT, BigUint::one()),
ty: u(32),
span: Span::Range {
lo: Some(k(u(32), 10)),
hi: Some(k(u(32), 40)),
},
weight: k(u(32), 31),
per_range: true,
};
let d = dist(
u(4),
vec![wide],
Some(DefaultWeight {
weight: k(u(32), 10),
per_range: true,
}),
);
assert_eq!(d.weights(&[]).unwrap(), vec![w(31, 31), w(10, 10)]);
}
#[test]
fn an_empty_range_and_a_full_default_have_weight_zero_not_a_zero_divisor() {
let x = Expr::Var(0, u(2));
let d = dist(
u(2),
vec![item(&x, 3, 1, 4, true), item(&x, 0, 3, 1, false)],
Some(DefaultWeight {
weight: k(u(32), 9),
per_range: true,
}),
);
assert_eq!(d.weights(&[]).unwrap(), vec![w(0, 1), w(1, 1), w(0, 1)]);
}
#[test]
fn weights_and_bounds_read_state_variables() {
let st = |i: u32| Expr::State(i, u(8));
let d = dist(
u(8),
vec![Item {
member: Expr::Const(Ty::BIT, BigUint::one()),
ty: u(8),
span: Span::Range {
lo: Some(st(0)),
hi: Some(st(1)),
},
weight: st(2),
per_range: true,
}],
None,
);
let states = [Some(0u32.into()), Some(9u32.into()), Some(20u32.into())];
assert_eq!(d.weights(&states).unwrap(), vec![w(20, 10)]);
let unknown = [Some(0u32.into()), None, Some(20u32.into())];
assert_eq!(
d.weights(&unknown),
Err(DistError::Eval(ConstraintEvalError::StateUnknown(1)))
);
}
#[test]
fn a_negative_weight_is_an_error() {
let x = Expr::Var(0, u(8));
let mut it = item(&x, 1, 1, 1, false);
it.weight = Expr::constant(Ty::signed(8), 0xFFu32); let d = dist(u(8), vec![it], None);
assert_eq!(d.weights(&[]), Err(DistError::NegativeWeight));
}
#[test]
fn a_signed_subject_counts_negative_values_in_the_default() {
let t = Ty::signed(4);
let neg2 = Expr::constant(t, 0xEu32);
let one = Expr::constant(t, 1u32);
let d = dist(
t,
vec![Item {
member: Expr::Const(Ty::BIT, BigUint::one()),
ty: t,
span: Span::Range {
lo: Some(neg2),
hi: Some(one),
},
weight: k(u(32), 1),
per_range: false,
}],
Some(DefaultWeight {
weight: k(u(32), 12),
per_range: true,
}),
);
assert_eq!(d.weights(&[]).unwrap(), vec![w(1, 1), w(12, 12)]);
}
#[test]
fn range_bounds_of_different_types_cover_the_intersection() {
let t = Ty::signed(4);
let d = dist(
t,
vec![Item {
member: Expr::Const(Ty::BIT, BigUint::one()),
ty: t,
span: Span::Range {
lo: Some(Expr::constant(t, 0xEu32)),
hi: Some(Expr::constant(u(4), 12u32)),
},
weight: k(u(32), 1),
per_range: false,
}],
Some(DefaultWeight {
weight: k(u(32), 8),
per_range: true,
}),
);
assert_eq!(d.weights(&[]).unwrap(), vec![w(1, 1), w(8, 8)]);
}
#[test]
fn the_conjunct_is_membership_of_exactly_one_item() {
let x = Expr::Var(0, u(2));
let d = dist(
u(2),
vec![item(&x, 1, 1, 1, false), item(&x, 2, 2, 1, false)],
Some(DefaultWeight {
weight: k(u(32), 1),
per_range: true,
}),
);
let c = d.conjunct(1);
let kt = d.selector_ty();
let ok = |xv: u64, kv: u64| {
let a = Assignment {
vars: vec![xv.into(), kv.into()],
states: vec![],
};
eval(&c, &a).unwrap() == BigUint::one()
};
for xv in 0..4 {
for kv in 0..(1u64 << kt.width) {
let want = match kv {
0 => xv == 1,
1 => xv == 2,
2 => xv == 0 || xv == 3,
_ => false,
};
assert_eq!(ok(xv, kv), want, "x={xv} K={kv}");
}
}
}
}