use core::fmt;
use std::borrow::Cow;
use std::collections::HashMap;
use std::sync::Arc;
use crate::error::Error;
use crate::expr::{Context, Expr};
use crate::facts::KnownBits;
use crate::hash::combine;
use crate::{BitVec, Width};
pub const MAX_ARGS: usize = 3;
pub const MAX_OUTPUTS: usize = 8;
pub const MAX_OPS: usize = 256;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ExtSig {
outputs: Vec<(Cow<'static, str>, Width)>,
}
impl ExtSig {
pub fn new(outputs: &[(&'static str, Width)]) -> ExtSig {
ExtSig {
outputs: outputs
.iter()
.map(|&(r, w)| (Cow::Borrowed(r), w))
.collect(),
}
}
pub fn outputs(&self) -> impl Iterator<Item = (&str, Width)> {
self.outputs.iter().map(|(r, w)| (&**r, *w))
}
pub fn len(&self) -> usize {
self.outputs.len()
}
pub fn is_empty(&self) -> bool {
self.outputs.is_empty()
}
pub fn width(&self, k: usize) -> Option<Width> {
self.outputs.get(k).map(|(_, w)| *w)
}
}
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct ExtTraits {
pub commutative: bool,
pub opaque: bool,
}
setters!(ExtTraits {
with_commutative: commutative: bool,
with_opaque: opaque: bool,
});
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum Invertible {
#[default]
No,
Injective,
Bijective,
}
pub trait ExtOp: Send + Sync + 'static {
fn name(&self) -> &str;
fn revision(&self) -> u32 {
1
}
fn signature(&self, args: &[Width]) -> Result<ExtSig, String>;
fn eval(&self, args: &[BitVec], out: &mut [BitVec]);
fn known_bits(&self, _args: &[KnownBits], _out: &mut [KnownBits]) {}
fn expand(&self, _cx: &mut Context, _args: &[Expr]) -> Option<Result<Vec<Expr>, Error>> {
None
}
fn smtlib(&self, _output: u8, _args: &[&str], _widths: &[Width]) -> Option<String> {
None
}
fn traits(&self) -> ExtTraits {
ExtTraits::default()
}
fn invertible(&self, _output: u8, _arg: u8, _args: &[KnownBits]) -> Invertible {
Invertible::No
}
fn invert(&self, _output: u8, _arg: u8, _args: &[BitVec], _value: &BitVec) -> Option<BitVec> {
None
}
}
#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub struct ExtArgs {
args: [Expr; MAX_ARGS],
len: u8,
}
impl ExtArgs {
pub(crate) fn new(args: &[Expr]) -> ExtArgs {
let mut a = [args[0]; MAX_ARGS];
a[..args.len()].copy_from_slice(args);
ExtArgs {
args: a,
len: args.len() as u8,
}
}
pub fn as_slice(&self) -> &[Expr] {
&self.args[..usize::from(self.len)]
}
}
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct ExtId {
index: u8,
tag: u32,
}
impl ExtId {
pub fn index(self) -> usize {
usize::from(self.index)
}
pub(crate) fn raw(self) -> u8 {
self.index
}
pub(crate) fn dangling() -> ExtId {
ExtId {
index: u8::MAX,
tag: 0,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum ExtError {
BadName(String),
Duplicate(String),
TooMany,
Contract(String, String),
}
impl fmt::Display for ExtError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
ExtError::BadName(n) => write!(f, "`{n}` is not a valid extension operation name"),
ExtError::Duplicate(n) => write!(f, "`{n}` is already registered"),
ExtError::TooMany => write!(f, "a registry holds at most {MAX_OPS} operations"),
ExtError::Contract(n, why) => write!(f, "`{n}` breaks its contract: {why}"),
}
}
}
impl std::error::Error for ExtError {}
#[derive(Clone, Default)]
pub struct Registry {
ops: Vec<Arc<dyn ExtOp>>,
by_name: HashMap<String, ExtId>,
hashes: Vec<u64>,
}
impl fmt::Debug for Registry {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_list()
.entries(
self.ops
.iter()
.map(|o| format!("{}#{}", o.name(), o.revision())),
)
.finish()
}
}
impl Registry {
pub fn builder() -> RegistryBuilder {
RegistryBuilder::default()
}
pub fn id(&self, name: &str) -> Option<ExtId> {
self.by_name.get(name).copied()
}
pub fn op(&self, id: ExtId) -> Option<&dyn ExtOp> {
(self.id_at(id.index)? == id).then(|| self.ops.get(id.index()).map(|o| &**o))?
}
pub(crate) fn id_at(&self, index: u8) -> Option<ExtId> {
let h = *self.hashes.get(usize::from(index))?;
Some(ExtId {
index,
tag: h as u32,
})
}
pub(crate) fn op_at(&self, index: u8) -> Option<&dyn ExtOp> {
self.ops.get(usize::from(index)).map(|o| &**o)
}
pub fn len(&self) -> usize {
self.ops.len()
}
pub fn is_empty(&self) -> bool {
self.ops.is_empty()
}
pub fn ids(&self) -> impl Iterator<Item = ExtId> + '_ {
(0..self.ops.len()).filter_map(|i| self.id_at(i as u8))
}
pub(crate) fn hash_at(&self, index: u8) -> u64 {
self.hashes.get(usize::from(index)).copied().unwrap_or(0)
}
}
#[derive(Default, Debug)]
pub struct RegistryBuilder {
reg: Registry,
}
impl RegistryBuilder {
pub fn register(self, op: impl ExtOp) -> Result<Self, ExtError> {
check(&op)?;
self.register_unchecked(op)
}
pub fn register_unchecked(mut self, op: impl ExtOp) -> Result<Self, ExtError> {
let name = op.name().to_string();
let good = name.chars().next().is_some_and(|c| c.is_ascii_alphabetic())
&& name
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '.')
&& !name.ends_with('.')
&& !name.contains("..");
if !good {
return Err(ExtError::BadName(name));
}
if self.reg.by_name.contains_key(&name) {
return Err(ExtError::Duplicate(name));
}
if self.reg.ops.len() >= MAX_OPS {
return Err(ExtError::TooMany);
}
let index = self.reg.ops.len() as u8;
let h = name
.bytes()
.fold(combine(0x6578_7400, u64::from(op.revision())), |h, b| {
combine(h, u64::from(b))
});
self.reg.hashes.push(h);
let id = ExtId {
index,
tag: h as u32,
};
self.reg.by_name.insert(name, id);
self.reg.ops.push(Arc::new(op));
Ok(self)
}
pub fn build(self) -> Registry {
self.reg
}
}
pub fn check(op: &dyn ExtOp) -> Result<(), ExtError> {
self_test(op).map_err(|why| ExtError::Contract(op.name().to_string(), why))
}
pub(crate) fn run_eval(op: &dyn ExtOp, args: &[BitVec]) -> Result<Vec<BitVec>, String> {
let widths: Vec<Width> = args.iter().map(|a| a.width()).collect();
let sig = op.signature(&widths)?;
if sig.is_empty() || sig.len() > MAX_OUTPUTS {
return Err(format!(
"{} outputs declared (1 to {MAX_OUTPUTS})",
sig.len()
));
}
let mut out: Vec<BitVec> = sig.outputs().map(|(_, w)| BitVec::zero(w)).collect();
op.eval(args, &mut out);
for (k, (v, (_, w))) in out.iter().zip(sig.outputs()).enumerate() {
if v.width() != w {
return Err(format!(
"output {k} has {} bits, declared {}",
v.width().bits(),
w.bits()
));
}
}
Ok(out)
}
pub(crate) fn run_known_bits(
op: &dyn ExtOp,
sig: &ExtSig,
args: &[KnownBits],
) -> Result<Vec<KnownBits>, String> {
let mut out: Vec<KnownBits> = sig.outputs().map(|(_, w)| KnownBits::unknown(w)).collect();
op.known_bits(args, &mut out);
for (k, (v, (_, w))) in out.iter().zip(sig.outputs()).enumerate() {
if v.width() != w {
return Err(format!(
"known bits of output {k} have {} bits, declared {}",
v.width().bits(),
w.bits()
));
}
}
Ok(out)
}
struct Rng(u64);
impl Rng {
fn next(&mut self) -> u64 {
self.0 = self.0.wrapping_add(0x9e37_79b9_7f4a_7c15);
let mut z = self.0;
z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
z ^ (z >> 31)
}
fn value(&mut self, w: Width) -> BitVec {
let limbs: Vec<u64> = (0..8)
.map(|_| match self.next() % 5 {
0 => 0,
1 => u64::MAX,
_ => self.next(),
})
.collect();
BitVec::wrapping_from_limbs(w, &limbs)
}
}
fn width_tuples() -> Vec<Vec<u16>> {
const W: [u16; 11] = [1, 3, 8, 13, 16, 32, 64, 65, 128, 256, 512];
let mut out: Vec<Vec<u16>> = W.iter().map(|&a| vec![a]).collect();
for &a in &W {
for &b in &W {
out.push(vec![a, b]);
}
}
for &a in &W {
for &b in &W {
for t in [vec![a, a, b], vec![a, b, a], vec![b, a, a]] {
if !out.contains(&t) {
out.push(t);
}
}
}
}
out
}
fn check_invertible(
op: &dyn ExtOp,
sig: &ExtSig,
widths: &[Width],
args: &[BitVec],
kb: &[KnownBits],
rng: &mut Rng,
) -> Result<(), String> {
use crate::facts::known::{bv_and, bv_not, bv_or};
for k in 0..sig.len() {
for a in 0..widths.len() {
let mut fixed = kb.to_vec();
fixed[a] = KnownBits::unknown(widths[a]);
let decl = op.invertible(k as u8, a as u8, &fixed);
if decl == Invertible::No {
continue;
}
if decl == Invertible::Bijective && sig.width(k) != Some(widths[a]) {
return Err(format!(
"declared bijective in argument {a}, but output {k} has another width"
));
}
let inst: Vec<BitVec> = args
.iter()
.zip(&fixed)
.map(|(v, f)| {
let r = rng.value(v.width());
bv_or(&bv_and(v, &f.known()), &bv_and(&r, &bv_not(&f.known())))
})
.collect();
let y = run_eval(op, &inst)?[k];
if op.invert(k as u8, a as u8, &inst, &y).as_ref() != Some(&inst[a]) {
return Err(format!(
"declared invertible, but invert does not recover argument {a} from output \
{k}"
));
}
if decl == Invertible::Bijective {
let target = rng.value(widths[a]);
let found = op.invert(k as u8, a as u8, &inst, &target);
let mut hit = inst.clone();
match found {
Some(v) if v.width() == widths[a] => hit[a] = v,
_ => {
return Err(format!(
"declared bijective, but invert finds no argument {a} for a value \
of output {k}"
));
}
}
if run_eval(op, &hit)?[k] != target {
return Err(format!(
"invert's argument {a} does not give the value asked of output {k}"
));
}
}
}
}
Ok(())
}
const TESTED_TUPLES: usize = 64;
fn self_test(op: &dyn ExtOp) -> Result<(), String> {
use crate::facts::known::{bv_and, bv_not};
let mut rng = Rng(0x5e1f_7e57 ^ u64::from(op.revision()));
let commutative = op.traits().commutative;
let mut accepted = 0;
for tuple in width_tuples() {
if accepted >= TESTED_TUPLES {
break;
}
let widths: Vec<Width> = tuple
.iter()
.map(|&w| Width::new(w).map_err(|e| e.to_string()))
.collect::<Result<_, _>>()?;
let Ok(sig) = op.signature(&widths) else {
continue;
};
if sig.is_empty() || sig.len() > MAX_OUTPUTS {
return Err(format!(
"{} outputs declared (1 to {MAX_OUTPUTS})",
sig.len()
));
}
accepted += 1;
let at = || format!(" at argument widths {tuple:?}");
for _ in 0..16 {
let args: Vec<BitVec> = widths.iter().map(|&w| rng.value(w)).collect();
let a = run_eval(op, &args).map_err(|e| e + &at())?;
if run_eval(op, &args)? != a {
return Err(format!("eval is not deterministic{}", at()));
}
if commutative && widths.len() >= 2 && widths[0] == widths[1] {
let mut swapped = args.clone();
swapped.swap(0, 1);
if run_eval(op, &swapped)? != a {
return Err(format!(
"declared commutative, but swapping the first two arguments changes \
the result{}",
at()
));
}
}
let masks: Vec<BitVec> = widths.iter().map(|&w| rng.value(w)).collect();
let kb: Vec<KnownBits> = args
.iter()
.zip(&masks)
.map(|(v, m)| {
KnownBits::new(bv_and(&bv_not(v), m), bv_and(v, m))
.unwrap_or(KnownBits::unknown(v.width()))
})
.collect();
let out = run_known_bits(op, &sig, &kb).map_err(|e| e + &at())?;
for _ in 0..4 {
let inst: Vec<BitVec> = args
.iter()
.zip(&masks)
.map(|(v, m)| {
let r = rng.value(v.width());
crate::facts::known::bv_or(&bv_and(v, m), &bv_and(&r, &bv_not(m)))
})
.collect();
for (k, (v, o)) in run_eval(op, &inst)?.iter().zip(&out).enumerate() {
if !o.contains(v) {
return Err(format!("known bits of output {k} exclude {v}{}", at()));
}
}
}
check_invertible(op, &sig, &widths, &args, &kb, &mut rng).map_err(|e| e + &at())?;
}
let mut cx = Context::new();
let syms: Vec<Expr> = widths
.iter()
.enumerate()
.map(|(i, &w)| cx.symbol(format!("a{i}").as_str(), w))
.collect::<Result<_, _>>()
.map_err(|e| e.to_string())?;
if let Some(ex) = op.expand(&mut cx, &syms) {
let ex = ex.map_err(|e| format!("expand failed: {e}{}", at()))?;
if ex.len() != sig.len() {
return Err(format!(
"expand gave {} outputs, not {}{}",
ex.len(),
sig.len(),
at()
));
}
for _ in 0..8 {
let args: Vec<BitVec> = widths.iter().map(|&w| rng.value(w)).collect();
let env: Vec<(crate::SymbolKey, BitVec)> = args
.iter()
.enumerate()
.map(|(i, v)| (crate::SymbolKey::from(format!("a{i}").as_str()), *v))
.collect();
let got = cx.eval(&ex, &env[..]).map_err(|e| e.to_string())?;
if got != run_eval(op, &args)? {
return Err(format!("expand disagrees with eval{}", at()));
}
}
}
}
if accepted == 0 {
return Err("no signature accepted at any tested argument widths".into());
}
Ok(())
}