use super::dist::{DefaultWeight, Dist, Item, SELECTOR_BASE, Span};
use super::func::{self, Func, FuncError};
use super::guard::{self, GuardValue};
use super::ir::{
ArrayId, ArrayNode, Assignment, BinOp, ConstraintEvalError, Expr, LoopId, ReduceOp, Ty, UnOp,
VarId, eval, mask,
};
use super::solve::{SolveError, Solver};
use num_bigint::BigUint;
use num_traits::{One, Zero};
use std::collections::HashMap;
pub const SIZE_TIERS: [u64; 2] = [1024, 65536];
pub const MAX_SIZE_RETRIES: usize = 256;
#[derive(Clone, Debug)]
pub struct ArrayState {
pub elem_ty: Ty,
pub active: bool,
pub cur: Vec<Option<BigUint>>,
pub domain: Option<Vec<BigUint>>,
pub rows: Option<RowsState>,
pub template_of: Option<ArrayId>,
}
#[derive(Clone, Debug)]
pub struct RowsState {
pub template: ArrayId,
pub cur_rows: Vec<Vec<Option<BigUint>>>,
pub row_base: ArrayId,
}
impl ArrayState {
pub fn flat(elem_ty: Ty, active: bool, cur: Vec<Option<BigUint>>) -> ArrayState {
ArrayState {
elem_ty,
active,
cur,
domain: None,
rows: None,
template_of: None,
}
}
pub fn with_domain(mut self, domain: Option<Vec<BigUint>>) -> ArrayState {
self.domain = domain;
self
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ExpandError {
IndexOutOfRange {
array: ArrayId,
index: i64,
size: usize,
},
LeafIndexOutOfRange {
index: i64,
size: usize,
},
NonConstantIndex,
GuardError,
StateUnknown(u32),
Func(FuncError),
DistRandomBound,
Malformed(&'static str),
}
#[derive(Clone, Debug, Default)]
pub struct Expanded {
pub conjuncts: Vec<Expr>,
pub elem_vars: Vec<(ArrayId, usize)>,
pub elem_tys: Vec<Ty>,
pub extra_states: Vec<Option<BigUint>>,
pub extra_dists: Vec<Dist>,
}
struct Expander<'a> {
arrays: &'a [ArrayState],
sizes: &'a [usize],
elem_base: Vec<u32>,
asg: Assignment,
synth: HashMap<(ArrayId, usize), u32>,
env: Vec<(LoopId, Expr)>,
funcs: &'a [Func],
dists: &'a [Dist],
extra_dists: Vec<Dist>,
instanced: Vec<InstancedDist>,
rebind: Vec<(ArrayId, ArrayId)>,
}
type InstancedDist = (usize, Vec<(LoopId, Expr)>, usize);
fn konst(t: Ty, v: u64) -> Expr {
Expr::constant(t, v)
}
fn has_vars(e: &Expr) -> bool {
let mut any = false;
e.for_each_var(&mut |_, _| any = true);
any
}
impl Expander<'_> {
fn array(&self, a: ArrayId) -> Result<&ArrayState, ExpandError> {
self.arrays.get(a as usize).ok_or(ExpandError::Malformed(
"an array id outside the program's table",
))
}
fn res(&self, a: ArrayId) -> ArrayId {
self.rebind
.iter()
.rev()
.find(|(t, _)| *t == a)
.map_or(a, |(_, r)| *r)
}
fn size_of(&self, a: ArrayId) -> Result<usize, ExpandError> {
let a = self.res(a);
self.sizes
.get(a as usize)
.copied()
.ok_or(ExpandError::Malformed("no size for an array"))
}
fn const_value(&self, e: &Expr) -> Result<BigUint, ExpandError> {
if has_vars(e) {
return Err(ExpandError::NonConstantIndex);
}
eval(e, &self.asg).map_err(|err| match err {
ConstraintEvalError::StateUnknown(s) => ExpandError::StateUnknown(s),
_ => ExpandError::NonConstantIndex,
})
}
fn index_value(&mut self, index: &Expr) -> Result<i64, ExpandError> {
let idx = self.expand(index)?;
let v = self.const_value(&idx)?;
let t = idx.ty();
Ok(if t.signed && v.bit(u64::from(t.width) - 1) {
-i64::try_from((BigUint::one() << t.width) - &v).unwrap_or(i64::MAX)
} else {
i64::try_from(&v).unwrap_or(i64::MAX)
})
}
fn element(&mut self, a: ArrayId, idx: usize, ty: Ty) -> Result<Expr, ExpandError> {
let a = self.res(a);
let arr = self.array(a)?;
if arr.elem_ty != ty {
return Err(ExpandError::Malformed("an element read at the wrong type"));
}
let (active, cur) = (arr.active, arr.cur.get(idx).cloned());
if active {
let id = self.elem_base[a as usize] + u32::try_from(idx).expect("an element index");
return Ok(Expr::Var(id, ty));
}
let Some(cur) = cur else {
return Err(ExpandError::Malformed("an element of a stale array view"));
};
if let Some(&id) = self.synth.get(&(a, idx)) {
return Ok(Expr::State(id, ty));
}
let id = u32::try_from(self.asg.states.len()).expect("a state id");
self.asg.states.push(cur);
self.synth.insert((a, idx), id);
Ok(Expr::State(id, ty))
}
fn lookup(&self, id: LoopId) -> Result<Expr, ExpandError> {
self.env
.iter()
.rev()
.find(|(l, _)| *l == id)
.map(|(_, e)| e.clone())
.ok_or(ExpandError::Malformed("a loop variable outside its loop"))
}
fn expand(&mut self, e: &Expr) -> Result<Expr, ExpandError> {
match e {
Expr::Guard(p, then, otherwise, t) => {
let p = self.expand(p)?;
match guard::value(&p, &self.asg) {
GuardValue::True => self.expand(then),
GuardValue::False => self.expand(otherwise),
GuardValue::Error => Err(ExpandError::GuardError),
GuardValue::Random => Ok(Expr::Guard(
Box::new(p),
Box::new(self.expand(then)?),
Box::new(self.expand(otherwise)?),
*t,
)),
}
}
Expr::Ite(c, x, y, t) => {
let c = self.expand(c)?;
if !has_vars(&c)
&& let Ok(v) = eval(&c, &self.asg)
{
return if v.is_zero() {
self.expand(y)
} else {
self.expand(x)
};
}
Ok(Expr::Ite(
Box::new(c),
Box::new(self.expand(x)?),
Box::new(self.expand(y)?),
*t,
))
}
Expr::Binary(op @ (BinOp::LogAnd | BinOp::LogOr | BinOp::LogImpl), l, r, t) => {
let l = self.expand(l)?;
if !has_vars(&l)
&& let Ok(v) = eval(&l, &self.asg)
{
let l_true = !v.is_zero();
let decided = match op {
BinOp::LogAnd | BinOp::LogImpl => !l_true,
_ => l_true,
};
if decided {
return Ok(konst(*t, u64::from(*op != BinOp::LogAnd)));
}
}
Ok(Expr::Binary(
*op,
Box::new(l),
Box::new(self.expand(r)?),
*t,
))
}
Expr::Array(n) => self.array_node(n),
Expr::Var(v, t) if *v >= SELECTOR_BASE => {
let j = (*v - SELECTOR_BASE) as usize;
match self.dists.get(j) {
Some(d) if d.subject.is_some() => {
let k = match self
.instanced
.iter()
.find(|(tj, env, _)| *tj == j && *env == self.env)
{
Some(&(_, _, k)) => k,
None => {
let inst = self.instantiate(d)?;
self.extra_dists.push(inst);
let k = self.extra_dists.len() - 1;
self.instanced.push((j, self.env.clone(), k));
k
}
};
Ok(Expr::Var(SELECTOR_BASE + (self.dists.len() + k) as u32, *t))
}
_ => Ok(e.clone()),
}
}
other => other.try_map_children(&mut |c| self.expand(c)),
}
}
fn instantiate(&mut self, d: &Dist) -> Result<Dist, ExpandError> {
let Some(subject) = &d.subject else {
return Err(ExpandError::Malformed("an instance of a plain dist"));
};
let subject = self.expand(subject)?;
let mut vars: Vec<VarId> = Vec::new();
subject.for_each_var(&mut |v, _| vars.push(v));
vars.sort_unstable();
vars.dedup();
let plain_var = match &subject {
Expr::Var(v, _) => Some(*v),
_ => None,
};
let fixed = |this: &mut Self, e: &Expr| -> Result<Expr, ExpandError> {
let x = this.expand(e)?;
if has_vars(&x) {
return Err(ExpandError::DistRandomBound);
}
Ok(x)
};
let mut items = Vec::with_capacity(d.items.len());
for it in &d.items {
let span = match &it.span {
Span::Point(v) => Span::Point(fixed(self, v)?),
Span::Range { lo, hi } => Span::Range {
lo: lo.as_ref().map(|b| fixed(self, b)).transpose()?,
hi: hi.as_ref().map(|b| fixed(self, b)).transpose()?,
},
};
items.push(Item {
member: konst(Ty::BIT, 1),
ty: it.ty,
span,
weight: fixed(self, &it.weight)?,
per_range: it.per_range,
});
}
let default = match &d.default {
Some(w) => Some(DefaultWeight {
weight: fixed(self, &w.weight)?,
per_range: w.per_range,
}),
None => None,
};
Ok(Dist {
subject_ty: d.subject_ty,
plain_var,
vars,
items,
default,
subject: None,
})
}
fn row_array(&mut self, template: ArrayId, row: &Expr) -> Result<ArrayId, ExpandError> {
let outer = self
.array(template)?
.template_of
.ok_or(ExpandError::Malformed("a row scope over a non-template"))?;
let k = self.index_value(row)?;
let size = self.size_of(outer)?;
if k < 0 || k as usize >= size {
return Err(ExpandError::IndexOutOfRange {
array: outer,
index: k,
size,
});
}
let base = self
.array(outer)?
.rows
.as_ref()
.ok_or(ExpandError::Malformed("rows of an array without rows"))?
.row_base;
Ok(base + k as ArrayId)
}
fn array_node(&mut self, n: &ArrayNode) -> Result<Expr, ExpandError> {
match n {
ArrayNode::Call { func, args, ty } => {
let args = args
.iter()
.map(|a| self.expand(a))
.collect::<Result<Vec<_>, _>>()?;
let call = Expr::Array(ArrayNode::Call {
func: *func,
args,
ty: *ty,
});
func::fold_ground_calls(&call, self.funcs, &self.asg.states).map_err(|e| match e {
FuncError::Eval(ConstraintEvalError::StateUnknown(s)) => {
ExpandError::StateUnknown(s)
}
other => ExpandError::Func(other),
})
}
ArrayNode::Size(a) => Ok(konst(Ty::signed(32), self.size_of(*a)? as u64)),
ArrayNode::Loop(id, t) => {
let v = self.lookup(*id)?;
if v.ty() != *t {
return Err(ExpandError::Malformed(
"a loop variable bound at another type",
));
}
Ok(v)
}
ArrayNode::Pick {
index,
low,
elems,
ty,
} => {
let i = self.index_value(index)?;
let pos = i
.checked_sub(*low)
.and_then(|p| usize::try_from(p).ok())
.filter(|&p| p < elems.len());
let Some(pos) = pos else {
return Err(ExpandError::LeafIndexOutOfRange {
index: i,
size: elems.len(),
});
};
let x = self.expand(&elems[pos])?;
if x.ty() != *ty {
return Err(ExpandError::Malformed("a pick element of another type"));
}
Ok(x)
}
ArrayNode::Elem { array, index, ty } => {
let i = self.index_value(index)?;
let size = self.size_of(*array)?;
if i < 0 || i as usize >= size {
return Err(ExpandError::IndexOutOfRange {
array: self.res(*array),
index: i,
size,
});
}
self.element(*array, i as usize, *ty)
}
ArrayNode::RowScope {
template,
row,
body,
} => {
let target = self.row_array(*template, row)?;
self.rebind.push((*template, target));
let r = self.expand(body);
self.rebind.pop();
r
}
ArrayNode::Distinct { scalars, arrays } => {
let mut members = Vec::new();
for s in scalars {
members.push(self.expand(s)?);
}
for a in arrays {
let size = self.size_of(*a)?;
let elem_ty = self.array(self.res(*a))?.elem_ty;
for i in 0..size {
members.push(self.element(*a, i, elem_ty)?);
}
}
if members.iter().any(|m| m.ty() != members[0].ty()) {
return Err(ExpandError::Malformed(
"the members of a unique group are of different types",
));
}
Ok(if members.len() < 2 {
konst(Ty::BIT, 1)
} else {
Expr::AllDistinct(members)
})
}
ArrayNode::Foreach { array, var, body } => {
let size = self.size_of(*array)?;
let mut acc: Option<Expr> = None;
for i in 0..size {
self.env.push((*var, konst(Ty::signed(32), i as u64)));
let b = self.expand(body);
self.env.pop();
let b = b?;
acc = Some(match acc {
None => b,
Some(a) => Expr::Binary(BinOp::LogAnd, Box::new(a), Box::new(b), Ty::BIT),
});
}
Ok(acc.unwrap_or_else(|| konst(Ty::BIT, 1)))
}
ArrayNode::Reduce {
op,
array,
var,
index_var,
body,
} => {
let size = self.size_of(*array)?;
let elem_ty = self.array(self.res(*array))?.elem_ty;
let ty = body.ty();
let mut acc: Option<Expr> = None;
for i in 0..size {
let item = self.element(*array, i, elem_ty)?;
self.env.push((*var, item));
self.env.push((*index_var, konst(Ty::signed(32), i as u64)));
let b = self.expand(body);
self.env.pop();
self.env.pop();
let b = b?;
acc = Some(match acc {
None => b,
Some(a) => {
let bop = match op {
ReduceOp::Sum => BinOp::Add,
ReduceOp::Product => BinOp::Mul,
ReduceOp::And => BinOp::And,
ReduceOp::Or => BinOp::Or,
ReduceOp::Xor => BinOp::Xor,
};
Expr::Binary(bop, Box::new(a), Box::new(b), ty)
}
});
}
Ok(acc.unwrap_or_else(|| {
Expr::Const(
ty,
match op {
ReduceOp::Product => BigUint::one(),
ReduceOp::And => mask(ty.width),
ReduceOp::Sum | ReduceOp::Or | ReduceOp::Xor => BigUint::zero(),
},
)
}))
}
}
}
}
fn domain_constraint(v: Expr, domain: &[BigUint]) -> Expr {
let t = v.ty();
let mut acc: Option<Expr> = None;
for d in domain {
let eq = Expr::Binary(
BinOp::Eq,
Box::new(v.clone()),
Box::new(Expr::Const(t, d.clone() & mask(t.width))),
Ty::BIT,
);
acc = Some(match acc {
None => eq,
Some(a) => Expr::Binary(BinOp::LogOr, Box::new(a), Box::new(eq), Ty::BIT),
});
}
acc.unwrap_or_else(|| konst(Ty::BIT, 0))
}
pub fn expand(
conjuncts: &[Expr],
arrays: &[ArrayState],
sizes: &[usize],
n_vars0: u32,
states: &[Option<BigUint>],
) -> Result<Expanded, ExpandError> {
expand_with(conjuncts, arrays, sizes, n_vars0, states, &[], &[])
}
pub fn expand_with(
conjuncts: &[Expr],
arrays: &[ArrayState],
sizes: &[usize],
n_vars0: u32,
states: &[Option<BigUint>],
funcs: &[Func],
dists: &[Dist],
) -> Result<Expanded, ExpandError> {
if arrays.len() != sizes.len() {
return Err(ExpandError::Malformed("one size per array"));
}
let mut out = Expanded::default();
let mut elem_base = vec![0u32; arrays.len()];
let mut next = n_vars0;
for (a, arr) in arrays.iter().enumerate() {
if !arr.active {
continue;
}
elem_base[a] = next;
for i in 0..sizes[a] {
out.elem_vars.push((a as ArrayId, i));
out.elem_tys.push(arr.elem_ty);
next += 1;
if let Some(dom) = &arr.domain {
let id = elem_base[a] + i as u32;
out.conjuncts
.push(domain_constraint(Expr::Var(id, arr.elem_ty), dom));
}
}
}
let mut ex = Expander {
arrays,
sizes,
elem_base,
asg: Assignment {
vars: Vec::new(),
states: states.to_vec(),
},
synth: HashMap::new(),
env: Vec::new(),
funcs,
dists,
extra_dists: Vec::new(),
instanced: Vec::new(),
rebind: Vec::new(),
};
for c in conjuncts {
let e = ex.expand(c)?;
out.conjuncts.push(e);
}
out.extra_dists = std::mem::take(&mut ex.extra_dists);
out.extra_states = ex.asg.states.split_off(states.len());
Ok(out)
}
#[derive(Clone, Debug)]
pub struct Lowered {
pub conjuncts: Vec<Expr>,
pub arrays: Vec<ArrayState>,
pub bases: Vec<(ArrayId, ArrayId)>,
}
struct Lowerer<'a> {
arrays: &'a [ArrayState],
sizes: &'a [usize],
states: &'a [Option<BigUint>],
bases: &'a [(ArrayId, ArrayId)],
env: Vec<(LoopId, i64)>,
rebind: Vec<(ArrayId, ArrayId)>,
}
impl Lowerer<'_> {
fn res(&self, a: ArrayId) -> ArrayId {
self.rebind
.iter()
.rev()
.find(|(t, _)| *t == a)
.map_or(a, |(_, r)| *r)
}
fn base_of(&self, outer: ArrayId) -> Option<ArrayId> {
self.bases
.iter()
.find(|(o, _)| *o == outer)
.map(|(_, b)| *b)
}
fn fixed(&self, row: &Expr) -> Option<i64> {
let row = replace_sizes(row, &mut |a| {
let a = self.res(a);
konst(
Ty::signed(32),
self.sizes.get(a as usize).copied().unwrap_or(0) as u64,
)
});
let asg = Assignment {
vars: Vec::new(),
states: self.states.to_vec(),
};
let v = eval(&row, &asg).ok()?;
let t = row.ty();
Some(if t.signed && v.bit(u64::from(t.width) - 1) {
-i64::try_from((BigUint::one() << t.width) - &v).ok()?
} else {
i64::try_from(&v).ok()?
})
}
fn low(&mut self, e: &Expr) -> Result<Expr, ExpandError> {
let Expr::Array(n) = e else {
return e.try_map_children(&mut |c| self.low(c));
};
match n {
ArrayNode::Loop(id, t) => Ok(match self.env.iter().rev().find(|(l, _)| l == id) {
Some(&(_, i)) => Expr::constant(*t, i as u64),
None => e.clone(),
}),
ArrayNode::Size(a) => Ok(Expr::Array(ArrayNode::Size(self.res(*a)))),
ArrayNode::Elem { array, index, ty } => Ok(Expr::Array(ArrayNode::Elem {
array: self.res(*array),
index: Box::new(self.low(index)?),
ty: *ty,
})),
ArrayNode::Distinct { scalars, arrays } => Ok(Expr::Array(ArrayNode::Distinct {
scalars: scalars
.iter()
.map(|x| self.low(x))
.collect::<Result<_, _>>()?,
arrays: arrays.iter().map(|a| self.res(*a)).collect(),
})),
ArrayNode::Reduce {
op,
array,
var,
index_var,
body,
} => Ok(Expr::Array(ArrayNode::Reduce {
op: *op,
array: self.res(*array),
var: *var,
index_var: *index_var,
body: Box::new(self.low(body)?),
})),
ArrayNode::Foreach { array, var, body } => {
let a = self.res(*array);
if self.base_of(a).is_none() {
return Ok(Expr::Array(ArrayNode::Foreach {
array: a,
var: *var,
body: Box::new(self.low(body)?),
}));
}
let mut acc: Option<Expr> = None;
for i in 0..self.sizes[a as usize] {
self.env.push((*var, i as i64));
let b = self.low(body);
self.env.pop();
let b = b?;
acc = Some(match acc {
None => b,
Some(x) => Expr::Binary(BinOp::LogAnd, Box::new(x), Box::new(b), Ty::BIT),
});
}
Ok(acc.unwrap_or_else(|| konst(Ty::BIT, 1)))
}
ArrayNode::RowScope {
template,
row,
body,
} => {
let row = self.low(row)?;
let outer = self
.arrays
.get(*template as usize)
.and_then(|t| t.template_of)
.ok_or(ExpandError::Malformed("a row scope over a non-template"))?;
let (Some(k), Some(base)) = (self.fixed(&row), self.base_of(outer)) else {
return Ok(Expr::Array(ArrayNode::RowScope {
template: *template,
row: Box::new(row),
body: Box::new(self.low(body)?),
}));
};
let size = self.sizes[outer as usize];
if k < 0 || k as usize >= size {
return Err(ExpandError::IndexOutOfRange {
array: outer,
index: k,
size,
});
}
self.rebind.push((*template, base + k as ArrayId));
let r = self.low(body);
self.rebind.pop();
r
}
ArrayNode::Call { .. } | ArrayNode::Pick { .. } => {
e.try_map_children(&mut |c| self.low(c))
}
}
}
}
pub fn lower_rows(
conjuncts: &[Expr],
arrays: &[ArrayState],
sizes: &[usize],
states: &[Option<BigUint>],
) -> Result<Lowered, ExpandError> {
let mut out: Vec<ArrayState> = arrays.to_vec();
let mut bases = Vec::new();
let mut sizes2: Vec<usize> = sizes.to_vec();
for (a, arr) in arrays.iter().enumerate() {
let Some(rows) = &arr.rows else { continue };
let n = sizes[a];
let base = out.len() as ArrayId;
bases.push((a as ArrayId, base));
for k in 0..n {
out.push(ArrayState {
elem_ty: arr.elem_ty,
active: arr.active,
cur: rows.cur_rows.get(k).cloned().unwrap_or_default(),
domain: arr.domain.clone(),
rows: None,
template_of: None,
});
sizes2.push(rows.cur_rows.get(k).map_or(0, Vec::len));
}
out[a].active = false;
out[a].cur = vec![Some(BigUint::zero()); n];
if let Some(r) = out[a].rows.as_mut() {
r.row_base = base;
}
}
let mut lw = Lowerer {
arrays,
sizes: &sizes2,
states,
bases: &bases,
env: Vec::new(),
rebind: Vec::new(),
};
let lowered = conjuncts
.iter()
.map(|c| lw.low(c))
.collect::<Result<Vec<_>, _>>()?;
Ok(Lowered {
conjuncts: lowered,
arrays: out,
bases,
})
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SizeError {
Unsatisfiable,
TooLarge,
GuardError,
Solve(SolveError),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SizeSolution {
pub sizes: Vec<usize>,
pub sized: Vec<ArrayId>,
pub tuple: Vec<usize>,
}
fn size_eligible(c: &Expr) -> bool {
!c.any_node(&mut |n| match n {
Expr::Array(
ArrayNode::Elem { .. }
| ArrayNode::Loop(..)
| ArrayNode::Foreach { .. }
| ArrayNode::Reduce { .. }
| ArrayNode::Distinct { .. }
| ArrayNode::Pick { .. }
| ArrayNode::RowScope { .. }
| ArrayNode::Call { .. },
) => true,
Expr::Var(v, _) => *v >= super::dist::SELECTOR_BASE,
_ => false,
})
}
fn replace_sizes(e: &Expr, f: &mut impl FnMut(ArrayId) -> Expr) -> Expr {
match e {
Expr::Array(ArrayNode::Size(a)) => f(*a),
other => other
.try_map_children::<()>(&mut |c| Ok(replace_sizes(c, f)))
.expect("infallible"),
}
}
fn find(parent: &mut [usize], x: usize) -> usize {
let mut r = x;
while parent[r] != r {
r = parent[r];
}
let mut c = x;
while parent[c] != r {
let n = parent[c];
parent[c] = r;
c = n;
}
r
}
pub fn solve_sizes(
conjuncts: &[Expr],
arrays: &[ArrayState],
var_tys: &[Ty],
states: &[Option<BigUint>],
blocked: &[Vec<usize>],
seed: &mut u64,
) -> Result<SizeSolution, SizeError> {
let current: Vec<usize> = arrays.iter().map(|a| a.cur.len()).collect();
let n_fields = var_tys.len() as u32;
let mut sized: Vec<ArrayId> = Vec::new();
let mut eligible: Vec<&Expr> = Vec::new();
for c in conjuncts {
if !size_eligible(c) {
continue;
}
c.any_node(&mut |n| {
if let Expr::Array(ArrayNode::Size(a)) = n
&& arrays.get(*a as usize).is_some_and(|x| x.active)
&& !sized.contains(a)
{
sized.push(*a);
}
false
});
eligible.push(c);
}
if sized.is_empty() {
return Ok(SizeSolution {
sizes: current,
sized,
tuple: Vec::new(),
});
}
sized.sort_unstable();
let size_var = |a: ArrayId| -> u32 {
n_fields + u32::try_from(sized.iter().position(|x| *x == a).expect("sized")).unwrap()
};
let total = n_fields as usize + sized.len();
let subbed: Vec<Expr> = eligible
.iter()
.map(|c| {
replace_sizes(c, &mut |a| {
if arrays[a as usize].active {
Expr::Var(size_var(a), Ty::signed(32))
} else {
konst(Ty::signed(32), current[a as usize] as u64)
}
})
})
.collect();
let mut parent: Vec<usize> = (0..total).collect();
let vars_of: Vec<Vec<VarId>> = subbed
.iter()
.map(|c| {
let mut v = Vec::new();
c.for_each_var(&mut |x, _| v.push(x));
v
})
.collect();
for vs in &vars_of {
for w in vs.windows(2) {
let (a, b) = (
find(&mut parent, w[0] as usize),
find(&mut parent, w[1] as usize),
);
parent[a] = b;
}
}
let size_roots: Vec<usize> = (0..sized.len())
.map(|k| find(&mut parent, n_fields as usize + k))
.collect();
let mut keep: Vec<Expr> = Vec::new();
let mut used = vec![false; total];
for k in 0..sized.len() {
used[n_fields as usize + k] = true;
}
for (c, vs) in subbed.iter().zip(&vars_of) {
let Some(&v0) = vs.first() else { continue };
let root = find(&mut parent, v0 as usize);
if size_roots.contains(&root) {
for &v in vs {
used[v as usize] = true;
}
keep.push(c.clone());
}
}
let mut new_id = vec![0u32; total];
let mut tys: Vec<Ty> = Vec::new();
for (v, u) in used.iter().enumerate() {
if *u {
new_id[v] = tys.len() as u32;
tys.push(if v < n_fields as usize {
var_tys[v]
} else {
Ty::signed(32)
});
}
}
let mut keep: Vec<Expr> = keep
.iter()
.map(|c| c.substitute_vars(&mut |v, t| Expr::Var(new_id[v as usize], t)))
.collect();
let svar = |k: usize| Expr::Var(new_id[n_fields as usize + k], Ty::signed(32));
for tup in blocked {
let mut same: Option<Expr> = None;
for (k, &sz) in tup.iter().enumerate() {
let eq = Expr::Binary(
BinOp::Eq,
Box::new(svar(k)),
Box::new(konst(Ty::signed(32), sz as u64)),
Ty::BIT,
);
same = Some(match same {
None => eq,
Some(a) => Expr::Binary(BinOp::LogAnd, Box::new(a), Box::new(eq), Ty::BIT),
});
}
if let Some(s) = same {
keep.push(Expr::Unary(UnOp::LogNot, Box::new(s), Ty::BIT));
}
}
let assignment = Assignment {
vars: Vec::new(),
states: states.to_vec(),
};
let order: Vec<VarId> = (0..tys.len() as u32).collect();
let tiers: [u64; 3] = [SIZE_TIERS[0], SIZE_TIERS[1], i32::MAX as u64];
for (ti, bound) in tiers.iter().enumerate() {
let mut cs = keep.clone();
for k in 0..sized.len() {
cs.push(Expr::Binary(
BinOp::Ge,
Box::new(svar(k)),
Box::new(konst(Ty::signed(32), 0)),
Ty::BIT,
));
cs.push(Expr::Binary(
BinOp::Le,
Box::new(svar(k)),
Box::new(konst(Ty::signed(32), *bound)),
Ty::BIT,
));
}
let mut resolved = Vec::with_capacity(cs.len());
for c in &cs {
resolved.push(guard::resolve(c, &assignment).map_err(|_| SizeError::GuardError)?);
}
let solver = match Solver::new(&tys, states, &resolved) {
Ok(s) => s,
Err(SolveError::Unsatisfiable) => continue,
Err(e) => return Err(SizeError::Solve(e)),
};
let values = match solver.sample(&order, seed) {
Ok(v) => v,
Err(SolveError::Unsatisfiable) => continue,
Err(e) => return Err(SizeError::Solve(e)),
};
if ti == 2 {
return Err(SizeError::TooLarge);
}
let mut sizes = current;
let mut tuple = Vec::with_capacity(sized.len());
for (k, &a) in sized.iter().enumerate() {
let v = &values[new_id[n_fields as usize + k] as usize];
let n = usize::try_from(v).expect("a bounded size");
sizes[a as usize] = n;
tuple.push(n);
}
return Ok(SizeSolution {
sizes,
sized,
tuple,
});
}
Err(SizeError::Unsatisfiable)
}
#[cfg(test)]
mod tests {
use super::*;
fn s32() -> Ty {
Ty::signed(32)
}
fn u(w: u32) -> Ty {
Ty::unsigned(w)
}
fn size(a: ArrayId) -> Expr {
Expr::Array(ArrayNode::Size(a))
}
fn elem(a: ArrayId, i: Expr, t: Ty) -> Expr {
Expr::Array(ArrayNode::Elem {
array: a,
index: Box::new(i),
ty: t,
})
}
fn lp(id: LoopId) -> Expr {
Expr::Array(ArrayNode::Loop(id, s32()))
}
fn bin(op: BinOp, a: Expr, b: Expr, t: Ty) -> Expr {
Expr::Binary(op, Box::new(a), Box::new(b), t)
}
fn arr(active: bool, cur: &[u64]) -> ArrayState {
ArrayState {
elem_ty: s32(),
active,
cur: cur.iter().map(|v| Some(BigUint::from(*v))).collect(),
domain: None,
rows: None,
template_of: None,
}
}
fn ev(e: &Expr, vars: &[u64], states: &[Option<BigUint>]) -> bool {
let a = Assignment {
vars: vars.iter().map(|v| BigUint::from(*v)).collect(),
states: states.to_vec(),
};
!eval(e, &a).expect("evaluates").is_zero()
}
#[test]
fn foreach_unrolls_over_the_size() {
let body = bin(
BinOp::Gt,
elem(0, lp(0), s32()),
bin(BinOp::Mul, konst(s32(), 2), lp(0), s32()),
Ty::BIT,
);
let c = Expr::Array(ArrayNode::Foreach {
array: 0,
var: 0,
body: Box::new(body),
});
let x = expand(&[c], &[arr(true, &[])], &[3], 0, &[]).expect("expands");
assert_eq!(x.elem_vars, vec![(0, 0), (0, 1), (0, 2)]);
assert_eq!(x.conjuncts.len(), 1);
assert!(ev(&x.conjuncts[0], &[1, 3, 5], &[]));
assert!(!ev(&x.conjuncts[0], &[1, 2, 5], &[]));
assert!(!ev(&x.conjuncts[0], &[0, 3, 5], &[]));
}
#[test]
fn a_constant_guard_keeps_the_last_neighbour_unread() {
let pred = bin(
BinOp::Lt,
lp(0),
bin(BinOp::Sub, size(0), konst(s32(), 1), s32()),
Ty::BIT,
);
let then = bin(
BinOp::Lt,
elem(0, lp(0), s32()),
elem(0, bin(BinOp::Add, lp(0), konst(s32(), 1), s32()), s32()),
Ty::BIT,
);
let body = Expr::Guard(
Box::new(pred),
Box::new(then),
Box::new(konst(Ty::BIT, 1)),
Ty::BIT,
);
let c = Expr::Array(ArrayNode::Foreach {
array: 0,
var: 0,
body: Box::new(body),
});
let x = expand(&[c], &[arr(true, &[])], &[4], 0, &[]).expect("expands");
assert_eq!(x.elem_vars.len(), 4);
assert!(ev(&x.conjuncts[0], &[1, 2, 3, 4], &[]));
assert!(!ev(&x.conjuncts[0], &[1, 3, 2, 4], &[]));
}
#[test]
fn an_unguarded_out_of_range_index_is_a_loud_error() {
let c = bin(
BinOp::Eq,
elem(0, konst(s32(), 5), s32()),
konst(s32(), 0),
Ty::BIT,
);
assert_eq!(
expand(&[c], &[arr(true, &[])], &[2], 0, &[]).unwrap_err(),
ExpandError::IndexOutOfRange {
array: 0,
index: 5,
size: 2
}
);
}
#[test]
fn a_pick_selects_the_element_its_index_names() {
let pick = |idx: Expr| {
Expr::Array(ArrayNode::Pick {
index: Box::new(idx),
low: 2,
elems: vec![Expr::Var(0, u(8)), Expr::Var(1, u(8)), Expr::Var(2, u(8))],
ty: u(8),
})
};
let states = [Some(BigUint::from(3u32))];
let c = bin(
BinOp::Eq,
pick(Expr::State(0, s32())),
konst(u(8), 9),
Ty::BIT,
);
let x = expand(&[c], &[], &[], 3, &states).expect("expands");
assert_eq!(x.conjuncts.len(), 1);
assert!(ev(&x.conjuncts[0], &[0, 9, 0], &states));
assert!(!ev(&x.conjuncts[0], &[9, 0, 9], &states));
}
#[test]
fn an_out_of_range_pick_is_an_error_only_when_live() {
let pick = Expr::Array(ArrayNode::Pick {
index: Box::new(Expr::State(0, s32())),
low: 0,
elems: vec![Expr::Var(0, u(8)), Expr::Var(1, u(8))],
ty: u(8),
});
let body = bin(BinOp::Eq, pick, konst(u(8), 1), Ty::BIT);
let guarded = Expr::Guard(
Box::new(bin(
BinOp::Lt,
Expr::State(0, s32()),
konst(s32(), 2),
Ty::BIT,
)),
Box::new(body.clone()),
Box::new(konst(Ty::BIT, 1)),
Ty::BIT,
);
let out = [Some(BigUint::from(7u32))];
let x = expand(&[guarded], &[], &[], 2, &out).expect("the guard drops the pick");
assert!(ev(&x.conjuncts[0], &[], &out));
assert_eq!(
expand(std::slice::from_ref(&body), &[], &[], 2, &out).unwrap_err(),
ExpandError::LeafIndexOutOfRange { index: 7, size: 2 }
);
let neg = [Some((BigUint::from(1u32) << 32) - 1u32)];
assert_eq!(
expand(&[body], &[], &[], 2, &neg).unwrap_err(),
ExpandError::LeafIndexOutOfRange { index: -1, size: 2 }
);
}
#[test]
fn an_inactive_array_is_state() {
let c = bin(
BinOp::Eq,
elem(0, konst(s32(), 1), s32()),
konst(s32(), 7),
Ty::BIT,
);
let x = expand(&[c], &[arr(false, &[3, 7])], &[2], 0, &[]).expect("expands");
assert!(x.elem_vars.is_empty());
assert_eq!(x.extra_states, vec![Some(BigUint::from(7u32))]);
assert!(ev(&x.conjuncts[0], &[], &x.extra_states));
}
#[test]
fn reductions_fold_and_empty_ones_are_the_identity() {
let sum = |op| {
Expr::Array(ArrayNode::Reduce {
op,
array: 0,
var: 0,
index_var: 1,
body: Box::new(Expr::Array(ArrayNode::Loop(0, s32()))),
})
};
let c = bin(BinOp::Eq, sum(ReduceOp::Sum), konst(s32(), 6), Ty::BIT);
let x = expand(&[c], &[arr(true, &[])], &[3], 0, &[]).expect("expands");
assert!(ev(&x.conjuncts[0], &[1, 2, 3], &[]));
assert!(!ev(&x.conjuncts[0], &[1, 2, 4], &[]));
for (op, want) in [
(ReduceOp::Sum, 0u64),
(ReduceOp::Product, 1),
(ReduceOp::Or, 0),
(ReduceOp::Xor, 0),
(ReduceOp::And, 0xFFFF_FFFF),
] {
let c = bin(BinOp::Eq, sum(op), konst(s32(), want), Ty::BIT);
let x = expand(&[c], &[arr(true, &[])], &[0], 0, &[]).expect("expands");
assert!(ev(&x.conjuncts[0], &[], &[]), "{op:?}");
}
}
#[test]
fn a_size_constraint_resizes_the_array() {
let c = bin(
BinOp::Eq,
size(0),
Expr::Resize(Box::new(Expr::Var(0, u(8))), s32()),
Ty::BIT,
);
let bound = bin(BinOp::Lt, Expr::Var(0, u(8)), konst(u(8), 5), Ty::BIT);
let mut seen = std::collections::BTreeSet::new();
let mut seed = 1u64;
for _ in 0..200 {
let s = solve_sizes(
&[c.clone(), bound.clone()],
&[arr(true, &[9, 9])],
&[u(8)],
&[],
&[],
&mut seed,
)
.expect("solves");
assert_eq!(s.sized, vec![0]);
assert!(s.sizes[0] < 5);
seen.insert(s.sizes[0]);
}
assert_eq!(seen.len(), 5, "every size 0..5 is reachable: {seen:?}");
}
#[test]
fn an_unconstrained_size_keeps_the_array() {
let mut seed = 1u64;
let s = solve_sizes(&[], &[arr(true, &[1, 2, 3])], &[], &[], &[], &mut seed).unwrap();
assert_eq!(s.sizes, vec![3]);
assert!(s.sized.is_empty());
}
#[test]
fn an_impossible_size_is_unsatisfiable_and_a_huge_one_is_too_large() {
let mut seed = 1u64;
let lt = |n: u64| bin(BinOp::Lt, size(0), konst(s32(), n), Ty::BIT);
let gt = |n: u64| bin(BinOp::Gt, size(0), konst(s32(), n), Ty::BIT);
let a = [arr(true, &[])];
assert_eq!(
solve_sizes(&[lt(3), gt(5)], &a, &[], &[], &[], &mut seed).unwrap_err(),
SizeError::Unsatisfiable
);
assert_eq!(
solve_sizes(&[gt(10_000_000)], &a, &[], &[], &[], &mut seed).unwrap_err(),
SizeError::TooLarge
);
let s = solve_sizes(&[gt(2000)], &a, &[], &[], &[], &mut seed).unwrap();
assert!(s.sizes[0] > 2000);
}
#[test]
fn blocked_tuples_are_excluded() {
let c = bin(BinOp::Le, size(0), konst(s32(), 1), Ty::BIT);
let a = [arr(true, &[])];
let mut seed = 3u64;
let s = solve_sizes(
std::slice::from_ref(&c),
&a,
&[],
&[],
&[vec![0]],
&mut seed,
)
.unwrap();
assert_eq!(s.sizes, vec![1]);
assert_eq!(
solve_sizes(&[c], &a, &[], &[], &[vec![0], vec![1]], &mut seed).unwrap_err(),
SizeError::Unsatisfiable
);
}
#[test]
fn the_size_problem_holds_only_connected_size_constraints() {
let elem_c = bin(
BinOp::Eq,
elem(0, konst(s32(), 0), s32()),
konst(s32(), 1),
Ty::BIT,
);
let unrelated = bin(BinOp::Eq, Expr::Var(0, u(8)), konst(u(8), 200), Ty::BIT);
let c = bin(BinOp::Eq, size(0), konst(s32(), 2), Ty::BIT);
let mut seed = 1u64;
let s = solve_sizes(
&[elem_c, unrelated, c],
&[arr(true, &[])],
&[u(8)],
&[],
&[],
&mut seed,
)
.unwrap();
assert_eq!(s.sizes, vec![2]);
}
}