use core::fmt;
use core::hash::Hash;
use std::collections::HashMap;
use std::sync::Mutex;
use crate::error::{Error, WidthError};
use crate::expr::{Context, Expr, View};
use crate::facts::KnownBits;
use crate::{ParseOptions, SymbolKey, Width};
pub trait Semantics {
type Value: Copy + Eq + Hash;
fn result(&self) -> Option<Self::Value>;
fn lower(&self, lw: &mut Lowering<'_, Self::Value>) -> Result<Expr, Error>;
}
pub trait Raise {
type Value: Copy;
fn emit(
&mut self,
cx: &Context,
e: Expr,
node: &View,
operands: &[Self::Value],
) -> Result<Self::Value, Error>;
}
pub struct Lowering<'cx, V> {
cx: &'cx mut Context,
exprs: HashMap<V, Expr>,
inputs: Vec<(V, Expr)>,
owners: HashMap<Expr, V>,
}
impl<V> fmt::Debug for Lowering<'_, V> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Lowering")
.field("values", &self.exprs.len())
.field("inputs", &self.inputs.len())
.finish_non_exhaustive()
}
}
pub struct Lowered<V> {
exprs: HashMap<V, Expr>,
inputs: Vec<(V, Expr)>,
owners: HashMap<Expr, V>,
}
impl<V> fmt::Debug for Lowered<V> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Lowered")
.field("values", &self.exprs.len())
.field("inputs", &self.inputs.len())
.finish_non_exhaustive()
}
}
impl<V> Default for Lowered<V> {
fn default() -> Self {
Lowered {
exprs: HashMap::new(),
inputs: Vec::new(),
owners: HashMap::new(),
}
}
}
impl<V: Copy + Eq + Hash> Lowered<V> {
pub fn get(&self, v: V) -> Option<Expr> {
self.exprs.get(&v).copied()
}
pub fn inputs(&self) -> &[(V, Expr)] {
&self.inputs
}
pub fn owner(&self, e: Expr) -> Option<V> {
self.owners.get(&e).copied()
}
pub fn attach(self, cx: &mut Context) -> Lowering<'_, V> {
Lowering {
cx,
exprs: self.exprs,
inputs: self.inputs,
owners: self.owners,
}
}
}
impl<'cx, V: Copy + Eq + Hash> Lowering<'cx, V> {
pub fn new(cx: &'cx mut Context) -> Self {
Lowered::default().attach(cx)
}
pub fn detach(self) -> Lowered<V> {
Lowered {
exprs: self.exprs,
inputs: self.inputs,
owners: self.owners,
}
}
pub fn context(&mut self) -> &mut Context {
self.cx
}
pub fn value(&mut self, v: V, width: Width) -> Result<Expr, Error> {
match self.exprs.get(&v) {
Some(&e) => {
let have = self.cx.width(e)?;
if have != width {
return Err(WidthError::Mismatch {
left: have.bits(),
right: width.bits(),
}
.into());
}
Ok(e)
}
None => self.input(v, width, None),
}
}
pub fn input(&mut self, v: V, width: Width, known: Option<KnownBits>) -> Result<Expr, Error> {
if let Some(&e) = self.exprs.get(&v) {
return Ok(e);
}
let s = self.cx.fresh_symbol(width)?;
self.bind_input(v, s, known)
}
pub fn input_named(
&mut self,
v: V,
key: impl Into<SymbolKey>,
width: Width,
known: Option<KnownBits>,
) -> Result<Expr, Error> {
if let Some(&e) = self.exprs.get(&v) {
return Ok(e);
}
let s = self.cx.symbol(key, width)?;
self.bind_input(v, s, known)
}
fn bind_input(&mut self, v: V, s: Expr, known: Option<KnownBits>) -> Result<Expr, Error> {
if let Some(k) = known {
self.cx.declare_known(s, k)?;
}
self.exprs.insert(v, s);
self.inputs.push((v, s));
self.owners.entry(s).or_insert(v);
Ok(s)
}
pub fn define(&mut self, v: V, e: Expr) -> Result<(), Error> {
self.cx.id(e)?;
self.exprs.insert(v, e);
self.owners.entry(e).or_insert(v);
Ok(())
}
pub fn lower<S: Semantics<Value = V> + ?Sized>(
&mut self,
inst: &S,
) -> Result<Option<Expr>, Error> {
let Some(v) = inst.result() else {
return Ok(None);
};
let e = inst.lower(self)?;
self.define(v, e)?;
Ok(Some(e))
}
pub fn lower_all<'i, S, I>(&mut self, insts: I) -> Result<(), Error>
where
S: Semantics<Value = V> + 'i,
I: IntoIterator<Item = &'i S>,
{
for inst in insts {
self.lower(inst)?;
}
Ok(())
}
pub fn get(&self, v: V) -> Option<Expr> {
self.exprs.get(&v).copied()
}
pub fn inputs(&self) -> &[(V, Expr)] {
&self.inputs
}
pub fn owner(&self, e: Expr) -> Option<V> {
self.owners.get(&e).copied()
}
pub fn raise<R: Raise<Value = V> + ?Sized>(&mut self, e: Expr, r: &mut R) -> Result<V, Error> {
let mut ops: Vec<V> = Vec::new();
for n in self.unraised(e)? {
let view = self.cx.view(n)?;
ops.clear();
for c in self.cx.children(n)? {
ops.push(self.owners[&c]);
}
let v = r.emit(self.cx, n, &view, &ops)?;
self.emitted(n, v)?;
}
Ok(self.owners[&e])
}
pub fn unraised(&mut self, e: Expr) -> Result<Vec<Expr>, Error> {
let i = self.cx.id(e)?;
if self.owner(e).is_some() {
return Ok(Vec::new());
}
let owners = &self.owners;
let order: Vec<Expr> = self
.cx
.post_order_ids_pruned(&[i], |n| owners.contains_key(&n))
.into_iter()
.map(|n| self.cx.handle(n))
.collect();
for &n in &order {
if matches!(self.cx.view(n)?, View::Sym(_)) {
return Err(Error::Unsupported(format!(
"symbol {} stands for no host value",
self.cx.display(n)
)));
}
}
Ok(order)
}
pub fn emitted(&mut self, e: Expr, v: V) -> Result<(), Error> {
self.cx.id(e)?;
self.exprs.entry(v).or_insert(e);
self.owners.insert(e, v);
Ok(())
}
}
pub struct Template {
text: String,
params: Vec<String>,
compiled: Mutex<Vec<(Vec<Width>, Compiled)>>,
}
struct Compiled {
cx: Context,
root: Expr,
params: Vec<u32>,
}
impl fmt::Debug for Template {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Template")
.field("text", &self.text)
.field("params", &self.params)
.finish_non_exhaustive()
}
}
impl Template {
pub fn new(text: &str, params: &[&str]) -> Result<Template, Error> {
for (k, p) in params.iter().enumerate() {
if p.is_empty() || params[..k].contains(p) {
return Err(Error::Unsupported(format!(
"template parameter `{p}` is empty or repeated"
)));
}
}
Ok(Template {
text: text.to_string(),
params: params.iter().map(|p| (*p).to_string()).collect(),
compiled: Mutex::new(Vec::new()),
})
}
pub fn check(&self, widths: &[Width]) -> Result<(), Error> {
if widths.len() != self.params.len() {
return Err(self.arity(widths.len()));
}
self.with_compiled(widths, |_| Ok(()))
}
fn arity(&self, got: usize) -> Error {
Error::Unsupported(format!(
"the template takes {} operands, not {got}",
self.params.len()
))
}
fn with_compiled<T>(
&self,
widths: &[Width],
f: impl FnOnce(&Compiled) -> Result<T, Error>,
) -> Result<T, Error> {
let mut compiled = self
.compiled
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let k = match compiled.iter().position(|(w, _)| w == widths) {
Some(k) => k,
None => {
let c = self.compile(widths)?;
compiled.push((widths.to_vec(), c));
compiled.len() - 1
}
};
f(&compiled[k].1)
}
fn compile(&self, widths: &[Width]) -> Result<Compiled, Error> {
let mut cx = Context::new();
let mut params = Vec::with_capacity(self.params.len());
for (p, &w) in self.params.iter().zip(widths) {
let s = cx.symbol(p.as_str(), w)?;
params.push(cx.id(s)?);
}
let default = widths.iter().copied().max().unwrap_or(Width::W64);
let root = cx.parse(&self.text, &ParseOptions::width(default))?;
Ok(Compiled { cx, root, params })
}
pub fn instantiate(&self, cx: &mut Context, args: &[Expr]) -> Result<Expr, Error> {
if args.len() != self.params.len() {
return Err(self.arity(args.len()));
}
let widths: Vec<Width> = args
.iter()
.map(|&a| cx.width(a))
.collect::<Result<_, _>>()?;
self.with_compiled(&widths, |c| {
let pre: Vec<(u32, u32)> = c
.params
.iter()
.zip(args)
.map(|(&p, &a)| Ok((p, cx.id(a)?)))
.collect::<Result<_, Error>>()?;
Ok(cx.import_mapped(&c.cx, &[c.root], &pre)?[0])
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{BinOp, BitVec, CmpOp, FnEnv};
struct Bin(u32, BinOp, u32, u32);
impl Semantics for Bin {
type Value = u32;
fn result(&self) -> Option<u32> {
Some(self.0)
}
fn lower(&self, lw: &mut Lowering<'_, u32>) -> Result<Expr, Error> {
let (a, b) = (lw.value(self.2, Width::W8)?, lw.value(self.3, Width::W8)?);
lw.context().bin(self.1, a, b)
}
}
struct Rec(Vec<String>);
impl Raise for Rec {
type Value = u32;
fn emit(&mut self, _: &Context, _: Expr, node: &View, ops: &[u32]) -> Result<u32, Error> {
let v = 100 + self.0.len() as u32;
let what = match node {
View::Const(c) => format!("const {c}"),
View::Bin(op, ..) => format!("{op:?} {ops:?}"),
other => format!("{other:?}"),
};
self.0.push(format!("%{v} = {what}"));
Ok(v)
}
}
#[test]
fn lowering_maps_values_and_inputs() {
let mut cx = Context::new();
let mut lw = Lowering::new(&mut cx);
lw.lower_all(&[Bin(2, BinOp::Add, 0, 1), Bin(3, BinOp::Xor, 2, 0)])
.unwrap();
let ins: Vec<u32> = lw.inputs().iter().map(|&(v, _)| v).collect();
assert_eq!(ins, [0, 1]);
assert!(matches!(lw.value(2, Width::W16), Err(Error::Width(_))));
let k = KnownBits::constant(&BitVec::from_u64(Width::W8, 7).unwrap());
let s = lw.input(9, Width::W8, Some(k)).unwrap();
assert_eq!(lw.context().declared_known(s).unwrap(), Some(k));
assert_eq!(lw.owner(s), Some(9));
}
#[test]
fn raising_emits_only_new_nodes_operands_first() {
let mut cx = Context::new();
let mut lw = Lowering::new(&mut cx);
lw.lower_all(&[Bin(2, BinOp::Add, 0, 1), Bin(3, BinOp::Mul, 2, 1)])
.unwrap();
let (v2, v0) = (lw.get(2).unwrap(), lw.get(0).unwrap());
let cx = lw.context();
let five = cx.constant_u64(Width::W8, 5).unwrap();
let mask = cx.bin(BinOp::And, v0, five).unwrap();
let e = cx.bin(BinOp::Xor, v2, mask).unwrap();
let mut rec = Rec(Vec::new());
let v = lw.raise(e, &mut rec).unwrap();
assert_eq!(v, 102);
assert_eq!(
rec.0,
[
"%100 = const 0x5:8",
"%101 = And [0, 100]",
"%102 = Xor [2, 101]"
]
);
assert_eq!(lw.raise(e, &mut rec).unwrap(), 102);
assert_eq!(rec.0.len(), 3);
let stray = lw.context().symbol("stray", Width::W8).unwrap();
assert!(lw.raise(stray, &mut rec).is_err());
}
#[test]
fn a_detached_lowering_resumes() {
let mut cx = Context::new();
let mut lw = Lowering::new(&mut cx);
lw.lower(&Bin(2, BinOp::Add, 0, 1)).unwrap();
let kept = lw.detach();
let one = cx.constant_u64(Width::W8, 1).unwrap();
let mut lw = kept.attach(&mut cx);
let v2 = lw.get(2).unwrap();
let e = lw.context().bin(BinOp::Add, v2, one).unwrap();
lw.define(3, e).unwrap();
assert_eq!(lw.inputs().len(), 2);
assert_eq!(lw.owner(v2), Some(2));
let mut rec = Rec(Vec::new());
assert_eq!(lw.raise(e, &mut rec).unwrap(), 3);
assert!(rec.0.is_empty());
}
#[test]
fn raising_in_steps() {
let mut cx = Context::new();
let mut lw = Lowering::new(&mut cx);
lw.lower(&Bin(2, BinOp::Add, 0, 1)).unwrap();
let (v2, v0) = (lw.get(2).unwrap(), lw.get(0).unwrap());
let cx = lw.context();
let e = cx.bin(BinOp::Mul, v2, v0).unwrap();
let todo = lw.unraised(e).unwrap();
assert_eq!(todo, [e]);
lw.emitted(e, 7).unwrap();
assert!(lw.unraised(e).unwrap().is_empty());
assert_eq!(lw.raise(e, &mut Rec(Vec::new())).unwrap(), 7);
let stray = lw.context().symbol("stray", Width::W8).unwrap();
let f = lw.context().bin(BinOp::Sub, e, stray).unwrap();
assert!(lw.unraised(f).is_err());
}
#[test]
fn raising_treats_owned_nodes_as_available_values() {
let mut cx = Context::new();
let hidden = cx.symbol("hidden", Width::W8).unwrap();
let one = cx.one(Width::W8).unwrap();
let owned = cx.add(hidden, one).unwrap();
let two = cx.constant_u64(Width::W8, 2).unwrap();
let root = cx.mul(owned, two).unwrap();
let mut lw = Lowering::new(&mut cx);
lw.define(7u32, owned).unwrap();
assert_eq!(lw.unraised(root).unwrap(), [two, root]);
let mut rec = Rec(Vec::new());
assert_eq!(lw.raise(root, &mut rec).unwrap(), 101);
assert_eq!(rec.0, ["%100 = const 0x2:8", "%101 = Mul [7, 100]"]);
let shared = lw.context().xor(owned, hidden).unwrap();
assert!(lw.raise(shared, &mut rec).is_err());
assert_eq!(rec.0.len(), 2, "failure precedes all emission");
lw.define(8, hidden).unwrap();
assert_eq!(lw.unraised(shared).unwrap(), [shared]);
assert_eq!(lw.raise(shared, &mut rec).unwrap(), 102);
}
#[test]
fn raising_rejects_stale_and_foreign_owned_handles() {
let mut cx = Context::new();
let e = cx.constant_u64(Width::W8, 7).unwrap();
let mut lw = Lowering::new(&mut cx);
lw.define(7u32, e).unwrap();
let kept = lw.detach();
let mut other = Context::new();
let mut lw = kept.attach(&mut other);
assert!(matches!(lw.unraised(e), Err(Error::ForeignExpr)));
assert!(matches!(
lw.raise(e, &mut Rec(Vec::new())),
Err(Error::ForeignExpr)
));
let kept = lw.detach();
cx.clear();
let mut lw = kept.attach(&mut cx);
assert!(matches!(lw.unraised(e), Err(Error::StaleExpr)));
assert!(matches!(
lw.raise(e, &mut Rec(Vec::new())),
Err(Error::StaleExpr)
));
}
#[test]
fn templates_agree_with_direct_construction_at_every_width() {
let smin = Template::new("select(a <s b, a, b)", &["a", "b"]).unwrap();
let absdiff = Template::new("select(a <u b, b - a, a - b)", &["a", "b"]).unwrap();
let pick = Template::new("select(c, a + 1, b)", &["c", "a", "b"]).unwrap();
for bits in [1u16, 7, 8, 32, 64] {
let w = Width::new(bits).unwrap();
let mut cx = Context::new();
let (a, b) = (cx.symbol("a", w).unwrap(), cx.symbol("b", w).unwrap());
let c = cx.symbol("c", Width::W1).unwrap();
let got = smin.instantiate(&mut cx, &[a, b]).unwrap();
let lt = cx.cmp(CmpOp::Slt, a, b).unwrap();
assert_eq!(got, cx.select(lt, a, b).unwrap());
let d = absdiff.instantiate(&mut cx, &[a, b]).unwrap();
let p = pick.instantiate(&mut cx, &[c, a, b]).unwrap();
assert_eq!(cx.width(p).unwrap(), w);
for k in 0..64u64 {
let x = BitVec::wrapping_from_u64(w, k.wrapping_mul(0x9e37_79b9_7f4a_7c15));
let y = BitVec::wrapping_from_u64(w, k.wrapping_mul(3));
let cond = k % 2 == 0;
let env = FnEnv(|key: &SymbolKey, _| match key.to_string().as_str() {
"a" => Some(x),
"b" => Some(y),
_ => Some(BitVec::from_bool(cond)),
});
let v = cx.eval(&[d, p], &env).unwrap();
let (xs, ys) = (x.to_u64().unwrap(), y.to_u64().unwrap());
let mask = u64::MAX >> (64 - bits);
assert_eq!(
v[0].to_u64(),
Some(xs.abs_diff(ys)),
"absdiff at {bits} bits"
);
let want = if cond { xs.wrapping_add(1) & mask } else { ys };
assert_eq!(v[1].to_u64(), Some(want), "pick at {bits} bits");
}
}
assert_eq!(
smin.compiled.lock().unwrap().len(),
5,
"compiled once per width"
);
let mut cx = Context::new();
let a = cx.symbol("a", Width::W8).unwrap();
assert!(smin.instantiate(&mut cx, &[a]).is_err());
assert!(Template::new("a + a", &["a", "a"]).is_err());
let broken = Template::new("a +", &["a"]).unwrap();
assert!(broken.check(&[Width::W8]).is_err());
assert!(pick.check(&[Width::W8, Width::W8, Width::W8]).is_err());
assert!(pick.check(&[Width::W1, Width::W8, Width::W8]).is_ok());
fn shared<T: Send + Sync>() {}
shared::<Template>();
}
}