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