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

forge/brutalizer/
visitor.rs
