Skip to main content

forge_lint/sol/gas/
cache_array_length.rs

1use super::CacheArrayLength;
2use crate::{
3    linter::{LateLintPass, LintContext},
4    sol::{Severity, SolLint},
5};
6use solar::{
7    ast::ElementaryType,
8    interface::{kw, sym},
9    sema::{
10        Gcx,
11        hir::{
12            self, BinOpKind, ExprKind, ItemId, LoopSource, Res, StateMutability, StmtKind,
13            VariableId,
14        },
15        ty::TyKind,
16    },
17};
18
19declare_forge_lint!(
20    CACHE_ARRAY_LENGTH,
21    Severity::Gas,
22    "cache-array-length",
23    "array length read in loop condition should be cached outside the loop"
24);
25
26#[derive(Clone, Copy)]
27struct LengthRead<'hir> {
28    expr: &'hir hir::Expr<'hir>,
29    base: &'hir hir::Expr<'hir>,
30}
31
32#[derive(Default)]
33struct LoopFacts {
34    written_vars: Vec<VariableId>,
35    mutates_array_length: bool,
36    has_state_mutating_call: bool,
37}
38
39impl LoopFacts {
40    const fn should_skip(&self) -> bool {
41        self.mutates_array_length || self.has_state_mutating_call
42    }
43
44    fn push_written_var(&mut self, var_id: VariableId) {
45        if !self.written_vars.contains(&var_id) {
46            self.written_vars.push(var_id);
47        }
48    }
49}
50
51impl<'hir> LateLintPass<'hir> for CacheArrayLength {
52    fn check_stmt(
53        &mut self,
54        ctx: &LintContext,
55        gcx: Gcx<'hir>,
56        hir: &'hir hir::Hir<'hir>,
57        stmt: &'hir hir::Stmt<'hir>,
58    ) {
59        let StmtKind::Loop(block, LoopSource::For | LoopSource::ForWithUpdate) = &stmt.kind else {
60            return;
61        };
62        let Some((condition, body)) = for_loop_parts(*block) else { return };
63
64        let mut reads = Vec::new();
65        collect_condition_length_reads(gcx, condition, &mut reads);
66        if reads.is_empty() {
67            return;
68        }
69
70        let mut facts = LoopFacts::default();
71        collect_stmt_facts(gcx, hir, body, &mut facts);
72        if facts.should_skip() {
73            return;
74        }
75
76        for read in reads {
77            if expr_is_loop_invariant(gcx, hir, read.base, &facts.written_vars) {
78                ctx.emit(&CACHE_ARRAY_LENGTH, read.expr.span);
79            }
80        }
81    }
82}
83
84fn for_loop_parts<'hir>(
85    block: hir::Block<'hir>,
86) -> Option<(&'hir hir::Expr<'hir>, &'hir hir::Stmt<'hir>)> {
87    let first = block.stmts.first()?;
88    match &first.kind {
89        StmtKind::If(condition, _, Some(else_stmt)) => {
90            matches!(&else_stmt.kind, StmtKind::Break).then_some((*condition, first))
91        }
92        _ => None,
93    }
94}
95
96fn collect_condition_length_reads<'hir>(
97    gcx: Gcx<'hir>,
98    expr: &'hir hir::Expr<'hir>,
99    reads: &mut Vec<LengthRead<'hir>>,
100) {
101    match &expr.peel_parens().kind {
102        ExprKind::Binary(lhs, op, rhs) if matches!(op.kind, BinOpKind::And | BinOpKind::Or) => {
103            collect_condition_length_reads(gcx, lhs, reads);
104            collect_condition_length_reads(gcx, rhs, reads);
105        }
106        ExprKind::Binary(lhs, op, rhs) if is_comparison(op.kind) => {
107            if matches!(lhs.peel_parens().kind, ExprKind::Ident(_)) {
108                collect_state_array_length_read(gcx, rhs, reads);
109            }
110            if matches!(rhs.peel_parens().kind, ExprKind::Ident(_)) {
111                collect_state_array_length_read(gcx, lhs, reads);
112            }
113        }
114        _ => {}
115    }
116}
117
118fn collect_state_array_length_read<'hir>(
119    gcx: Gcx<'hir>,
120    expr: &'hir hir::Expr<'hir>,
121    reads: &mut Vec<LengthRead<'hir>>,
122) {
123    let expr = expr.peel_parens();
124    if let ExprKind::Member(base, member) = &expr.kind
125        && member.name == sym::length
126        && is_state_array(gcx, base)
127    {
128        reads.push(LengthRead { expr, base });
129    }
130}
131
132fn collect_stmt_facts<'hir>(
133    gcx: Gcx<'hir>,
134    hir: &'hir hir::Hir<'hir>,
135    stmt: &'hir hir::Stmt<'hir>,
136    facts: &mut LoopFacts,
137) {
138    match &stmt.kind {
139        StmtKind::DeclSingle(var_id) => {
140            if let Some(expr) = hir.variable(*var_id).initializer {
141                collect_expr_facts(gcx, hir, expr, facts);
142            }
143        }
144        StmtKind::DeclMulti(_, expr)
145        | StmtKind::Emit(expr)
146        | StmtKind::Revert(expr)
147        | StmtKind::Expr(expr) => collect_expr_facts(gcx, hir, expr, facts),
148        StmtKind::Return(expr) => {
149            if let Some(expr) = expr {
150                collect_expr_facts(gcx, hir, expr, facts);
151            }
152        }
153        StmtKind::Block(block) | StmtKind::UncheckedBlock(block) | StmtKind::Loop(block, _) => {
154            for stmt in block.stmts {
155                collect_stmt_facts(gcx, hir, stmt, facts);
156            }
157        }
158        StmtKind::If(condition, then_stmt, else_stmt) => {
159            collect_expr_facts(gcx, hir, condition, facts);
160            collect_stmt_facts(gcx, hir, then_stmt, facts);
161            if let Some(else_stmt) = else_stmt {
162                collect_stmt_facts(gcx, hir, else_stmt, facts);
163            }
164        }
165        StmtKind::Try(stmt_try) => {
166            collect_expr_facts(gcx, hir, &stmt_try.expr, facts);
167            for clause in stmt_try.clauses {
168                for stmt in clause.block.stmts {
169                    collect_stmt_facts(gcx, hir, stmt, facts);
170                }
171            }
172        }
173        StmtKind::Break
174        | StmtKind::Continue
175        | StmtKind::Placeholder
176        | StmtKind::AssemblyBlock(_)
177        | StmtKind::Switch(_)
178        | StmtKind::Err(_) => {}
179    }
180}
181
182fn collect_expr_facts<'hir>(
183    gcx: Gcx<'hir>,
184    hir: &'hir hir::Hir<'hir>,
185    expr: &'hir hir::Expr<'hir>,
186    facts: &mut LoopFacts,
187) {
188    let expr = expr.peel_parens();
189    if array_length_mutated(gcx, expr) {
190        facts.mutates_array_length = true;
191    }
192
193    match &expr.kind {
194        ExprKind::Array(exprs) => {
195            for expr in *exprs {
196                collect_expr_facts(gcx, hir, expr, facts);
197            }
198        }
199        ExprKind::Assign(lhs, _, rhs) => {
200            collect_written_vars(lhs, facts);
201            collect_expr_facts(gcx, hir, lhs, facts);
202            collect_expr_facts(gcx, hir, rhs, facts);
203        }
204        ExprKind::Binary(lhs, _, rhs) => {
205            collect_expr_facts(gcx, hir, lhs, facts);
206            collect_expr_facts(gcx, hir, rhs, facts);
207        }
208        ExprKind::Call(callee, args, named_args) => {
209            if call_may_mutate_state(gcx, hir, callee) {
210                facts.has_state_mutating_call = true;
211            }
212            collect_expr_facts(gcx, hir, callee, facts);
213            for arg in args.exprs() {
214                collect_expr_facts(gcx, hir, arg, facts);
215            }
216            if let Some(named_args) = named_args {
217                for arg in named_args.args {
218                    collect_expr_facts(gcx, hir, &arg.value, facts);
219                }
220            }
221        }
222        ExprKind::Delete(inner) => {
223            collect_written_vars(inner, facts);
224            collect_expr_facts(gcx, hir, inner, facts);
225        }
226        ExprKind::Payable(inner) => collect_expr_facts(gcx, hir, inner, facts),
227        ExprKind::Unary(op, inner) => {
228            if op.kind.has_side_effects() {
229                collect_written_vars(inner, facts);
230            }
231            collect_expr_facts(gcx, hir, inner, facts);
232        }
233        ExprKind::Index(base, index) => {
234            collect_expr_facts(gcx, hir, base, facts);
235            if let Some(index) = index {
236                collect_expr_facts(gcx, hir, index, facts);
237            }
238        }
239        ExprKind::Slice(base, start, end) => {
240            collect_expr_facts(gcx, hir, base, facts);
241            if let Some(start) = start {
242                collect_expr_facts(gcx, hir, start, facts);
243            }
244            if let Some(end) = end {
245                collect_expr_facts(gcx, hir, end, facts);
246            }
247        }
248        ExprKind::Member(base, _) => collect_expr_facts(gcx, hir, base, facts),
249        ExprKind::Ternary(condition, then_expr, else_expr) => {
250            collect_expr_facts(gcx, hir, condition, facts);
251            collect_expr_facts(gcx, hir, then_expr, facts);
252            collect_expr_facts(gcx, hir, else_expr, facts);
253        }
254        ExprKind::Tuple(exprs) => {
255            for expr in exprs.iter().flatten() {
256                collect_expr_facts(gcx, hir, expr, facts);
257            }
258        }
259        ExprKind::Ident(_)
260        | ExprKind::Lit(_)
261        | ExprKind::New(_)
262        | ExprKind::TypeCall(_)
263        | ExprKind::Type(_)
264        | ExprKind::YulMember(..)
265        | ExprKind::Err(_) => {}
266    }
267}
268
269fn collect_written_vars(expr: &hir::Expr<'_>, facts: &mut LoopFacts) {
270    match &expr.peel_parens().kind {
271        ExprKind::Ident(resolutions) => {
272            if let Some(var_id) = variable_resolution(resolutions) {
273                facts.push_written_var(var_id);
274            }
275        }
276        ExprKind::Index(base, _) => {
277            collect_written_vars(base, facts);
278        }
279        ExprKind::Slice(base, _, _) => {
280            collect_written_vars(base, facts);
281        }
282        ExprKind::Member(base, _) | ExprKind::Payable(base) => collect_written_vars(base, facts),
283        ExprKind::Tuple(exprs) => {
284            for expr in exprs.iter().flatten() {
285                collect_written_vars(expr, facts);
286            }
287        }
288        _ => {}
289    }
290}
291
292fn array_length_mutated<'hir>(gcx: Gcx<'hir>, expr: &'hir hir::Expr<'hir>) -> bool {
293    match &expr.kind {
294        ExprKind::Assign(lhs, _, _) | ExprKind::Delete(lhs) => is_array_like(gcx, lhs),
295        ExprKind::Call(callee, _, _) => {
296            let ExprKind::Member(base, member) = &callee.peel_parens().kind else { return false };
297            matches!(member.name, sym::push | kw::Pop) && is_array_like(gcx, base)
298        }
299        _ => false,
300    }
301}
302
303fn call_may_mutate_state<'hir>(
304    gcx: Gcx<'hir>,
305    hir: &'hir hir::Hir<'hir>,
306    callee: &'hir hir::Expr<'hir>,
307) -> bool {
308    match &callee.peel_parens().kind {
309        ExprKind::Type(_) => false,
310        ExprKind::Ident(resolutions) => resolutions
311            .iter()
312            .find_map(|res| {
313                if let Res::Item(ItemId::Function(function_id)) = res {
314                    Some(hir.function(*function_id).mutates_state())
315                } else {
316                    None
317                }
318            })
319            .unwrap_or(true),
320        ExprKind::Member(base, member)
321            if matches!(member.name, sym::push | kw::Pop) && is_array_like(gcx, base) =>
322        {
323            false
324        }
325        _ => match gcx.type_of_expr(callee.peel_parens().id).map(|ty| ty.peel_refs().kind) {
326            Some(TyKind::Fn(function)) => function.state_mutability >= StateMutability::Payable,
327            _ => true,
328        },
329    }
330}
331
332fn expr_is_loop_invariant<'hir>(
333    gcx: Gcx<'hir>,
334    hir: &'hir hir::Hir<'hir>,
335    expr: &'hir hir::Expr<'hir>,
336    written_vars: &[VariableId],
337) -> bool {
338    match &expr.peel_parens().kind {
339        ExprKind::Ident(resolutions) => {
340            variable_resolution(resolutions).is_none_or(|var_id| !written_vars.contains(&var_id))
341        }
342        ExprKind::Lit(_) | ExprKind::Type(_) | ExprKind::TypeCall(_) => true,
343        ExprKind::Array(exprs) => {
344            exprs.iter().all(|expr| expr_is_loop_invariant(gcx, hir, expr, written_vars))
345        }
346        ExprKind::Binary(lhs, _, rhs) => {
347            expr_is_loop_invariant(gcx, hir, lhs, written_vars)
348                && expr_is_loop_invariant(gcx, hir, rhs, written_vars)
349        }
350        ExprKind::Call(callee, args, named_args) => {
351            call_is_safe_to_cache(gcx, hir, callee)
352                && expr_is_loop_invariant(gcx, hir, callee, written_vars)
353                && args.exprs().all(|arg| expr_is_loop_invariant(gcx, hir, arg, written_vars))
354                && named_args.is_none_or(|named_args| {
355                    named_args
356                        .args
357                        .iter()
358                        .all(|arg| expr_is_loop_invariant(gcx, hir, &arg.value, written_vars))
359                })
360        }
361        ExprKind::Index(base, index) => {
362            expr_is_loop_invariant(gcx, hir, base, written_vars)
363                && index.is_none_or(|index| expr_is_loop_invariant(gcx, hir, index, written_vars))
364        }
365        ExprKind::Slice(base, start, end) => {
366            expr_is_loop_invariant(gcx, hir, base, written_vars)
367                && start.is_none_or(|start| expr_is_loop_invariant(gcx, hir, start, written_vars))
368                && end.is_none_or(|end| expr_is_loop_invariant(gcx, hir, end, written_vars))
369        }
370        ExprKind::Member(base, _) | ExprKind::Payable(base) => {
371            expr_is_loop_invariant(gcx, hir, base, written_vars)
372        }
373        ExprKind::Ternary(condition, then_expr, else_expr) => {
374            expr_is_loop_invariant(gcx, hir, condition, written_vars)
375                && expr_is_loop_invariant(gcx, hir, then_expr, written_vars)
376                && expr_is_loop_invariant(gcx, hir, else_expr, written_vars)
377        }
378        ExprKind::Tuple(exprs) => {
379            exprs.iter().flatten().all(|expr| expr_is_loop_invariant(gcx, hir, expr, written_vars))
380        }
381        ExprKind::Unary(op, inner) => {
382            !op.kind.has_side_effects() && expr_is_loop_invariant(gcx, hir, inner, written_vars)
383        }
384        ExprKind::Assign(_, _, _)
385        | ExprKind::Delete(_)
386        | ExprKind::New(_)
387        | ExprKind::YulMember(..)
388        | ExprKind::Err(_) => false,
389    }
390}
391
392fn call_is_safe_to_cache<'hir>(
393    gcx: Gcx<'hir>,
394    hir: &'hir hir::Hir<'hir>,
395    callee: &'hir hir::Expr<'hir>,
396) -> bool {
397    match &callee.peel_parens().kind {
398        ExprKind::Type(_) => true,
399        ExprKind::Ident(resolutions) => resolutions
400            .iter()
401            .find_map(|res| {
402                if let Res::Item(ItemId::Function(function_id)) = res {
403                    Some(hir.function(*function_id).state_mutability <= StateMutability::View)
404                } else {
405                    None
406                }
407            })
408            .unwrap_or(false),
409        _ => match gcx.type_of_expr(callee.peel_parens().id).map(|ty| ty.peel_refs().kind) {
410            Some(TyKind::Fn(function)) => function.state_mutability <= StateMutability::View,
411            _ => false,
412        },
413    }
414}
415
416const fn is_comparison(op: BinOpKind) -> bool {
417    matches!(
418        op,
419        BinOpKind::Lt
420            | BinOpKind::Le
421            | BinOpKind::Gt
422            | BinOpKind::Ge
423            | BinOpKind::Eq
424            | BinOpKind::Ne
425    )
426}
427
428fn is_array_like<'hir>(gcx: Gcx<'hir>, expr: &'hir hir::Expr<'hir>) -> bool {
429    let Some(ty) = gcx.type_of_expr(expr.peel_parens().id) else { return false };
430    matches!(ty.peel_refs().kind, TyKind::DynArray(_) | TyKind::Elementary(ElementaryType::Bytes))
431}
432
433fn is_state_array<'hir>(gcx: Gcx<'hir>, expr: &'hir hir::Expr<'hir>) -> bool {
434    let ExprKind::Ident(resolutions) = &expr.peel_parens().kind else { return false };
435    let Some(var_id) = variable_resolution(resolutions) else { return false };
436    gcx.hir.variable(var_id).is_state_variable()
437        && matches!(
438            gcx.type_of_expr(expr.peel_parens().id).map(|ty| ty.peel_refs().kind),
439            Some(TyKind::DynArray(_))
440        )
441}
442
443fn variable_resolution(resolutions: &[Res]) -> Option<VariableId> {
444    resolutions.iter().find_map(|res| {
445        if let Res::Item(ItemId::Variable(var_id)) = res { Some(*var_id) } else { None }
446    })
447}