forge/brutalizer/
visitor.rs

1use std::ops::ControlFlow;
2
3use solar::ast::{CallArgsKind, Expr, ExprKind, ItemFunction, SourceUnit, Type, visit::Visit};
4
5use super::{
6    assembly::assembly_transforms,
7    transform::{Transform, span_text},
8    value::{brutalize_cast, brutalize_payable_address, deterministic_mask},
9};
10
11pub(super) fn collect_transforms<'ast>(
12    source: &str,
13    ast: &'ast SourceUnit<'ast>,
14) -> Vec<Transform> {
15    let mut visitor = BrutalizerVisitor::new(source);
16    let _ = visitor.visit_source_unit(ast);
17    visitor.transforms
18}
19
20struct BrutalizerVisitor<'src> {
21    transforms: Vec<Transform>,
22    source: &'src str,
23}
24
25impl<'src> BrutalizerVisitor<'src> {
26    const fn new(source: &'src str) -> Self {
27        Self { transforms: Vec::new(), source }
28    }
29}
30
31impl<'ast, 'src> Visit<'ast> for BrutalizerVisitor<'src> {
32    type BreakValue = ();
33
34    fn visit_expr(&mut self, expr: &'ast Expr<'ast>) -> ControlFlow<Self::BreakValue> {
35        if let Some((ty, arg_text)) = cast_call(self.source, expr) {
36            let mask = deterministic_mask(expr.span);
37            if let Some(replacement) = brutalize_cast(ty, arg_text, &mask) {
38                self.transforms.push(Transform::Replace { span: expr.span, replacement });
39                return ControlFlow::Continue(());
40            }
41        }
42        if let Some(arg_text) = payable_call(self.source, expr) {
43            let mask = deterministic_mask(expr.span);
44            let replacement = brutalize_payable_address(arg_text, &mask);
45            self.transforms.push(Transform::Replace { span: expr.span, replacement });
46            return ControlFlow::Continue(());
47        }
48
49        self.walk_expr(expr)
50    }
51
52    fn visit_item_function(
53        &mut self,
54        func: &'ast ItemFunction<'ast>,
55    ) -> ControlFlow<Self::BreakValue> {
56        self.transforms.extend(assembly_transforms(func));
57        self.walk_item_function(func)
58    }
59}
60
61fn payable_call<'src>(source: &'src str, expr: &Expr<'_>) -> Option<&'src str> {
62    let ExprKind::Payable(call_args) = &expr.kind else { return None };
63    let CallArgsKind::Unnamed(args_exprs) = &call_args.kind else { return None };
64    let arg_text = span_text(source, args_exprs.first()?.span)?;
65    (!arg_text.is_empty()).then_some(arg_text)
66}
67
68fn cast_call<'ast, 'src>(
69    source: &'src str,
70    expr: &'ast Expr<'ast>,
71) -> Option<(&'ast Type<'ast>, &'src str)> {
72    let ExprKind::Call(callee, call_args) = &expr.kind else { return None };
73    let ExprKind::Type(ty) = &callee.kind else { return None };
74    let CallArgsKind::Unnamed(args_exprs) = &call_args.kind else { return None };
75    let arg_text = span_text(source, args_exprs.first()?.span)?;
76    (!arg_text.is_empty()).then_some((ty, arg_text))
77}
forge/brutalizer/visitor.rs

forge/brutalizer/
visitor.rs
