vj-sim 0.1.0

//! Array constraints (plan C11, C12, C14; LRM 18.4, 18.5.7).
//!
//! The builder turns `A.size`, `A[i]`, `foreach (A[i]) ...` and `A.sum()
//! with (...)` into array-level IR ([`ArrayNode`]) that does not know how many
//! elements `A` has. Two steps at `randomize()` time turn that into the plain
//! scalar IR the BDD solver understands:
//!
//! 1. **Sizes** ([`solve_sizes`]). LRM 18.4: "a constraint on the size of a
//!    dynamic array or queue ... resizes the array"; an array whose size no
//!    constraint mentions keeps its size. Every `rand`-active array whose
//!    `size` appears in a constraint with no element/loop reference gets a
//!    32-bit size variable, and the constraints connected to it (through shared
//!    variables) form a small sub-problem solved first.
//! 2. **Expansion** ([`expand`]). With the sizes fixed, `foreach` unrolls,
//!    reductions fold, `size` is a constant and every element of a `rand`-
//!    active array is a fresh random variable (LRM 18.4: the elements are
//!    all randomized); an element of an inactive or non-`rand` array is a
//!    state variable holding its current value.
//!
//! The size sub-problem cannot see element constraints, so a size that
//! satisfies it may leave the full problem unsatisfiable: the caller blocks
//! that size tuple ([`solve_sizes`]'s `blocked`) and asks again, up to
//! [`MAX_SIZE_RETRIES`] times (an exhausted search is its own loud outcome,
//! not "unsatisfiable").
//!
//! # Limits (stated, not silent)
//!
//! * An array size is drawn uniformly from the sizes that satisfy its
//!   constraints and are at most [`SIZE_TIERS`]`[0]` (1024); if there are
//!   none, at most `SIZE_TIERS[1]` (65536). A constraint that only allows
//!   larger sizes is [`SizeError::TooLarge`], never a silent clamp.
//! * An element index must not depend on a random variable (the builder
//!   rejects those); after loop-variable substitution it must be a constant.
//! * An index outside `0..size` where the constraint is live is
//!   [`ExpandError::IndexOutOfRange`]; one in an arm a constant guard
//!   (`if (i < A.size-1) A[i] < A[i+1]`) never selects is not evaluated.

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;

/// Size tiers tried in order (module docs).
pub const SIZE_TIERS: [u64; 2] = [1024, 65536];

/// How many size tuples [`solve_sizes`] is asked to block before the search
/// is declared exhausted.
pub const MAX_SIZE_RETRIES: usize = 256;

/// One array of a program as seen by one `randomize()` call.
#[derive(Clone, Debug)]
pub struct ArrayState {
    /// Element type (the solver's 2-state view).
    pub elem_ty: Ty,
    /// Whether the array's elements (and size) are being randomized this call:
    /// a `rand` property whose `rand_mode` bit is set.
    pub active: bool,
    /// The current elements; `None` is an element with an X/Z bit.
    pub cur: Vec<Option<BigUint>>,
    /// The enum values an element must be one of (a `rand` enum array).
    pub domain: Option<Vec<BigUint>>,
}

/// Why [`expand`] failed.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ExpandError {
    /// An element index outside `0..size` in a live constraint.
    IndexOutOfRange {
        /// The array.
        array: ArrayId,
        /// The index (may be negative).
        index: i64,
        /// The array's size in this call.
        size: usize,
    },
    /// An index outside the elements of a fixed-size array held in a class
    /// property ([`ArrayNode::Pick`]) in a live constraint.
    LeafIndexOutOfRange {
        /// The index (may be negative).
        index: i64,
        /// The array's size.
        size: usize,
    },
    /// An element index that is not a constant once the loop variables are
    /// substituted.
    NonConstantIndex,
    /// A guard of an implication or `if` constraint is `ERROR` (LRM 18.5.12).
    GuardError,
    /// A state an index or constant guard reads is X/Z.
    StateUnknown(u32),
    /// A function called from a constraint failed on fixed arguments.
    Func(FuncError),
    /// A weight or bound of a `dist` inside a `foreach` over a dynamic array
    /// reads a random element once the loop is unrolled (LRM 18.5.3: weights
    /// and bounds are fixed values).
    DistRandomBound,
    /// The array-level IR is malformed (a builder bug).
    Malformed(&'static str),
}

/// The result of [`expand`].
#[derive(Clone, Debug, Default)]
pub struct Expanded {
    /// Scalar constraints over the field variables `0..n_vars0` and the element
    /// variables `n_vars0..`.
    pub conjuncts: Vec<Expr>,
    /// Element variable `n_vars0 + k` is `(array, index)` of `elem_vars[k]`.
    pub elem_vars: Vec<(ArrayId, usize)>,
    /// Types of the element variables.
    pub elem_tys: Vec<Ty>,
    /// State values appended after the caller's states (an element of an
    /// inactive or non-`rand` array).
    pub extra_states: Vec<Option<BigUint>>,
    /// The instances of the program's `dist` templates ([`Dist::subject`]),
    /// one per iteration that met one. The instance `k` has the selector
    /// `SELECTOR_BASE + n_program_dists + k`, and its variable ids are already
    /// the expansion's (element variables numbered after the field variables).
    pub extra_dists: Vec<Dist>,
}

