vj-sim 0.3.0

//! The bit-blaster, generic over what a "bit" is (plan C3 and the SAT fallback).
//!
//! Every [`Expr`] node becomes a vector of bits, one per result bit, least
//! significant first; the whole constraint set becomes one bit that is true
//! exactly for the legal assignments. Each IR operator has one gadget here
//! (ripple-carry adder, comparator chain, barrel shifter, shift-add
//! multiplier, restoring divider, ...), written once against the [`Gates`]
//! trait. Two back ends implement it: [`super::bdd`] (a bit is a BDD, so the
//! result is a canonical function that can be counted and sampled exactly) and
//! [`super::sat`] (a bit is a CNF literal, so the result is a clause set a
//! CDCL solver can search when the BDD would be too large). Sharing the
//! gadgets is what keeps the two paths' notion of "legal" identical: the
//! unit tests of both check every gadget against the IR's own direct
//! evaluator ([`super::ir::eval`]).
//!
//! # Definedness (division by zero)
//!
//! [`ir::eval`](super::ir::eval) reports `x / 0` and `x % 0` as an error, and
//! `&&`, `||`, `->` and `?:` only evaluate the operand they need. The blaster
//! mirrors that with a *definedness* bit carried beside every value: `/` and
//! `%` contribute "divisor is non-zero", short-circuit operators gate the
//! right operand's definedness by the left operand's value, and the final
//! conjunct is `value & defined`. So an assignment on which `eval` would
//! return `DivideByZero` is not a model; it is "not legal", the same way the
//! LRM treats a constraint expression it cannot evaluate as unsatisfied.

use super::bdd::BddBuildError;
use super::ir::{BinOp, Expr, UnOp, VarId};
use num_bigint::BigUint;

/// The single-bit operations the gadgets are built from.
///
/// Every operation may fail (a back end has a size budget); `during` names the
/// IR operation being blasted, for the error.
pub(crate) trait Gates {
    /// One bit of a blasted value.
    type Bit: Clone;

    /// The constant bit.
    fn constant(&self, v: bool) -> Self::Bit;
    /// The value of `a` when it is a constant.
    fn as_const(&self, a: &Self::Bit) -> Option<bool>;
    /// Bit `i` (LSB = 0) of random variable `var`.
    fn input(&self, var: VarId, i: usize) -> Self::Bit;
    /// `!a`.
    fn not(&self, a: &Self::Bit) -> Self::Bit;
    /// `a & b`.
    fn and(
        &self,
        during: &'static str,
        a: &Self::Bit,
        b: &Self::Bit,
    ) -> Result<Self::Bit, BddBuildError>;
    /// `a | b`.
    fn or(
        &self,
        during: &'static str,
        a: &Self::Bit,
        b: &Self::Bit,
    ) -> Result<Self::Bit, BddBuildError>;
    /// `a ^ b`.
    fn xor(
        &self,
        during: &'static str,
        a: &Self::Bit,
        b: &Self::Bit,
    ) -> Result<Self::Bit, BddBuildError>;
    /// `s ? a : b`.
    fn mux(
        &self,
        during: &'static str,
        s: &Self::Bit,
        a: &Self::Bit,
        b: &Self::Bit,
    ) -> Result<Self::Bit, BddBuildError>;
}

/// A blasted value plus the condition under which it is defined (`None` =
/// always defined).
pub(crate) struct Val<B> {
    /// The value's bits, least significant first.
    pub bits: Vec<B>,
    /// When the value exists (no division by zero was needed to compute it).
    pub def: Option<B>,
}

/// The expression walker over a [`Gates`] back end.
pub(crate) struct Blaster<'a, G: Gates> {
    pub(crate) g: &'a G,
    /// The declared width of each random variable, indexed by [`VarId`].
    pub(crate) widths: &'a [usize],
    pub(crate) states: &'a [Option<BigUint>],
}

type Bits<G> = Vec<<G as Gates>::Bit>;

