1use super::EnumerableLoopRemoval;
2use crate::{
3 linter::{LateLintPass, LintContext},
4 sol::{Severity, SolLint, analysis::primitives::branch_always_exits},
5};
6use alloy_primitives::U256;
7use solar::{
8 ast::{LitKind, UnOpKind},
9 interface::{Span, Symbol},
10 sema::{
11 Gcx,
12 hir::{
13 self, CallArgs, CallArgsKind, Expr, ExprKind, FunctionId, Hir, ItemId, LoopSource, Res,
14 Stmt, StmtKind, VarKind, VariableId, Visit,
15 },
16 ty::TyKind,
17 },
18};
19use std::{collections::HashSet, convert::Infallible, ops::ControlFlow};
20
21declare_forge_lint!(
22 ENUMERABLE_LOOP_REMOVAL,
23 Severity::High,
24 "enumerable-loop-removal",
25 "`remove` on an EnumerableSet inside a loop that iterates it with `at` corrupts the iteration"
26);
27
28impl<'hir> LateLintPass<'hir> for EnumerableLoopRemoval {
37 fn check_function(
38 &mut self,
39 ctx: &LintContext,
40 gcx: Gcx<'hir>,
41 hir: &'hir Hir<'hir>,
42 func: &'hir hir::Function<'hir>,
43 ) {
44 if let Some(body) = &func.body {
45 let mut finder =
46 LoopFinder { gcx, hir, ctx, bindings: Vec::new(), emitted: HashSet::new() };
47 finder.walk_body(body.stmts);
48 }
49 }
50}
51
52struct LoopFinder<'ctx, 's, 'c, 'hir> {
57 gcx: Gcx<'hir>,
58 hir: &'hir Hir<'hir>,
59 ctx: &'ctx LintContext<'s, 'c>,
60 bindings: Vec<(VariableId, Option<SetPath>)>,
65 emitted: HashSet<Span>,
67}
68
69impl<'hir> LoopFinder<'_, '_, '_, 'hir> {
70 fn walk_body(&mut self, stmts: &'hir [Stmt<'hir>]) {
71 for stmt in stmts {
72 self.walk_stmt(stmt);
73 }
74 }
75
76 fn walk_stmt(&mut self, stmt: &'hir Stmt<'hir>) {
77 if let StmtKind::Block(block) = &stmt.kind
80 && let Some(last) = block.stmts.last()
81 && let StmtKind::Loop(body, LoopSource::For) = &last.kind
82 {
83 let init = &block.stmts[..block.stmts.len() - 1];
84 self.walk_body(init);
85 self.enter_loop(init, body.stmts, LoopSource::For);
86 return;
87 }
88 match &stmt.kind {
89 StmtKind::Block(block) | StmtKind::UncheckedBlock(block) => {
91 self.walk_body(block.stmts);
92 }
93 StmtKind::Loop(body, source) => self.enter_loop(&[], body.stmts, *source),
95 StmtKind::If(_, then, else_) => {
96 self.poison_writes(std::slice::from_ref(stmt));
100 let mark = self.bindings.len();
101 self.walk_stmt(then);
102 self.bindings.truncate(mark);
103 if let Some(else_) = else_ {
104 self.walk_stmt(else_);
105 self.bindings.truncate(mark);
106 }
107 }
108 StmtKind::Try(try_) => {
109 self.poison_writes(std::slice::from_ref(stmt));
111 let mark = self.bindings.len();
112 for clause in try_.clauses {
113 self.walk_body(clause.block.stmts);
114 self.bindings.truncate(mark);
115 }
116 }
117 _ => self.apply_bindings(stmt),
118 }
119 }
120
121 fn enter_loop(
126 &mut self,
127 init: &'hir [Stmt<'hir>],
128 body: &'hir [Stmt<'hir>],
129 source: LoopSource,
130 ) {
131 self.poison_writes(init);
132 self.poison_writes(body);
133 self.analyze_loop(real_body(source, body));
137 let mark = self.bindings.len();
138 self.walk_body(body);
139 self.bindings.truncate(mark);
140 }
141
142 fn apply_bindings(&mut self, stmt: &'hir Stmt<'hir>) {
147 self.poison_writes(std::slice::from_ref(stmt));
148 match &stmt.kind {
149 StmtKind::DeclSingle(variable_id) => {
150 if let Some(initializer) = self.hir.variable(*variable_id).initializer {
151 let resolved = set_path(self.hir, initializer, &self.bindings, &mut Vec::new());
152 self.bindings.push((*variable_id, resolved));
153 }
154 }
155 StmtKind::Expr(expr) => {
156 if let ExprKind::Assign(target, None, value) = &expr.peel_parens().kind
157 && let ExprKind::Ident(resolutions) = &target.peel_parens().kind
158 {
159 let resolved = set_path(self.hir, value, &self.bindings, &mut Vec::new());
160 for res in *resolutions {
161 if let Res::Item(ItemId::Variable(variable_id)) = res {
162 self.bindings.push((*variable_id, resolved.clone()));
163 }
164 }
165 }
166 }
167 _ => {}
168 }
169 }
170
171 fn poison_writes(&mut self, stmts: &'hir [Stmt<'hir>]) {
173 let mut written = Vec::new();
174 collect_variables(self.hir, stmts, &mut written);
175 for variable_id in written {
176 self.bindings.push((variable_id, None));
177 }
178 }
179
180 fn analyze_loop(&mut self, body: &'hir [Stmt<'hir>]) {
184 if !body_is_straight_line(body) {
187 return;
188 }
189 let cadence = ascending_cadence(self.hir, body);
193 if cadence.is_empty() {
194 return;
195 }
196 let mut ats = AtCollector {
198 gcx: self.gcx,
199 hir: self.hir,
200 bindings: &self.bindings,
201 cadence: &cadence,
202 iterated: Vec::new(),
203 };
204 for stmt in body {
205 let _ = ats.visit_stmt(stmt);
206 }
207 if ats.iterated.is_empty() {
208 return;
209 }
210 let mut removes = Vec::new();
214 let mut scan = RemoveScanner {
215 gcx: self.gcx,
216 hir: self.hir,
217 bindings: &self.bindings,
218 out: &mut removes,
219 };
220 for stmt in body {
221 let _ = scan.visit_stmt(stmt);
222 }
223 for (removed, span) in removes {
224 let corrupts = ats.iterated.iter().any(|iterated| paths_alias(&removed, iterated));
225 if corrupts && self.emitted.insert(span) {
226 self.ctx.emit(&ENUMERABLE_LOOP_REMOVAL, span);
227 }
228 }
229 }
230}
231
232const fn real_body<'hir>(source: LoopSource, body: &'hir [Stmt<'hir>]) -> &'hir [Stmt<'hir>] {
239 match source {
240 LoopSource::For | LoopSource::While => {
241 if let [only] = body
242 && let StmtKind::If(_, then, Some(else_)) = &only.kind
243 && matches!(else_.kind, StmtKind::Break)
244 {
245 return std::slice::from_ref(*then);
246 }
247 body
248 }
249 LoopSource::DoWhile => {
250 if let Some((last, rest)) = body.split_last()
251 && let StmtKind::If(_, then, Some(else_)) = &last.kind
252 && matches!(then.kind, StmtKind::Continue)
253 && matches!(else_.kind, StmtKind::Break)
254 {
255 return rest;
256 }
257 body
258 }
259 }
260}
261
262fn body_is_straight_line(stmts: &[Stmt<'_>]) -> bool {
268 stmts.iter().all(|stmt| {
269 !branch_always_exits(stmt)
270 && match &stmt.kind {
271 StmtKind::Block(block) | StmtKind::UncheckedBlock(block) => {
272 body_is_straight_line(block.stmts)
273 }
274 StmtKind::If(..)
275 | StmtKind::Try(..)
276 | StmtKind::Loop(..)
277 | StmtKind::AssemblyBlock(..)
278 | StmtKind::Break
279 | StmtKind::Continue => false,
280 _ => true,
281 }
282 })
283}
284
285fn ascending_cadence<'hir>(hir: &'hir Hir<'hir>, body: &'hir [Stmt<'hir>]) -> Vec<VariableId> {
292 let mut cadence = Vec::new();
293 let mut other_writes = HashSet::new();
294 collect_cadence_writes(hir, body, &mut cadence, &mut other_writes);
295 cadence.retain(|variable_id| !other_writes.contains(variable_id));
296 cadence
297}
298
299fn collect_cadence_writes<'hir>(
303 hir: &'hir Hir<'hir>,
304 stmts: &'hir [Stmt<'hir>],
305 cadence: &mut Vec<VariableId>,
306 other_writes: &mut HashSet<VariableId>,
307) {
308 for stmt in stmts {
309 match &stmt.kind {
310 StmtKind::Block(block) | StmtKind::UncheckedBlock(block) => {
311 collect_cadence_writes(hir, block.stmts, cadence, other_writes);
312 }
313 _ => {
314 let ascending = match &stmt.kind {
315 StmtKind::Expr(expr) => ascending_step(expr.peel_parens()),
316 _ => None,
317 };
318 let mut written = Vec::new();
319 match &stmt.kind {
320 StmtKind::DeclSingle(variable_id) => written.push(*variable_id),
321 StmtKind::DeclMulti(variable_ids, _) => {
322 written.extend(variable_ids.iter().flatten().copied());
323 }
324 _ => {}
325 }
326 collect_variables(hir, std::slice::from_ref(stmt), &mut written);
327 for variable_id in written {
328 if ascending == Some(variable_id) {
329 if !cadence.contains(&variable_id) {
330 cadence.push(variable_id);
331 }
332 } else {
333 other_writes.insert(variable_id);
334 }
335 }
336 }
337 }
338 }
339}
340
341fn ascending_step<'hir>(expr: &'hir Expr<'hir>) -> Option<VariableId> {
343 match &expr.kind {
344 ExprKind::Unary(op, operand) if matches!(op.kind, UnOpKind::PreInc | UnOpKind::PostInc) => {
346 bare_identifier(operand)
347 }
348 ExprKind::Assign(lhs, Some(op), rhs)
350 if op.kind == hir::BinOpKind::Add && is_positive_literal(rhs) =>
351 {
352 bare_identifier(lhs)
353 }
354 ExprKind::Assign(lhs, None, rhs) => {
356 let target = bare_identifier(lhs)?;
357 let ExprKind::Binary(left, op, right) = &rhs.peel_parens().kind else { return None };
358 (op.kind == hir::BinOpKind::Add
359 && ((bare_identifier(left) == Some(target) && is_positive_literal(right))
360 || (is_positive_literal(left) && bare_identifier(right) == Some(target))))
361 .then_some(target)
362 }
363 _ => None,
364 }
365}
366
367fn bare_identifier(expr: &Expr<'_>) -> Option<VariableId> {
369 let ExprKind::Ident(resolutions) = &expr.peel_parens().kind else { return None };
370 resolutions.iter().find_map(|res| match res {
371 Res::Item(ItemId::Variable(variable_id)) => Some(*variable_id),
372 _ => None,
373 })
374}
375
376fn is_positive_literal(expr: &Expr<'_>) -> bool {
378 let ExprKind::Lit(lit) = &expr.peel_parens().kind else { return false };
379 matches!(&lit.kind, LitKind::Number(value) if !value.is_zero())
380}
381
382struct AtCollector<'a, 'hir> {
385 gcx: Gcx<'hir>,
386 hir: &'hir Hir<'hir>,
387 bindings: &'a Bindings,
388 cadence: &'a [VariableId],
389 iterated: Vec<Option<SetPath>>,
390}
391
392impl<'hir> Visit<'hir> for AtCollector<'_, 'hir> {
393 type BreakValue = Infallible;
394
395 fn hir(&self) -> &'hir Hir<'hir> {
396 self.hir
397 }
398
399 fn visit_expr(&mut self, expr: &'hir Expr<'hir>) -> ControlFlow<Self::BreakValue> {
400 if let Some(call) = enumerable_set_call(self.gcx, self.hir, self.bindings, expr)
401 && call.name == SetOp::At
402 && nth_argument(self.hir, call.function_id, call.args, call.index_arg, INDEX_PARAMETER)
403 .and_then(bare_identifier)
404 .is_some_and(|index| self.cadence.contains(&index))
405 {
406 self.iterated.push(call.set);
407 }
408 walk_literal_reachable_expr(self, expr)
409 }
410}
411
412struct RemoveScanner<'a, 'hir> {
414 gcx: Gcx<'hir>,
415 hir: &'hir Hir<'hir>,
416 bindings: &'a Bindings,
417 out: &'a mut Vec<(Option<SetPath>, Span)>,
418}
419
420impl<'hir> Visit<'hir> for RemoveScanner<'_, 'hir> {
421 type BreakValue = Infallible;
422
423 fn hir(&self) -> &'hir Hir<'hir> {
424 self.hir
425 }
426
427 fn visit_expr(&mut self, expr: &'hir Expr<'hir>) -> ControlFlow<Self::BreakValue> {
428 if let Some(call) = enumerable_set_call(self.gcx, self.hir, self.bindings, expr)
429 && call.name == SetOp::Remove
430 {
431 self.out.push((call.set, expr.span));
432 }
433 walk_literal_reachable_expr(self, expr)
434 }
435}
436
437fn walk_literal_reachable_expr<'hir, V>(
440 visitor: &mut V,
441 expr: &'hir Expr<'hir>,
442) -> ControlFlow<Infallible>
443where
444 V: Visit<'hir, BreakValue = Infallible>,
445{
446 match &expr.kind {
447 ExprKind::Binary(left, op, right)
448 if matches!(op.kind, hir::BinOpKind::And | hir::BinOpKind::Or) =>
449 {
450 visitor.visit_expr(left)?;
451 let short_circuits = matches!(
452 (op.kind, literal_bool(left)),
453 (hir::BinOpKind::And, Some(false)) | (hir::BinOpKind::Or, Some(true))
454 );
455 if !short_circuits {
456 visitor.visit_expr(right)?;
457 }
458 ControlFlow::Continue(())
459 }
460 ExprKind::Ternary(condition, true_expr, false_expr) => {
461 visitor.visit_expr(condition)?;
462 match literal_bool(condition) {
463 Some(true) => visitor.visit_expr(true_expr),
464 Some(false) => visitor.visit_expr(false_expr),
465 None => {
466 visitor.visit_expr(true_expr)?;
467 visitor.visit_expr(false_expr)
468 }
469 }
470 }
471 _ => visitor.walk_expr(expr),
472 }
473}
474
475fn literal_bool(expr: &Expr<'_>) -> Option<bool> {
477 let ExprKind::Lit(lit) = &expr.peel_parens().kind else { return None };
478 let LitKind::Bool(value) = lit.kind else { return None };
479 Some(value)
480}
481
482fn collect_variables<'hir>(
485 hir: &'hir Hir<'hir>,
486 stmts: &'hir [Stmt<'hir>],
487 out: &mut Vec<VariableId>,
488) {
489 struct Collector<'a, 'hir> {
490 hir: &'hir Hir<'hir>,
491 out: &'a mut Vec<VariableId>,
492 }
493 impl<'hir> Visit<'hir> for Collector<'_, 'hir> {
494 type BreakValue = Infallible;
495 fn hir(&self) -> &'hir Hir<'hir> {
496 self.hir
497 }
498 fn visit_expr(&mut self, expr: &'hir Expr<'hir>) -> ControlFlow<Self::BreakValue> {
499 let written = match &expr.kind {
500 ExprKind::Assign(lhs, ..) => Some(*lhs),
501 ExprKind::Delete(operand) => Some(*operand),
502 ExprKind::Unary(op, operand)
503 if matches!(
504 op.kind,
505 UnOpKind::PreInc | UnOpKind::PreDec | UnOpKind::PostInc | UnOpKind::PostDec
506 ) =>
507 {
508 Some(*operand)
509 }
510 _ => None,
511 };
512 if let Some(written) = written {
513 collect_lvalue_variables(written, self.out);
514 }
515 self.walk_expr(expr)
516 }
517 }
518 let mut collector = Collector { hir, out };
519 for stmt in stmts {
520 let _ = collector.visit_stmt(stmt);
521 }
522}
523
524fn collect_lvalue_variables(expr: &Expr<'_>, out: &mut Vec<VariableId>) {
527 match &expr.peel_parens().kind {
528 ExprKind::Ident(resolutions) => {
529 out.extend(resolutions.iter().filter_map(|res| match res {
530 Res::Item(ItemId::Variable(variable_id)) => Some(*variable_id),
531 _ => None,
532 }));
533 }
534 ExprKind::Tuple(exprs) => {
535 for expr in exprs.iter().flatten() {
536 collect_lvalue_variables(expr, out);
537 }
538 }
539 _ => {}
540 }
541}
542
543#[derive(PartialEq, Eq, Clone, Copy)]
544enum SetOp {
545 At,
546 Remove,
547}
548
549const INDEX_PARAMETER: usize = 1;
551
552struct SetCall<'hir> {
556 name: SetOp,
557 function_id: FunctionId,
558 args: &'hir CallArgs<'hir>,
559 set: Option<SetPath>,
560 index_arg: usize,
561}
562
563fn enumerable_set_call<'hir>(
567 gcx: Gcx<'hir>,
568 hir: &'hir Hir<'hir>,
569 bindings: &Bindings,
570 expr: &'hir Expr<'hir>,
571) -> Option<SetCall<'hir>> {
572 let ExprKind::Call(callee, args, _) = &expr.kind else { return None };
573 let ty = gcx.type_of_expr(callee.peel_parens().id)?;
574 let TyKind::Fn(function_ty) = ty.kind else { return None };
575 let function_id = function_ty.function_id?;
576 let function = hir.function(function_id);
577 let contract = hir.contract(function.contract?);
578 if !contract.kind.is_library() || contract.name.as_str() != "EnumerableSet" {
579 return None;
580 }
581 let name = match function.name?.as_str() {
582 "at" => SetOp::At,
583 "remove" => SetOp::Remove,
584 _ => return None,
585 };
586 let (set_expr, index_arg) = match &callee.peel_parens().kind {
589 ExprKind::Member(receiver, _) if is_enumerable_set_value(gcx, hir, receiver) => {
590 (Some(&**receiver), 0)
591 }
592 _ => (nth_argument(hir, function_id, args, 0, 0), 1),
593 };
594 let set = set_expr.and_then(|expr| set_path(hir, expr, bindings, &mut Vec::new()));
595 Some(SetCall { name, function_id, args, set, index_arg })
596}
597
598#[derive(PartialEq, Eq, Clone, Copy)]
600enum Step {
601 Field(Symbol),
602 Key(U256),
603}
604
605#[derive(PartialEq, Eq, Clone)]
609struct SetPath {
610 base: VariableId,
611 steps: Vec<Step>,
612}
613
614type Bindings = [(VariableId, Option<SetPath>)];
617
618fn set_path(
622 hir: &Hir<'_>,
623 expr: &Expr<'_>,
624 bindings: &Bindings,
625 seen: &mut Vec<VariableId>,
626) -> Option<SetPath> {
627 match &expr.peel_parens().kind {
628 ExprKind::Ident(resolutions) => {
629 let variable_id = resolutions.iter().find_map(|res| match res {
630 Res::Item(ItemId::Variable(variable_id)) => Some(*variable_id),
631 _ => None,
632 })?;
633 if seen.contains(&variable_id) {
634 return None;
635 }
636 seen.push(variable_id);
637 let variable = hir.variable(variable_id);
638 if matches!(variable.kind, VarKind::Statement) {
644 if let Some((_, binding)) =
645 bindings.iter().rev().find(|(bound, _)| *bound == variable_id)
646 {
647 return binding.clone();
648 }
649 if let Some(initializer) = variable.initializer {
650 return set_path(hir, initializer, bindings, seen);
651 }
652 return None;
655 }
656 Some(SetPath { base: variable_id, steps: Vec::new() })
657 }
658 ExprKind::Member(base, field) => {
659 let mut path = set_path(hir, base, bindings, seen)?;
660 path.steps.push(Step::Field(field.name));
661 Some(path)
662 }
663 ExprKind::Index(base, Some(index)) => {
664 let ExprKind::Lit(lit) = &index.peel_parens().kind else { return None };
665 let LitKind::Number(key) = &lit.kind else { return None };
666 let mut path = set_path(hir, base, bindings, seen)?;
667 path.steps.push(Step::Key(*key));
668 Some(path)
669 }
670 _ => None,
671 }
672}
673
674fn paths_alias(removed: &Option<SetPath>, iterated: &Option<SetPath>) -> bool {
677 match (removed, iterated) {
678 (Some(removed), Some(iterated)) => removed == iterated,
679 _ => true,
680 }
681}
682
683fn nth_argument<'hir>(
687 hir: &'hir Hir<'hir>,
688 function_id: FunctionId,
689 args: &'hir CallArgs<'hir>,
690 arg: usize,
691 parameter: usize,
692) -> Option<&'hir Expr<'hir>> {
693 match &args.kind {
694 CallArgsKind::Unnamed(exprs) => exprs.get(arg),
695 CallArgsKind::Named(named) => {
696 let parameter = *hir.function(function_id).parameters.get(parameter)?;
697 let name = hir.variable(parameter).name?;
698 named
699 .iter()
700 .find(|argument| argument.name.as_str() == name.as_str())
701 .map(|argument| &argument.value)
702 }
703 }
704}
705
706fn is_enumerable_set_value<'hir>(
710 gcx: Gcx<'hir>,
711 hir: &'hir Hir<'hir>,
712 receiver: &Expr<'_>,
713) -> bool {
714 let Some(ty) = gcx.type_of_expr(receiver.peel_parens().id) else { return false };
715 let TyKind::Struct(id) = ty.peel_refs().kind else { return false };
716 let Some(contract_id) = hir.strukt(id).contract else { return false };
717 hir.contract(contract_id).name.as_str() == "EnumerableSet"
718}