struct Expander<'a> {
    arrays: &'a [ArrayState],
    sizes: &'a [usize],
    /// `elem_base[a]`: the variable id of element 0 of active array `a`.
    elem_base: Vec<u32>,
    asg: Assignment,
    synth: HashMap<(ArrayId, usize), u32>,
    env: Vec<(LoopId, Expr)>,
    funcs: &'a [Func],
    /// The program's dists; those with a `subject` are templates.
    dists: &'a [Dist],
    extra_dists: Vec<Dist>,
    /// The instance made for a template under a loop binding: the template's
    /// index, the bindings, and the instance's index in `extra_dists`. A
    /// template's selector occurs once per item code in its membership
    /// conjunct; all of them name one instance.
    instanced: Vec<InstancedDist>,
}

/// A template dist's instance: template index, loop bindings, index in `extra_dists`.
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 size_of(&self, a: ArrayId) -> Result<usize, ExpandError> {
        self.sizes
            .get(a as usize)
            .copied()
            .ok_or(ExpandError::Malformed("no size for an array"))
    }

    /// The constant value of the (var-free) expression `e`.
    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,
        })
    }

    /// The value of the index expression `index` (expanded here), signed per
    /// its type; it must be free of random variables.
    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)
        })
    }

    /// Element `idx` of array `a`: a random variable, or a state variable
    /// holding its current value.
    fn element(&mut self, a: ArrayId, idx: usize, ty: Ty) -> Result<Expr, ExpandError> {
        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)),
        }
    }

    /// The instance of the template `d` (a `dist` inside a `foreach`) for the
    /// current loop bindings: its subject, bounds and weights expanded. The
    /// items' `member` predicates are not needed (the membership conjunct is
    /// the expanded template's own) and are left true.
    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 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,
                });
                // The arguments are now scalar IR, so a call over fixed values
                // folds; one over random variables stays.
                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: *array,
                        index: i,
                        size,
                    });
                }
                self.element(*array, i as usize, *ty)
            }
            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(*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(*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(),
                        },
                    )
                }))
            }
        }
    }
}

/// `v == d0 || v == d1 || ...` for an enum element.
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))
}

/// Expands the array-level `conjuncts` for the array sizes `sizes` (module
/// docs). `n_vars0` is how many field variables the conjuncts already
/// number; element variables are numbered after them. `states` are the
/// caller's state values (an element of an inactive array becomes a state
/// appended after them).
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, &[], &[])
}

/// [`expand`] for a program with functions: a call whose arguments the
/// expansion makes fixed (an array size, an element of an inactive array, a
/// state) is evaluated into a constant; a call over random variables stays for
/// the staged solve ([`super::func_order`]).
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(),
    };
    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)
}

/// Why [`solve_sizes`] failed.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SizeError {
    /// No size assignment satisfies the size constraints.
    Unsatisfiable,
    /// Only sizes beyond [`SIZE_TIERS`] satisfy them.
    TooLarge,
    /// A guard in a size constraint is `ERROR`.
    GuardError,
    /// The size sub-problem failed otherwise.
    Solve(SolveError),
}

/// The array sizes chosen for one attempt.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SizeSolution {
    /// The size of every array (current size unless a constraint chose one).
    pub sizes: Vec<usize>,
    /// The arrays a constraint sized, in the order of [`Self::tuple`].
    pub sized: Vec<ArrayId>,
    /// The chosen sizes of [`Self::sized`] (what a retry blocks).
    pub tuple: Vec<usize>,
}

/// Whether `c` takes part in the size sub-problem: no element, loop or `dist`
/// selector reference.
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 { .. }
            // An index is only known once the node is expanded.
            | ArrayNode::Pick { .. }
            // A call over a size (the pre-pass folded every fixed one) is
            // only decidable once the arguments are known.
            | 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
}

/// Chooses the sizes of the `rand`-active arrays whose `size` the size-eligible
/// `conjuncts` constrain (module docs). `var_tys[v]` is the type of field
/// variable `v` (`conjuncts` are over `0..var_tys.len()`), `states` the state
/// values, `blocked` size tuples ruled out by earlier attempts (in the
/// `sized` order every call reproduces).
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;
    // 1. Which active arrays does an eligible conjunct size?
    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()
    };
    // 2. Substitute the size variables, then keep the conjuncts connected to one.
    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());
        }
    }
    // 3. Compact the variables.
    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));
    // Blocked tuples.
    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,
        }
    }
    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()
    }

    /// LRM 18.5.7.1: `foreach (A[j]) A[j] > 2*j` unrolls to one comparison per
    /// element, over one fresh variable each.
    #[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], &[]));
    }

    /// The ascending-sort example: `if (k < A.size-1) A[k] < A[k+1]` never
    /// reads `A[size]` because the constant guard drops that arm.
    #[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
            }
        );
    }

    /// A pick (a fixed-size array in a class property, XC-059) expands to the
    /// element its state index names, counted from the array's lowest index.
    #[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));
    }

    /// Only the chosen element is expanded, so a constant-false guard keeps an
    /// out-of-range pick unread; a live one is a loud error naming the index
    /// and the size (LRM 18.5.7.1).
    #[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 }
        );
        // A negative index (an `int` state of -1) is out of range too.
        let neg = [Some((BigUint::from(1u32) << 32) - 1u32)];
        assert_eq!(
            expand(&[body], &[], &[], 2, &neg).unwrap_err(),
            ExpandError::LeafIndexOutOfRange { index: -1, size: 2 }
        );
    }

    /// An inactive array is state: its elements read as the current values
    /// and its size as the current length.
    #[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:?}");
        }
    }

    /// LRM 18.4: `data.size == len` sizes the array to a drawn `len`.
    #[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
        );
        // Beyond the first tier but within the second.
        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
        );
    }

    /// A conjunct with an element reference is not part of the size problem,
    /// and a disconnected one is left out.
    #[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]);
    }
}