impl<G: Gates> Blaster<'_, G> {
    // ---- single-bit helpers ----

    fn iff(&self, d: &'static str, a: &G::Bit, b: &G::Bit) -> Result<G::Bit, BddBuildError> {
        Ok(self.g.not(&self.g.xor(d, a, b)?))
    }

    // ---- vector helpers ----

    fn zeros(&self, w: usize) -> Bits<G> {
        vec![self.g.constant(false); w]
    }

    fn const_bits(&self, v: &BigUint, w: u32) -> Bits<G> {
        (0..w)
            .map(|i| self.g.constant(v.bit(u64::from(i))))
            .collect()
    }

    /// "Some bit is set".
    pub(crate) fn truth(&self, a: &[G::Bit], d: &'static str) -> Result<G::Bit, BddBuildError> {
        let mut acc = self.g.constant(false);
        for x in a {
            acc = self.g.or(d, &acc, x)?;
        }
        Ok(acc)
    }

    fn map2(
        &self,
        d: &'static str,
        a: &[G::Bit],
        b: &[G::Bit],
        f: impl Fn(&'static str, &G::Bit, &G::Bit) -> Result<G::Bit, BddBuildError>,
    ) -> Result<Bits<G>, BddBuildError> {
        a.iter().zip(b).map(|(x, y)| f(d, x, y)).collect()
    }

    /// Ripple-carry `a + b + cin`; the carry-out is computed only if asked.
    fn add_c(
        &self,
        d: &'static str,
        a: &[G::Bit],
        b: &[G::Bit],
        cin: G::Bit,
        want_cout: bool,
    ) -> Result<(Bits<G>, G::Bit), BddBuildError> {
        let mut carry = cin;
        let mut out = Vec::with_capacity(a.len());
        for i in 0..a.len() {
            let x = self.g.xor(d, &a[i], &b[i])?;
            out.push(self.g.xor(d, &x, &carry)?);
            if want_cout || i + 1 < a.len() {
                let gen_ = self.g.and(d, &a[i], &b[i])?;
                let p = self.g.and(d, &carry, &x)?;
                carry = self.g.or(d, &gen_, &p)?;
            }
        }
        Ok((out, carry))
    }

    fn add(&self, d: &'static str, a: &[G::Bit], b: &[G::Bit]) -> Result<Bits<G>, BddBuildError> {
        Ok(self.add_c(d, a, b, self.g.constant(false), false)?.0)
    }

    fn sub(&self, d: &'static str, a: &[G::Bit], b: &[G::Bit]) -> Result<Bits<G>, BddBuildError> {
        let nb: Bits<G> = b.iter().map(|x| self.g.not(x)).collect();
        Ok(self.add_c(d, a, &nb, self.g.constant(true), false)?.0)
    }

    fn neg(&self, d: &'static str, a: &[G::Bit]) -> Result<Bits<G>, BddBuildError> {
        self.sub(d, &self.zeros(a.len()), a)
    }

    /// `a == b`.
    fn eq(&self, d: &'static str, a: &[G::Bit], b: &[G::Bit]) -> Result<G::Bit, BddBuildError> {
        let mut acc = self.g.constant(true);
        // Most significant bit first: a BDD's variables follow the bit order,
        // so each conjunct then lands above the chain built so far and the
        // `and` is constant-time, where LSB first re-walks the whole chain
        // (quadratic in the width, and the bulk of a debug randomize call).
        for (x, y) in a.iter().zip(b).rev() {
            let e = self.iff(d, x, y)?;
            acc = self.g.and(d, &acc, &e)?;
        }
        Ok(acc)
    }

    /// `a < b`, unsigned or signed (two's complement).
    fn lt(
        &self,
        d: &'static str,
        a: &[G::Bit],
        b: &[G::Bit],
        signed: bool,
    ) -> Result<G::Bit, BddBuildError> {
        let w = a.len();
        let mut lt = self.g.constant(false);
        for i in 0..w {
            // Signed order is unsigned order with the sign bits swapped.
            let (x, y) = if signed && i + 1 == w {
                (b[i].clone(), a[i].clone())
            } else {
                (a[i].clone(), b[i].clone())
            };
            let strictly = self.g.and(d, &self.g.not(&x), &y)?;
            let same = self.iff(d, &x, &y)?;
            let carried = self.g.and(d, &same, &lt)?;
            lt = self.g.or(d, &strictly, &carried)?;
        }
        Ok(lt)
    }

    /// `cond ? neg(x) : x`, bitwise.
    fn cond_neg(
        &self,
        d: &'static str,
        cond: &G::Bit,
        x: &[G::Bit],
    ) -> Result<Bits<G>, BddBuildError> {
        let n = self.neg(d, x)?;
        x.iter()
            .zip(&n)
            .map(|(orig, neg)| self.g.mux(d, cond, neg, orig))
            .collect()
    }

    /// Barrel shifter. The amount is unsigned; at or past the width
    /// everything is shifted out (`Ashr` of a signed value fills with the
    /// sign), exactly like the evaluator.
    fn shift(
        &self,
        x: &[G::Bit],
        amt: &[G::Bit],
        op: BinOp,
        signed: bool,
    ) -> Result<Bits<G>, BddBuildError> {
        let d = "shift";
        let w = x.len();
        let fill = if op == BinOp::Ashr && signed {
            x[w - 1].clone()
        } else {
            self.g.constant(false)
        };
        let nstages = (0..usize::BITS as usize)
            .take_while(|k| (1usize << k) < w)
            .count();
        let mut cur: Bits<G> = x.to_vec();
        for (k, sel) in amt.iter().enumerate().take(nstages) {
            if self.g.as_const(sel) == Some(false) {
                continue;
            }
            let s = 1usize << k;
            let mut next = Vec::with_capacity(w);
            for i in 0..w {
                let moved = if op == BinOp::Shl {
                    if i >= s {
                        cur[i - s].clone()
                    } else {
                        self.g.constant(false)
                    }
                } else if i + s < w {
                    cur[i + s].clone()
                } else {
                    fill.clone()
                };
                next.push(self.g.mux(d, sel, &moved, &cur[i])?);
            }
            cur = next;
        }
        // Any set amount bit worth at least `w` shifts everything out.
        let big = self.truth(amt.get(nstages..).unwrap_or(&[]), d)?;
        if self.g.as_const(&big) != Some(false) {
            cur = cur
                .iter()
                .map(|c| self.g.mux(d, &big, &fill, c))
                .collect::<Result<_, _>>()?;
        }
        Ok(cur)
    }

    /// Shift-add multiplier, truncated to the operand width.
    fn mul(&self, a: &[G::Bit], b: &[G::Bit]) -> Result<Bits<G>, BddBuildError> {
        let d = "multiplication";
        let w = a.len();
        let mut acc = self.zeros(w);
        for (j, bj) in b.iter().enumerate() {
            if self.g.as_const(bj) == Some(false) {
                continue;
            }
            // Partial product `(a & bj) << j`, only the bits from `j` up.
            let part: Bits<G> = (0..w - j)
                .map(|i| self.g.and(d, &a[i], bj))
                .collect::<Result<_, _>>()?;
            let (sum, _) = self.add_c(d, &acc[j..], &part, self.g.constant(false), false)?;
            acc[j..].clone_from_slice(&sum);
        }
        Ok(acc)
    }

    /// Unsigned restoring divider; `(quotient, remainder)`. Division by zero
    /// yields junk here and is excluded by the caller's definedness gate.
    fn divmod_unsigned(
        &self,
        a: &[G::Bit],
        dv: &[G::Bit],
    ) -> Result<(Bits<G>, Bits<G>), BddBuildError> {
        let d = "division";
        let w = a.len();
        let mut dext: Bits<G> = dv.to_vec();
        dext.push(self.g.constant(false));
        let mut q = self.zeros(w);
        let mut r = self.zeros(w);
        for i in (0..w).rev() {
            // t = (r << 1) | a[i], w + 1 bits.
            let mut t: Bits<G> = Vec::with_capacity(w + 1);
            t.push(a[i].clone());
            t.extend(r.iter().cloned());
            let ge = self.g.not(&self.lt(d, &t, &dext, false)?);
            let diff = self.sub(d, &t, &dext)?;
            r = (0..w)
                .map(|k| self.g.mux(d, &ge, &diff[k], &t[k]))
                .collect::<Result<_, _>>()?;
            q[i] = ge;
        }
        Ok((q, r))
    }

    fn divmod(
        &self,
        a: &[G::Bit],
        dv: &[G::Bit],
        signed: bool,
    ) -> Result<(Bits<G>, Bits<G>), BddBuildError> {
        if !signed {
            return self.divmod_unsigned(a, dv);
        }
        let d = "signed division";
        let w = a.len();
        let (sa, sd) = (a[w - 1].clone(), dv[w - 1].clone());
        let abs_a = self.cond_neg(d, &sa, a)?;
        let abs_d = self.cond_neg(d, &sd, dv)?;
        let (q, r) = self.divmod_unsigned(&abs_a, &abs_d)?;
        let flip = self.g.xor(d, &sa, &sd)?;
        Ok((self.cond_neg(d, &flip, &q)?, self.cond_neg(d, &sa, &r)?))
    }

    // ---- definedness ----

    fn def_and(
        &self,
        a: Option<G::Bit>,
        b: Option<G::Bit>,
    ) -> Result<Option<G::Bit>, BddBuildError> {
        Ok(match (a, b) {
            (None, x) | (x, None) => x,
            (Some(x), Some(y)) => Some(self.g.and("definedness", &x, &y)?),
        })
    }

    // ---- the expression walker ----

    /// The conjunct `c` as one bit: its truthiness gated by its definedness.
    pub(crate) fn conjunct(&self, c: &Expr) -> Result<G::Bit, BddBuildError> {
        let v = self.val(c)?;
        let t = self.truth(&v.bits, "conjunct")?;
        match &v.def {
            Some(d) => self.g.and("conjunct definedness", &t, d),
            None => Ok(t),
        }
    }

    pub(crate) fn val(&self, e: &Expr) -> Result<Val<G::Bit>, BddBuildError> {
        let plain = |bits: Bits<G>| Val { bits, def: None };
        match e {
            Expr::Const(t, v) => Ok(plain(self.const_bits(v, t.width))),
            Expr::Var(id, t) => {
                let width = *self
                    .widths
                    .get(*id as usize)
                    .ok_or(BddBuildError::UnboundVar(*id))?;
                if width != t.width as usize {
                    return Err(BddBuildError::VarWidthMismatch {
                        var: *id,
                        declared: u32::try_from(width).expect("a width"),
                        used: t.width,
                    });
                }
                Ok(plain((0..width).map(|i| self.g.input(*id, i)).collect()))
            }
            Expr::State(id, t) => match self.states.get(*id as usize) {
                None => Err(BddBuildError::UnboundState(*id)),
                Some(None) => Err(BddBuildError::StateUnknown(*id)),
                Some(Some(v)) => Ok(plain(self.const_bits(v, t.width))),
            },
            Expr::Unary(op, x, _) => {
                let a = self.val(x)?;
                let d = "unary operator";
                let bits = match op {
                    UnOp::BitNot => a.bits.iter().map(|b| self.g.not(b)).collect(),
                    UnOp::Neg => self.neg(d, &a.bits)?,
                    UnOp::LogNot | UnOp::RedNor => vec![self.g.not(&self.truth(&a.bits, d)?)],
                    UnOp::RedOr => vec![self.truth(&a.bits, d)?],
                    UnOp::RedAnd | UnOp::RedNand => {
                        let mut acc = self.g.constant(true);
                        for b in &a.bits {
                            acc = self.g.and(d, &acc, b)?;
                        }
                        vec![if *op == UnOp::RedAnd {
                            acc
                        } else {
                            self.g.not(&acc)
                        }]
                    }
                    UnOp::RedXor | UnOp::RedXnor => {
                        let mut acc = self.g.constant(false);
                        for b in &a.bits {
                            acc = self.g.xor(d, &acc, b)?;
                        }
                        vec![if *op == UnOp::RedXor {
                            acc
                        } else {
                            self.g.not(&acc)
                        }]
                    }
                };
                Ok(Val { bits, def: a.def })
            }
            Expr::Binary(op @ (BinOp::LogAnd | BinOp::LogOr | BinOp::LogImpl), l, r, _) => {
                let d = "logical operator";
                let (a, b) = (self.val(l)?, self.val(r)?);
                let ta = self.truth(&a.bits, d)?;
                let tb = self.truth(&b.bits, d)?;
                // `decided` is true when the left operand alone fixes the
                // result; then the right operand is never evaluated.
                let (value, decided) = match op {
                    BinOp::LogAnd => (self.g.and(d, &ta, &tb)?, self.g.not(&ta)),
                    BinOp::LogOr => (self.g.or(d, &ta, &tb)?, ta.clone()),
                    _ => (self.g.or(d, &self.g.not(&ta), &tb)?, self.g.not(&ta)),
                };
                let def = match b.def {
                    None => a.def,
                    Some(db) => {
                        // Defined iff the left is, and (decided or right defined).
                        let right_ok = self.g.or(d, &decided, &db)?;
                        self.def_and(a.def, Some(right_ok))?
                    }
                };
                Ok(Val {
                    bits: vec![value],
                    def,
                })
            }
            Expr::Binary(op, l, r, t) => {
                let (a, b) = (self.val(l)?, self.val(r)?);
                let d = "binary operator";
                let mut def = self.def_and(a.def, b.def)?;
                let signed = l.ty().signed;
                let bits: Bits<G> = match op {
                    BinOp::Add => self.add(d, &a.bits, &b.bits)?,
                    BinOp::Sub => self.sub(d, &a.bits, &b.bits)?,
                    BinOp::Mul => self.mul(&a.bits, &b.bits)?,
                    BinOp::Div | BinOp::Mod => {
                        let nonzero = self.truth(&b.bits, d)?;
                        def = self.def_and(def, Some(nonzero))?;
                        let (q, r) = self.divmod(&a.bits, &b.bits, signed)?;
                        if *op == BinOp::Div { q } else { r }
                    }
                    BinOp::And => self.map2(d, &a.bits, &b.bits, |d, x, y| self.g.and(d, x, y))?,
                    BinOp::Or => self.map2(d, &a.bits, &b.bits, |d, x, y| self.g.or(d, x, y))?,
                    BinOp::Xor => self.map2(d, &a.bits, &b.bits, |d, x, y| self.g.xor(d, x, y))?,
                    BinOp::Xnor => self.map2(d, &a.bits, &b.bits, |d, x, y| self.iff(d, x, y))?,
                    BinOp::Shl | BinOp::Shr | BinOp::Ashr => {
                        self.shift(&a.bits, &b.bits, *op, signed)?
                    }
                    BinOp::Eq => vec![self.eq(d, &a.bits, &b.bits)?],
                    BinOp::Ne => vec![self.g.not(&self.eq(d, &a.bits, &b.bits)?)],
                    BinOp::Lt => vec![self.lt(d, &a.bits, &b.bits, signed)?],
                    BinOp::Gt => vec![self.lt(d, &b.bits, &a.bits, signed)?],
                    BinOp::Le => vec![self.g.not(&self.lt(d, &b.bits, &a.bits, signed)?)],
                    BinOp::Ge => vec![self.g.not(&self.lt(d, &a.bits, &b.bits, signed)?)],
                    BinOp::LogEquiv => {
                        let ta = self.truth(&a.bits, d)?;
                        let tb = self.truth(&b.bits, d)?;
                        vec![self.iff(d, &ta, &tb)?]
                    }
                    BinOp::LogAnd | BinOp::LogOr | BinOp::LogImpl => {
                        unreachable!("short-circuit operators have their own arm")
                    }
                };
                debug_assert_eq!(bits.len(), t.width as usize);
                Ok(Val { bits, def })
            }
            Expr::Ite(c, x, y, _) | Expr::Guard(c, x, y, _) => {
                let d = "conditional";
                let (cv, a, b) = (self.val(c)?, self.val(x)?, self.val(y)?);
                let sel = self.truth(&cv.bits, d)?;
                let bits = (0..a.bits.len())
                    .map(|i| self.g.mux(d, &sel, &a.bits[i], &b.bits[i]))
                    .collect::<Result<_, _>>()?;
                // Only the taken arm is evaluated.
                let arms = match (a.def, b.def) {
                    (None, None) => None,
                    (da, db) => {
                        let t = da.unwrap_or_else(|| self.g.constant(true));
                        let f = db.unwrap_or_else(|| self.g.constant(true));
                        Some(self.g.mux(d, &sel, &t, &f)?)
                    }
                };
                Ok(Val {
                    bits,
                    def: self.def_and(cv.def, arms)?,
                })
            }
            Expr::Resize(x, t) => {
                let a = self.val(x)?;
                let from = x.ty();
                let w = t.width as usize;
                let mut bits = a.bits;
                if w <= bits.len() {
                    bits.truncate(w);
                } else {
                    let ext = if from.signed {
                        bits[bits.len() - 1].clone()
                    } else {
                        self.g.constant(false)
                    };
                    bits.resize(w, ext);
                }
                Ok(Val { bits, def: a.def })
            }
            Expr::Concat(parts, _) => {
                let mut bits: Bits<G> = Vec::new();
                let mut def = None;
                // First element is most significant: emit the last first.
                for p in parts.iter().rev() {
                    let v = self.val(p)?;
                    bits.extend(v.bits);
                    def = self.def_and(def, v.def)?;
                }
                Ok(Val { bits, def })
            }
            Expr::Slice(x, lo, t) => {
                let a = self.val(x)?;
                let lo = *lo as usize;
                Ok(Val {
                    bits: a.bits[lo..lo + t.width as usize].to_vec(),
                    def: a.def,
                })
            }
            Expr::AllDistinct(members) => {
                let d = "uniqueness group";
                let vals = members
                    .iter()
                    .map(|m| self.val(m))
                    .collect::<Result<Vec<_>, _>>()?;
                let mut def = None;
                for v in &vals {
                    def = self.def_and(def, v.def.clone())?;
                }
                // No two members equal: the conjunction of every pair's `!=`.
                let mut all = self.g.constant(true);
                for (i, a) in vals.iter().enumerate() {
                    for b in &vals[i + 1..] {
                        let ne = self.g.not(&self.eq(d, &a.bits, &b.bits)?);
                        all = self.g.and(d, &all, &ne)?;
                    }
                }
                Ok(Val {
                    bits: vec![all],
                    def,
                })
            }
            Expr::Array(_) => Err(BddBuildError::Unexpanded),
        }
    }
}