use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use alloc::string::String;
use alloc::vec::Vec;
use syn::spanned::Spanned as _;
use syn::visit::{self, Visit};
use syn::visit_mut::{self, VisitMut};
use crate::alpha::Renamer;
use crate::schedule::{self, BitwiseOp, BlockFacts, Lbl, Port, TotalBinary, Value};
use crate::scope::PrimTy;
struct Bound {
depth: usize,
nodes: usize,
valid: bool,
}
impl<'ast> Visit<'ast> for Bound {
fn visit_expr(&mut self, expr: &'ast syn::Expr) {
if !self.valid {
return;
}
self.nodes += 1;
self.depth += 1;
if self.depth > schedule::MAX_BLOCK_DEPTH || self.nodes > schedule::MAX_REGION_NODES {
self.valid = false;
} else {
visit::visit_expr(self, expr);
}
self.depth -= 1;
}
fn visit_block(&mut self, _: &'ast syn::Block) {}
}
pub(crate) fn within_bounds(expr: &syn::Expr) -> bool {
let mut bound = Bound {
depth: 0,
nodes: 0,
valid: true,
};
bound.visit_expr(expr);
bound.valid
}
pub(crate) fn within_struct_bounds(expr: &syn::ExprStruct) -> bool {
let mut bound = Bound {
depth: 1,
nodes: 1,
valid: true,
};
bound.visit_expr_struct(expr);
bound.valid
}
pub(crate) fn normalize(renamer: &Renamer<'_>, block: &mut syn::Block, facts: &BlockFacts) {
if !facts.can_rewrite() {
return;
}
let mut env = BTreeMap::new();
renamer.frame_env(&mut env);
let mut walker = Walker {
env,
locals: 0,
dedup: true,
};
for stmt in &mut block.stmts {
let definite = matches!(
schedule::stmt_kind(renamer, stmt),
schedule::StmtKind::Boundary(true)
);
match stmt {
syn::Stmt::Local(local) => {
let plan = facts.plan(&local.pat);
if local.attrs.is_empty()
&& plan.is_none_or(|plan| plan.active)
&& let Some(init) = &mut local.init
&& init.diverge.is_none()
{
walker.visit_expr_mut(&mut init.expr);
}
for name in schedule::pat_names(&local.pat) {
let value = if let Some(plan) = plan {
Value {
port: plan.port.clone(),
prim: if plan.active { plan.prim } else { None },
helper: None,
}
} else {
walker.locals += 1;
Value {
port: Port::Local(walker.locals),
prim: None,
helper: None,
}
};
walker.env.insert(name, value);
}
}
syn::Stmt::Expr(expr, _) => {
walker.visit_expr_mut(expr);
if let Some(name) = schedule::assign_target(stmt) {
walker.locals += 1;
walker.env.insert(
name,
Value {
port: Port::Local(walker.locals),
prim: None,
helper: None,
},
);
}
}
syn::Stmt::Item(_) | syn::Stmt::Macro(_) => {}
}
if definite {
break;
}
}
}
struct Walker {
env: BTreeMap<String, Value>,
locals: usize,
dedup: bool,
}
fn empty_expr() -> syn::Expr {
syn::Expr::Verbatim(proc_macro2::TokenStream::new())
}
fn tail_expr_mut(block: &mut syn::Block) -> Option<&mut syn::Expr> {
let [syn::Stmt::Expr(expr, None)] = block.stmts.as_mut_slice() else {
return None;
};
Some(expr)
}
impl Walker {
fn scalar_in_place(&self, expr: &mut syn::Expr, algebra: bool) -> Option<Lbl> {
let owned = core::mem::replace(expr, empty_expr());
let (normalized, label) = self.scalar(owned, algebra);
*expr = normalized;
label
}
fn scalar(&self, expr: syn::Expr, algebra: bool) -> (syn::Expr, Option<Lbl>) {
if !crate::drift::expr_attrs(&expr).is_empty() {
return (expr, None);
}
match expr {
syn::Expr::Binary(mut binary) => {
let left = self.scalar_in_place(&mut binary.left, algebra);
let right = self.scalar_in_place(&mut binary.right, algebra);
let (Some(left), Some(right)) = (left, right) else {
return (syn::Expr::Binary(binary), None);
};
if self.dedup
&& let Some(label) = schedule::true_comparison(binary.op, &left, &right)
{
let expr = if left == Lbl::Bool(true) {
*binary.right
} else {
*binary.left
};
return (expr, Some(label));
}
match schedule::total_binary(binary.op, &left, &right) {
Some(TotalBinary::Flat(op, prim)) if !algebra => {
let mut leaves = Vec::new();
schedule::collect_leaves(op, prim, left, &mut leaves);
schedule::collect_leaves(op, prim, right, &mut leaves);
(
syn::Expr::Binary(binary),
Some(schedule::flat_label(op, prim, leaves)),
)
}
Some(TotalBinary::Flat(op, prim)) => {
let span = binary.span();
let mut pieces = Vec::new();
extend(op, prim, left, *binary.left, &mut pieces);
extend(op, prim, right, *binary.right, &mut pieces);
pieces.sort_by(|a, b| a.0.cmp(&b.0));
if self.dedup && op != BitwiseOp::Xor {
pieces.dedup_by(|a, b| a.0 == b.0);
}
fold_literals(op, prim, &mut pieces, span);
if pieces.len() == 1 {
let (label, expr) = pieces.pop().expect("a single operand");
return (expr, Some(label));
}
let count = pieces.len();
let mut labels = Vec::with_capacity(count);
let expression = {
let mut expressions = pieces.into_iter().map(|(label, expr)| {
labels.push(label);
expr
});
emit(op, &mut expressions, count, span)
};
(expression, Some(Lbl::Flat(op.tag(), labels, prim)))
}
Some(TotalBinary::Ordered { tag, reverse, prim }) => {
let (left, right) = if let Some(reverse) = reverse {
if algebra {
core::mem::swap(&mut binary.left, &mut binary.right);
binary.op = reverse.bin_op();
}
(right, left)
} else {
(left, right)
};
(
syn::Expr::Binary(binary),
Some(Lbl::Op(tag, Box::new(left), Box::new(right), prim)),
)
}
None => (syn::Expr::Binary(binary), None),
}
}
syn::Expr::Unary(mut unary) if self.dedup && matches!(unary.op, syn::UnOp::Not(_)) => {
if let Some(inner) = self.scalar_in_place(&mut unary.expr, algebra)
&& inner.prim() == Some(PrimTy::Bool)
{
if !matches!(
inner,
Lbl::Not(_) | Lbl::Bool(_) | Lbl::Op('∧' | '∨', _, _, _)
) {
return (syn::Expr::Unary(unary), Some(schedule::not_label(inner)));
}
let (expr, label) = negate(*unary.expr, inner);
return (expr, Some(label));
}
(syn::Expr::Unary(unary), None)
}
syn::Expr::If(branch) if self.dedup => self.branch(branch, algebra),
syn::Expr::Group(group) => self.scalar(*group.expr, algebra),
expr @ (syn::Expr::Lit(_) | syn::Expr::Path(_) | syn::Expr::Unary(_)) => {
let label = schedule::scalar_leaf(&self.env, &expr);
(expr, label)
}
expr => (expr, None),
}
}
fn condition(&self, expr: syn::Expr, algebra: bool) -> (syn::Expr, Option<Lbl>, bool) {
match expr {
syn::Expr::Unary(mut unary)
if unary.attrs.is_empty() && matches!(unary.op, syn::UnOp::Not(_)) =>
{
let owned = core::mem::replace(&mut *unary.expr, empty_expr());
let (expr, label, inverted) = self.condition(owned, algebra);
if label
.as_ref()
.is_some_and(|label| label.prim() == Some(PrimTy::Bool))
{
(expr, label, !inverted)
} else {
*unary.expr = expr;
(syn::Expr::Unary(unary), None, false)
}
}
syn::Expr::Group(group) if group.attrs.is_empty() => {
self.condition(*group.expr, algebra)
}
expr => {
let (expr, label) = self.scalar(expr, algebra);
(expr, label, false)
}
}
}
fn tail(&self, block: &mut syn::Block, algebra: bool) -> Option<Lbl> {
let expr = tail_expr_mut(block)?;
self.scalar_in_place(expr, algebra)
}
fn branch(&self, mut branch: syn::ExprIf, algebra: bool) -> (syn::Expr, Option<Lbl>) {
let Some((_, otherwise)) = &mut branch.else_branch else {
return (syn::Expr::If(branch), None);
};
let syn::Expr::Block(else_block) = &mut **otherwise else {
return (syn::Expr::If(branch), None);
};
if !schedule::branches_unattributed(&branch.attrs, &branch.then_branch, else_block) {
return (syn::Expr::If(branch), None);
}
let mut then = self.tail(&mut branch.then_branch, algebra);
let mut otherwise = self.tail(&mut else_block.block, algebra);
let complementary = matches!(
(&then, &otherwise),
(Some(Lbl::Bool(true)), Some(Lbl::Bool(false)))
| (Some(Lbl::Bool(false)), Some(Lbl::Bool(true)))
);
let owned = core::mem::replace(&mut *branch.cond, empty_expr());
let (cond, label, inverted) = self.condition(owned, algebra && complementary);
*branch.cond = cond;
let Some(cond_label) = label.filter(|label| label.prim() == Some(PrimTy::Bool)) else {
return (syn::Expr::If(branch), None);
};
if inverted {
core::mem::swap(&mut branch.then_branch, &mut else_block.block);
core::mem::swap(&mut then, &mut otherwise);
}
let (Some(then), Some(otherwise)) = (then, otherwise) else {
return (syn::Expr::If(branch), None);
};
let Some(value) = schedule::branch_value(&cond_label, &then, &otherwise) else {
return (syn::Expr::If(branch), None);
};
if let Lbl::Bool(selected) = cond_label {
let block = if selected {
&mut branch.then_branch
} else {
&mut else_block.block
};
let expr = tail_expr_mut(block).expect("proven scalar tail");
return (core::mem::replace(expr, empty_expr()), Some(value));
}
if then == Lbl::Bool(true) {
(*branch.cond, Some(value))
} else {
let (expr, label) = negate(*branch.cond, cond_label);
(expr, Some(label))
}
}
}
impl VisitMut for Walker {
fn visit_expr_mut(&mut self, expr: &mut syn::Expr) {
if !crate::drift::expr_attrs(expr).is_empty() {
return;
}
match expr {
syn::Expr::Binary(_)
| syn::Expr::Unary(_)
| syn::Expr::If(_)
| syn::Expr::Lit(_)
| syn::Expr::Path(_)
| syn::Expr::Group(_) => {
let owned = core::mem::replace(expr, empty_expr());
*expr = self.scalar(owned, true).0;
}
syn::Expr::Reference(reference) => {
let dedup = core::mem::replace(&mut self.dedup, false);
self.visit_expr_mut(&mut reference.expr);
self.dedup = dedup;
}
syn::Expr::RawAddr(raw) => {
let dedup = core::mem::replace(&mut self.dedup, false);
self.visit_expr_mut(&mut raw.expr);
self.dedup = dedup;
}
syn::Expr::MethodCall(method) => {
let dedup = core::mem::replace(&mut self.dedup, false);
self.visit_expr_mut(&mut method.receiver);
self.dedup = dedup;
for argument in &mut method.args {
self.visit_expr_mut(argument);
}
}
syn::Expr::Block(_)
| syn::Expr::Const(_)
| syn::Expr::Unsafe(_)
| syn::Expr::Async(_)
| syn::Expr::TryBlock(_)
| syn::Expr::Closure(_)
| syn::Expr::Match(_)
| syn::Expr::While(_)
| syn::Expr::ForLoop(_)
| syn::Expr::Loop(_)
| syn::Expr::Let(_)
| syn::Expr::Macro(_)
| syn::Expr::Verbatim(_) => {}
_ => visit_mut::visit_expr_mut(self, expr),
}
}
}
fn negate(expr: syn::Expr, label: Lbl) -> (syn::Expr, Lbl) {
let span = expr.span();
match (expr, label) {
(syn::Expr::Unary(unary), Lbl::Not(inner)) => (*unary.expr, *inner),
(syn::Expr::Binary(mut binary), Lbl::Op(tag @ ('∧' | '∨'), left, right, _)) => {
let (left_expr, left) =
negate(core::mem::replace(&mut *binary.left, empty_expr()), *left);
let (right_expr, right) =
negate(core::mem::replace(&mut *binary.right, empty_expr()), *right);
*binary.left = left_expr;
*binary.right = right_expr;
let tag = if tag == '∧' { '∨' } else { '∧' };
binary.op = if tag == '∧' {
syn::BinOp::And(syn::token::AndAnd::default())
} else {
syn::BinOp::Or(syn::token::OrOr::default())
};
(
syn::Expr::Binary(binary),
Lbl::Op(tag, Box::new(left), Box::new(right), PrimTy::Bool),
)
}
(_, Lbl::Bool(value)) => (
syn::Expr::Lit(syn::ExprLit {
attrs: Vec::new(),
lit: syn::Lit::Bool(syn::LitBool::new(!value, span)),
}),
Lbl::Bool(!value),
),
(expr, label) => (
syn::Expr::Unary(syn::ExprUnary {
attrs: Vec::new(),
op: syn::UnOp::Not(syn::token::Not { spans: [span] }),
expr: Box::new(expr),
}),
schedule::not_label(label),
),
}
}
fn extend(
op: BitwiseOp,
prim: PrimTy,
label: Lbl,
expr: syn::Expr,
out: &mut Vec<(Lbl, syn::Expr)>,
) {
match label {
Lbl::Flat(inner_op, labels, inner_prim) if inner_op == op.tag() && inner_prim == prim => {
let mut expressions = Vec::with_capacity(labels.len());
decompose(op, expr, &mut expressions);
assert_eq!(
labels.len(),
expressions.len(),
"one expression per operand"
);
out.extend(labels.into_iter().zip(expressions));
}
label => out.push((label, expr)),
}
}
fn decompose(op: BitwiseOp, expr: syn::Expr, out: &mut Vec<syn::Expr>) {
match expr {
syn::Expr::Binary(binary) if BitwiseOp::from_bin_op(binary.op) == Some(op) => {
decompose(op, *binary.left, out);
decompose(op, *binary.right, out);
}
expr => out.push(expr),
}
}
fn fold_literals(
op: BitwiseOp,
prim: PrimTy,
pieces: &mut Vec<(Lbl, syn::Expr)>,
span: proc_macro2::Span,
) {
let literal = |label: &Lbl| label.literal_bits().is_some();
if pieces.iter().filter(|piece| literal(&piece.0)).count() < 2 {
return;
}
let mut acc = None;
pieces.retain(|piece| {
let Some(bits) = piece.0.literal_bits() else {
return true;
};
acc = Some(match acc {
None => bits,
Some(left) => op.fold(left, bits),
});
false
});
let bits = acc.expect("at least two literal operands");
let label = schedule::literal_label(prim, bits);
pieces.push((label, crate::constants::literal_expr(prim, bits, span)));
pieces.sort_by(|a, b| a.0.cmp(&b.0));
}
fn emit(
op: BitwiseOp,
expressions: &mut impl Iterator<Item = syn::Expr>,
count: usize,
span: proc_macro2::Span,
) -> syn::Expr {
if count == 1 {
return expressions
.next()
.expect("an operand for each counted leaf");
}
let left = emit(op, expressions, count / 2, span);
let right = emit(op, expressions, count - count / 2, span);
let op = match op {
BitwiseOp::And => syn::BinOp::BitAnd(syn::token::And { spans: [span] }),
BitwiseOp::Or => syn::BinOp::BitOr(syn::token::Or { spans: [span] }),
BitwiseOp::Xor => syn::BinOp::BitXor(syn::token::Caret { spans: [span] }),
};
syn::Expr::Binary(syn::ExprBinary {
attrs: Vec::new(),
left: Box::new(left),
op,
right: Box::new(right),
})
}