use alloc::boxed::Box;
use alloc::format;
use alloc::vec::Vec;
use syn::spanned::Spanned;
use syn::visit_mut::{self, VisitMut};
use crate::alpha::Renamer;
use crate::scope::{IntWidth, PrimTy, primitive_suffix};
const MAX_DEPTH: usize = 64;
#[derive(Clone, Copy, PartialEq, Eq)]
struct FixedInt {
width: IntWidth,
signed: bool,
}
impl FixedInt {
fn suffix(self) -> &'static str {
PrimTy::Int {
width: self.width,
signed: self.signed,
}
.suffix()
}
}
#[derive(Clone, Copy)]
struct Integer {
ty: FixedInt,
bits: u128,
}
impl Integer {
fn width(self) -> u32 {
self.ty.width.bits()
}
fn signed(self) -> bool {
self.ty.signed
}
fn mask(self) -> u128 {
u128::MAX >> (128 - self.width())
}
fn signed_value(self) -> i128 {
let sign = 1u128 << (self.width() - 1);
let bits = if self.bits & sign == 0 {
self.bits
} else {
self.bits | !self.mask()
};
i128::from_le_bytes(bits.to_le_bytes())
}
fn checked_signed(self, value: i128) -> Option<Self> {
let limit = i128::MAX >> (128 - self.width());
if value < -limit - 1 || value > limit {
return None;
}
Some(Self {
bits: u128::from_le_bytes(value.to_le_bytes()) & self.mask(),
..self
})
}
fn checked_unsigned(self, value: u128) -> Option<Self> {
(value <= self.mask()).then_some(Self {
bits: value,
..self
})
}
fn shift(self, right: Self, left: bool) -> Option<Self> {
if right.signed() && right.signed_value() < 0 {
return None;
}
let count = u32::try_from(right.bits).ok()?;
if count >= self.width() {
return None;
}
let bits = if left {
self.bits.checked_shl(count)? & self.mask()
} else if self.signed() {
u128::from_le_bytes(self.signed_value().checked_shr(count)?.to_le_bytes()) & self.mask()
} else {
self.bits.checked_shr(count)?
};
Some(Self { bits, ..self })
}
fn binary(self, right: Self, op: BinaryOp) -> Option<Self> {
let op = match op {
BinaryOp::Shl => return self.shift(right, true),
BinaryOp::Shr => return self.shift(right, false),
BinaryOp::BitAnd => {
return Some(Self {
bits: self.bits & right.bits,
..self
});
}
BinaryOp::BitOr => {
return Some(Self {
bits: self.bits | right.bits,
..self
});
}
BinaryOp::BitXor => {
return Some(Self {
bits: self.bits ^ right.bits,
..self
});
}
BinaryOp::Math(op) => op,
};
if self.signed() {
let (left, right) = (self.signed_value(), right.signed_value());
let minimum = -(i128::MAX >> (128 - self.width())) - 1;
if matches!(op, MathOp::Div | MathOp::Rem) && left == minimum && right == -1 {
return None;
}
let value = match op {
MathOp::Add => left.checked_add(right),
MathOp::Sub => left.checked_sub(right),
MathOp::Mul => left.checked_mul(right),
MathOp::Div => left.checked_div(right),
MathOp::Rem => left.checked_rem(right),
}?;
self.checked_signed(value)
} else {
let value = match op {
MathOp::Add => self.bits.checked_add(right.bits),
MathOp::Sub => self.bits.checked_sub(right.bits),
MathOp::Mul => self.bits.checked_mul(right.bits),
MathOp::Div => self.bits.checked_div(right.bits),
MathOp::Rem => self.bits.checked_rem(right.bits),
}?;
self.checked_unsigned(value)
}
}
fn expression(self, span: proc_macro2::Span) -> syn::Expr {
let negative = self.signed() && self.signed_value() < 0;
let magnitude = if negative {
self.signed_value().unsigned_abs()
} else {
self.bits
};
let literal = syn::Expr::Lit(syn::ExprLit {
attrs: Vec::new(),
lit: syn::Lit::Int(syn::LitInt::new(
&format!("{magnitude}{}", self.ty.suffix()),
span,
)),
});
if negative {
syn::Expr::Unary(syn::ExprUnary {
attrs: Vec::new(),
op: syn::UnOp::Neg(syn::token::Minus { spans: [span] }),
expr: Box::new(literal),
})
} else {
literal
}
}
}
#[derive(Clone, Copy)]
enum MathOp {
Add,
Sub,
Mul,
Div,
Rem,
}
#[derive(Clone, Copy)]
enum BinaryOp {
Math(MathOp),
BitAnd,
BitOr,
BitXor,
Shl,
Shr,
}
fn arithmetic(op: syn::BinOp) -> Option<BinaryOp> {
Some(match op {
syn::BinOp::Add(_) => BinaryOp::Math(MathOp::Add),
syn::BinOp::Sub(_) => BinaryOp::Math(MathOp::Sub),
syn::BinOp::Mul(_) => BinaryOp::Math(MathOp::Mul),
syn::BinOp::Div(_) => BinaryOp::Math(MathOp::Div),
syn::BinOp::Rem(_) => BinaryOp::Math(MathOp::Rem),
syn::BinOp::BitAnd(_) => BinaryOp::BitAnd,
syn::BinOp::BitOr(_) => BinaryOp::BitOr,
syn::BinOp::BitXor(_) => BinaryOp::BitXor,
syn::BinOp::Shl(_) => BinaryOp::Shl,
syn::BinOp::Shr(_) => BinaryOp::Shr,
_ => return None,
})
}
fn scalar_const(block: &syn::Block) -> Option<&syn::Expr> {
match block.stmts.as_slice() {
[syn::Stmt::Expr(expr, None)] => Some(expr),
_ => None,
}
}
enum Node<'a> {
Literal(&'a syn::LitInt),
Nested(&'a syn::Expr),
Unary(&'a syn::ExprUnary),
Binary(&'a syn::ExprBinary),
}
fn node(expr: &syn::Expr) -> Option<Node<'_>> {
match expr {
syn::Expr::Lit(syn::ExprLit {
lit: syn::Lit::Int(literal),
..
}) => Some(Node::Literal(literal)),
syn::Expr::Paren(paren) => Some(Node::Nested(&paren.expr)),
syn::Expr::Group(group) => Some(Node::Nested(&group.expr)),
syn::Expr::Const(constant) => scalar_const(&constant.block).map(Node::Nested),
syn::Expr::Unary(unary) => Some(Node::Unary(unary)),
syn::Expr::Binary(binary) => Some(Node::Binary(binary)),
_ => None,
}
}
fn closed(expr: &syn::Expr, depth: usize) -> bool {
if depth >= MAX_DEPTH {
return true;
}
match node(expr) {
Some(Node::Literal(_)) => true,
Some(Node::Nested(expr)) => closed(expr, depth + 1),
Some(Node::Unary(unary)) if matches!(unary.op, syn::UnOp::Neg(_) | syn::UnOp::Not(_)) => {
closed(&unary.expr, depth + 1)
}
Some(Node::Binary(binary)) if arithmetic(binary.op).is_some() => {
closed(&binary.left, depth + 1) && closed(&binary.right, depth + 1)
}
_ => false,
}
}
fn type_free_value(expr: &syn::Expr, depth: usize) -> Option<u8> {
if depth >= MAX_DEPTH || !crate::drift::expr_attrs(expr).is_empty() {
return None;
}
let value = match node(expr)? {
Node::Literal(literal) if literal.suffix().is_empty() => {
literal.base10_parse::<u8>().ok()?
}
Node::Nested(expr) => type_free_value(expr, depth + 1)?,
Node::Binary(binary) => {
let left = type_free_value(&binary.left, depth + 1)?;
let right = type_free_value(&binary.right, depth + 1)?;
match arithmetic(binary.op)? {
BinaryOp::Math(MathOp::Add) => left.checked_add(right)?,
BinaryOp::Math(MathOp::Sub) => left.checked_sub(right)?,
BinaryOp::Math(MathOp::Mul) => left.checked_mul(right)?,
BinaryOp::Math(MathOp::Div) => left.checked_div(right)?,
BinaryOp::Math(MathOp::Rem) => left.checked_rem(right)?,
BinaryOp::BitAnd => left & right,
BinaryOp::BitOr => left | right,
BinaryOp::BitXor => left ^ right,
BinaryOp::Shl if right < 8 => u8::try_from(u16::from(left) << right).ok()?,
BinaryOp::Shr if right < 8 => left >> right,
BinaryOp::Shl | BinaryOp::Shr => return None,
}
}
Node::Literal(_) | Node::Unary(_) => return None,
};
(value <= 127).then_some(value)
}
fn integer_type(ty: PrimTy) -> Option<FixedInt> {
match ty {
PrimTy::Int { width, signed } => Some(FixedInt { width, signed }),
PrimTy::Bool => None,
}
}
pub(crate) fn literal_expr(ty: PrimTy, bits: u128, span: proc_macro2::Span) -> syn::Expr {
if ty == PrimTy::Bool {
return syn::Expr::Lit(syn::ExprLit {
attrs: Vec::new(),
lit: syn::Lit::Bool(syn::LitBool::new(bits != 0, span)),
});
}
let ty = integer_type(ty).expect("an integer type");
Integer { ty, bits }.expression(span)
}
#[derive(Clone, Copy)]
enum TypeProof {
Unsuffixed,
Fixed(FixedInt),
}
impl TypeProof {
fn fixed(self) -> Option<FixedInt> {
match self {
Self::Unsuffixed => None,
Self::Fixed(ty) => Some(ty),
}
}
}
fn proven_type(expr: &syn::Expr, expected: Option<FixedInt>, depth: usize) -> Option<TypeProof> {
if depth >= MAX_DEPTH || !crate::drift::expr_attrs(expr).is_empty() {
return None;
}
match node(expr)? {
Node::Literal(literal) => {
if literal.suffix().is_empty() {
Some(expected.map_or(TypeProof::Unsuffixed, TypeProof::Fixed))
} else {
let ty = integer_type(primitive_suffix(literal.suffix())?)?;
if expected.is_some_and(|expected| expected != ty) {
return None;
}
Some(TypeProof::Fixed(ty))
}
}
Node::Nested(expr) => proven_type(expr, expected, depth + 1),
Node::Unary(unary) => proven_type(&unary.expr, expected, depth + 1),
Node::Binary(binary) => {
let left = proven_type(&binary.left, expected, depth + 1)?;
if matches!(binary.op, syn::BinOp::Shl(_) | syn::BinOp::Shr(_)) {
proven_type(&binary.right, None, depth + 1)?.fixed()?;
return Some(left);
}
let right = proven_type(&binary.right, expected, depth + 1)?;
match (left, right) {
(TypeProof::Fixed(left), TypeProof::Fixed(right)) if left != right => None,
(TypeProof::Fixed(ty), _) | (_, TypeProof::Fixed(ty)) => Some(TypeProof::Fixed(ty)),
(TypeProof::Unsuffixed, TypeProof::Unsuffixed) => Some(TypeProof::Unsuffixed),
}
}
}
}
fn literal(expr: &syn::Expr) -> Option<&syn::LitInt> {
match expr {
syn::Expr::Lit(syn::ExprLit {
lit: syn::Lit::Int(literal),
..
}) => Some(literal),
syn::Expr::Group(group) => literal(&group.expr),
_ => None,
}
}
pub(crate) fn signed_literal(expr: &syn::Expr) -> Option<(PrimTy, u128)> {
if !matches!(expr, syn::Expr::Unary(unary) if matches!(unary.op, syn::UnOp::Neg(_)) && literal(&unary.expr).is_some())
{
return None;
}
let ty = proven_type(expr, None, 0)?.fixed()?;
let value = evaluate(expr, ty, 0)?;
Some((
PrimTy::Int {
width: ty.width,
signed: ty.signed,
},
value.bits,
))
}
fn evaluate(expr: &syn::Expr, ty: FixedInt, depth: usize) -> Option<Integer> {
let model = Integer { ty, bits: 0 };
match node(expr)? {
Node::Literal(literal) => {
let magnitude = literal.base10_parse::<u128>().ok()?;
if model.signed() {
model.checked_signed(i128::try_from(magnitude).ok()?)
} else {
model.checked_unsigned(magnitude)
}
}
Node::Nested(expr) => evaluate(expr, ty, depth + 1),
Node::Unary(unary) => {
if matches!(unary.op, syn::UnOp::Neg(_))
&& model.signed()
&& let Some(literal) = literal(&unary.expr)
{
let magnitude = literal.base10_parse::<u128>().ok()?;
let limit = 1u128 << (model.width() - 1);
if magnitude > limit {
return None;
}
return Some(Integer {
bits: magnitude.wrapping_neg() & model.mask(),
..model
});
}
let value = evaluate(&unary.expr, ty, depth + 1)?;
match unary.op {
syn::UnOp::Neg(_) if value.signed() => {
value.checked_signed(value.signed_value().checked_neg()?)
}
syn::UnOp::Not(_) => Some(Integer {
bits: !value.bits & value.mask(),
..value
}),
_ => None,
}
}
Node::Binary(binary) => {
let left = evaluate(&binary.left, ty, depth + 1)?;
let right_ty = if matches!(binary.op, syn::BinOp::Shl(_) | syn::BinOp::Shr(_)) {
proven_type(&binary.right, None, depth + 1)?.fixed()?
} else {
ty
};
left.binary(
evaluate(&binary.right, right_ty, depth + 1)?,
arithmetic(binary.op)?,
)
}
}
}
fn literal_form(renamer: &Renamer<'_>, expr: &syn::Expr) -> Option<syn::Expr> {
let span = expr.span();
match expr {
syn::Expr::Lit(syn::ExprLit {
lit: syn::Lit::Byte(byte),
..
}) if byte.suffix().is_empty() => Some(literal_expr(
PrimTy::Int {
width: IntWidth::W8,
signed: false,
},
u128::from(byte.value()),
span,
)),
syn::Expr::Cast(cast) => {
let syn::Expr::Lit(syn::ExprLit { attrs, lit }) = &*cast.expr else {
return None;
};
let value = match lit {
syn::Lit::Int(int) if attrs.is_empty() && int.suffix().is_empty() => {
int.base10_parse::<u128>().ok()?
}
syn::Lit::Char(c)
if attrs.is_empty() && c.suffix().is_empty() && c.value().is_ascii() =>
{
u128::from(u32::from(c.value()))
}
syn::Lit::Byte(byte)
if attrs.is_empty() && byte.suffix().is_empty() && byte.value().is_ascii() =>
{
u128::from(byte.value())
}
_ => return None,
};
let syn::Type::Path(path) = &*cast.ty else {
return None;
};
let name = path.path.get_ident().filter(|_| path.qself.is_none())?;
if let Some(signed) = renamer.pointer_sized(name) {
let limit = if signed {
i16::MAX.unsigned_abs()
} else {
u16::MAX
};
if value > u128::from(limit) {
return None;
}
let suffix = if signed { "isize" } else { "usize" };
return Some(syn::Expr::Lit(syn::ExprLit {
attrs: Vec::new(),
lit: syn::Lit::Int(syn::LitInt::new(&format!("{value}{suffix}"), span)),
}));
}
let ty = integer_type(renamer.type_name_prim(name)?)?;
let model = Integer { ty, bits: 0 };
let value = if ty.signed {
model.checked_signed(i128::try_from(value).ok()?)
} else {
model.checked_unsigned(value)
}?;
Some(value.expression(span))
}
syn::Expr::Path(path) if path.qself.is_none() && path.path.leading_colon.is_none() => {
let mut segments = path.path.segments.iter();
let (Some(ty), Some(limit), None) = (segments.next(), segments.next(), segments.next())
else {
return None;
};
if !ty.arguments.is_none() || !limit.arguments.is_none() {
return None;
}
let prim @ PrimTy::Int { width, signed } = renamer.type_name_prim(&ty.ident)? else {
return None;
};
let all = u128::MAX >> (128 - width.bits());
let bits = if limit.ident == "MAX" {
if signed { all >> 1 } else { all }
} else if limit.ident == "MIN" {
if signed { 1 << (width.bits() - 1) } else { 0 }
} else {
return None;
};
Some(literal_expr(prim, bits, span))
}
_ => None,
}
}
pub(crate) fn normalize_expr(
renamer: &Renamer<'_>,
expr: &mut syn::Expr,
expected: Option<PrimTy>,
declared: bool,
) {
Normalizer {
renamer,
expected: expected.and_then(integer_type),
returns: None,
type_free: !declared,
contextual: true,
depth: 0,
}
.visit_expr_mut(expr);
}
pub(crate) fn normalize_block(renamer: &Renamer<'_>, block: &mut syn::Block) {
let returns = renamer.body_return(block).and_then(integer_type);
let mut normalizer = Normalizer {
renamer,
expected: None,
returns,
type_free: true,
contextual: true,
depth: 0,
};
if returns.is_some() {
normalizer.tail_block(block, returns);
} else {
normalizer.visit_block_mut(block);
}
}
struct Normalizer<'a, 'env> {
renamer: &'a Renamer<'env>,
expected: Option<FixedInt>,
returns: Option<FixedInt>,
type_free: bool,
contextual: bool,
depth: usize,
}
impl VisitMut for Normalizer<'_, '_> {
fn visit_expr_mut(&mut self, expr: &mut syn::Expr) {
if self.depth >= MAX_DEPTH || !crate::drift::expr_attrs(expr).is_empty() {
return;
}
if let Some(literal) = literal_form(self.renamer, expr) {
*expr = literal;
return;
}
if closed(expr, 0) {
if let Some(ty) = proven_type(expr, self.expected, 0).and_then(TypeProof::fixed)
&& let Some(value) = evaluate(expr, ty, 0)
{
if matches!(expr, syn::Expr::Lit(syn::ExprLit { lit: syn::Lit::Int(lit), .. }) if lit.suffix() == ty.suffix())
{
return;
}
*expr = value.expression(expr.span());
} else if self.type_free
&& literal(expr).is_none()
&& let Some(value) = type_free_value(expr, 0)
{
*expr = syn::Expr::Lit(syn::ExprLit {
attrs: Vec::new(),
lit: syn::Lit::Int(syn::LitInt::new(&format!("{value}"), expr.span())),
});
}
return;
}
let expected = self.expected.take();
let type_free = core::mem::replace(&mut self.type_free, true);
self.depth += 1;
if let syn::Expr::If(branch) = expr
&& expected.is_some()
&& let Some((_, otherwise)) = &mut branch.else_branch
{
self.visit_expr_mut(&mut branch.cond);
self.tail_block(&mut branch.then_branch, expected);
self.expected = expected;
self.visit_expr_mut(otherwise);
} else if let syn::Expr::Block(block) = expr
&& expected.is_some()
&& block.label.is_none()
{
self.tail_block(&mut block.block, expected);
} else if let syn::Expr::Return(r#return) = expr
&& let Some(value) = &mut r#return.expr
&& self.returns.is_some()
{
self.expected = self.returns;
self.visit_expr_mut(value);
} else if matches!(expr, syn::Expr::Closure(_) | syn::Expr::Async(_)) {
let returns = self.returns.take();
visit_mut::visit_expr_mut(self, expr);
self.returns = returns;
} else {
visit_mut::visit_expr_mut(self, expr);
}
self.depth -= 1;
self.expected = expected;
self.type_free = type_free;
}
fn visit_block_mut(&mut self, block: &mut syn::Block) {
if self.depth >= MAX_DEPTH || crate::schedule::unknown_macro_reaches(self.renamer, block) {
return;
}
let contextual = self.contextual;
self.contextual &= !block.stmts.iter().any(|stmt| {
matches!(stmt, syn::Stmt::Item(syn::Item::Use(_)))
|| (self.depth != 0 && matches!(stmt, syn::Stmt::Item(_)))
});
self.depth += 1;
for stmt in &mut block.stmts {
self.visit_stmt_mut(stmt);
}
self.depth -= 1;
self.contextual = contextual;
}
fn visit_local_mut(&mut self, local: &mut syn::Local) {
if !local.attrs.is_empty() {
return;
}
let mut annotation = None;
if let Some(init) = &mut local.init {
let expected = self.expected;
annotation = match &local.pat {
syn::Pat::Type(typed) if self.contextual => self.renamer.param_prim(&typed.ty),
_ => None,
};
self.expected = annotation.and_then(integer_type);
if !matches!(local.pat, syn::Pat::Type(_))
|| self.expected.is_some()
|| !closed(&init.expr, 0)
{
self.visit_expr_mut(&mut init.expr);
}
self.expected = expected;
if let Some((_, diverge)) = &mut init.diverge {
self.visit_expr_mut(diverge);
}
}
drop_stated_annotation(local, annotation);
self.visit_pat_mut(&mut local.pat);
}
fn visit_type_mut(&mut self, ty: &mut syn::Type) {
if crate::drift::type_attrs(ty).is_empty() {
visit_mut::visit_type_mut(self, ty);
}
}
fn visit_generic_argument_mut(&mut self, arg: &mut syn::GenericArgument) {
visit_mut::visit_generic_argument_mut(self, arg);
if let syn::GenericArgument::Const(expr) = arg {
crate::drift::fold_scalar_argument_block(expr);
}
}
fn visit_item_mut(&mut self, _: &mut syn::Item) {}
fn visit_macro_mut(&mut self, _: &mut syn::Macro) {}
fn visit_attribute_mut(&mut self, _: &mut syn::Attribute) {}
}
impl Normalizer<'_, '_> {
fn tail_block(&mut self, block: &mut syn::Block, expected: Option<FixedInt>) {
self.visit_block_mut(block);
if let Some(syn::Stmt::Expr(tail, None)) = block.stmts.last_mut() {
self.expected = expected;
self.visit_expr_mut(tail);
self.expected = None;
}
}
}
fn drop_stated_annotation(local: &mut syn::Local, annotation: Option<PrimTy>) {
let (Some(annotation @ PrimTy::Int { .. }), Some(init)) = (annotation, &local.init) else {
return;
};
let stated = match &*init.expr {
syn::Expr::Lit(_) => literal(&init.expr)
.filter(|literal| !literal.suffix().is_empty())
.and_then(|_| proven_type(&init.expr, None, 0)?.fixed())
.map(|ty| PrimTy::Int {
width: ty.width,
signed: ty.signed,
}),
_ => signed_literal(&init.expr).map(|(prim, _)| prim),
};
if init.diverge.is_some() || stated != Some(annotation) {
return;
}
if let syn::Pat::Type(typed) = &mut local.pat
&& typed.attrs.is_empty()
&& matches!(typed.pat.as_ref(), syn::Pat::Ident(ident)
if ident.attrs.is_empty() && ident.by_ref.is_none() && ident.subpat.is_none())
{
let binding = core::mem::replace(
&mut *typed.pat,
syn::Pat::Verbatim(proc_macro2::TokenStream::new()),
);
local.pat = binding;
}
}