Skip to main content

forge_lint/sol/high/
enumerable_loop_removal.rs

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
28// The detector reports only the shape it can judge without a flow analysis: a loop whose own
29// index is written exclusively by simple unconditional increments, reads the set with `at` at
30// that bare index, and removes from the same set in a straight-line body. Control flow,
31// descending traversal, composite indices, and other shapes that need value or path reasoning
32// are deliberately unreported, even when they corrupt iteration. Set operands that cannot be
33// identified statically are conservatively treated as possible aliases, so false positives are
34// still possible on set identity.
35
36impl<'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
52/// Walks a function body in statement order and, for each loop, flags the EnumerableSet `remove`
53/// calls that corrupt that loop's own iteration. The walk keeps, at every point, what each local
54/// `storage` reference last named, so each loop is judged against the bindings standing where it
55/// runs rather than against every binding of the function.
56struct LoopFinder<'ctx, 's, 'c, 'hir> {
57    gcx: Gcx<'hir>,
58    hir: &'hir Hir<'hir>,
59    ctx: &'ctx LintContext<'s, 'c>,
60    /// What each local `storage` reference names where the walk stands, the latest entry
61    /// winning: the path its last straight-line binding resolved to, or `None` once a write
62    /// leaves it without one answer, from a conditional branch, a loop body, or a shape the
63    /// analysis cannot read.
64    bindings: Vec<(VariableId, Option<SetPath>)>,
65    // A loop nested in a flagged loop sees the same calls: dedupe emissions by span.
66    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        // A `for` desugars to `Block { init; Loop(For) }`; its index lives partly in the init,
78        // which runs once, on the straight line entering the loop.
79        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            // A bare block runs on the straight line: what it binds stays bound past it.
90            StmtKind::Block(block) | StmtKind::UncheckedBlock(block) => {
91                self.walk_body(block.stmts);
92            }
93            // A bare loop is a `while`, a `do-while`, or a `for` with no init.
94            StmtKind::Loop(body, source) => self.enter_loop(&[], body.stmts, *source),
95            StmtKind::If(_, then, else_) => {
96                // The condition and either branch may write a reference, and which of them ran
97                // is not read here: everything the statement writes stops naming one thing, and
98                // what a branch binds for its own statements ends with the branch.
99                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                // Clauses are branches the same way.
110                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    /// Analyzes one loop, then walks inside it for the nested ones. A write anywhere in the
122    /// loop may have run on an earlier turn by the time any statement of it runs again, so
123    /// everything the loop writes, init included, stops naming one thing before the loop is
124    /// judged, and stays so past it.
125    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        // Analyze the user-written body, peeled out of the synthetic condition guard the lowering
134        // wraps every loop in (`if (cond) { body } else break`), so the guard's `break` and the
135        // `for`'s next-step are read for what they are, not as user control flow.
136        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    /// Applies one straight-line statement to the bindings. Everything it writes stops naming
143    /// one thing first; a declaration or a plain assignment then binds its reference to what
144    /// the right-hand side names right here, resolved eagerly because a later write to a
145    /// reference the right-hand side reads must not reach back into this binding.
146    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    /// Marks everything the statements write as no longer naming one thing.
172    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    /// Flags the removals in `body` that corrupt this loop's iteration, under the shrunk rule:
181    /// an unconditional ascending cadence, a straight-line body, and a `remove` on a set the
182    /// loop reads with `at` at that cadence.
183    fn analyze_loop(&mut self, body: &'hir [Stmt<'hir>]) {
184        // A control-flow construct in the body would make the corruption depend on where
185        // control flows, which this detector does not track. Stay silent instead of guessing.
186        if !body_is_straight_line(body) {
187            return;
188        }
189        // The loop's own index must step upward on the straight line of the body, never under a
190        // branch. Without one, the iteration order is not the ascending walk the swap-and-pop
191        // corruption needs.
192        let cadence = ascending_cadence(self.hir, body);
193        if cadence.is_empty() {
194            return;
195        }
196        // Which sets this loop iterates with `at` at that cadence.
197        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        // Every reachable removal in the body, including calls nested in short-circuit or ternary
211        // expressions. Statement-level control flow was rejected above, and expression-level
212        // reachability is folded only when a literal condition proves an arm cannot execute.
213        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
232/// The user-written body of a loop, peeled out of the synthetic condition guard the AST lowering
233/// wraps it in. A `for`/`while` lowers to a single `if (cond) { body } else break`, its `body`
234/// holding the user statements (and, for a `for`, the next-step); a `do-while` appends an
235/// `if (cond) continue else break` after the user statements. Peeling these makes the guard's
236/// `break`/`continue` and the next-step read for what they are, not as user control flow. A body
237/// that does not match the exact synthetic shape is returned unchanged.
238const 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
262/// Whether every statement of a loop body runs on one straight line: no branch (`if`/`try`), no
263/// jump (`break`/`continue`), no terminal statement, no inline assembly, and no nested loop. Bare
264/// blocks are transparent. Any of these would let control skip a removal, skip the cadence step,
265/// or leave the loop before a shifted slot is read, none of which this detector tracks, so their
266/// presence makes it silent.
267fn 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
285/// The loop's own indices that step upward unconditionally: a bare identifier advanced by `i++`,
286/// `i += <positive literal>`, or `i = i + <positive literal>` (with either addition order) as a
287/// straight-line statement of the body (a `for`'s desugared post-step lands here, a `while`'s
288/// in-body counter too). A step under a branch, a reset (`i = 0`), a no-op (`i += 0`), a decrement,
289/// or composite arithmetic (`i = (i + 2) - 1`) does not qualify: the walk is only known to ascend
290/// for the simple forms.
291fn 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
299/// Walks the straight-line statements of a body, bare blocks included, and records variables
300/// written by a supported ascending step separately from every other write or declaration. A
301/// cadence is valid only when all of its writes are supported ascending steps.
302fn 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
341/// The bare identifier an expression steps upward, if it is one of the simple ascending forms.
342fn ascending_step<'hir>(expr: &'hir Expr<'hir>) -> Option<VariableId> {
343    match &expr.kind {
344        // `i++` / `++i`.
345        ExprKind::Unary(op, operand) if matches!(op.kind, UnOpKind::PreInc | UnOpKind::PostInc) => {
346            bare_identifier(operand)
347        }
348        // `i += <positive literal>`.
349        ExprKind::Assign(lhs, Some(op), rhs)
350            if op.kind == hir::BinOpKind::Add && is_positive_literal(rhs) =>
351        {
352            bare_identifier(lhs)
353        }
354        // `i = i + <positive literal>` / `i = <positive literal> + i`.
355        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
367/// The variable a bare identifier expression resolves to, or `None` for anything else.
368fn 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
376/// Whether an expression is a positive integer literal: `1`, `2`, never `0`.
377fn 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
382/// Collects the sets a loop iterates with `at` at its ascending cadence. The index must be the
383/// bare cadence identifier; copies and composite expressions are deliberately left out.
384struct 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
412/// Scans a straight-line body for EnumerableSet `remove` calls, with the span to report.
413struct 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
437/// Walks the children of an expression while skipping short-circuit and ternary arms that a
438/// literal boolean condition proves unreachable. Unknown conditions keep both possible arms.
439fn 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
475/// The value of a bare boolean literal expression.
476fn 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
482/// The variables a statement list writes through expressions, through the loops under it as well:
483/// assignments (tuple targets included), increments, decrements, and deletes, wherever they sit.
484fn 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
524/// Collects the variables written through an assignment target. Tuple assignment targets can
525/// contain several identifiers; member and indexed targets do not write the base variable itself.
526fn 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
549/// `at`'s index is its second parameter, wherever the call form puts the argument.
550const INDEX_PARAMETER: usize = 1;
551
552/// A resolved EnumerableSet call: which operation, on which set, and where its index sits. The
553/// method form binds the set to the receiver, so its arguments start one position later than the
554/// parameters they fill.
555struct SetCall<'hir> {
556    name: SetOp,
557    function_id: FunctionId,
558    args: &'hir CallArgs<'hir>,
559    set: Option<SetPath>,
560    index_arg: usize,
561}
562
563/// The EnumerableSet `at` or `remove` a call dispatches to. Resolving through the type checker
564/// covers the `using for` method form, the library-qualified form and import aliases. The library
565/// is identified only by its kind and exact `EnumerableSet` name, not its source or behavior.
566fn 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    // The set operand is the bound receiver in the method form and the first argument in the
587    // library-qualified form; the index of `at` sits right after it.
588    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/// One step of the storage path naming a set: a struct field or a literal mapping key.
599#[derive(PartialEq, Eq, Clone, Copy)]
600enum Step {
601    Field(Symbol),
602    Key(U256),
603}
604
605/// The storage location a set expression names: a base variable and the steps taken from it.
606/// `holders`, `pair.a` and `sets[1]` each name one, and two of them are the same set exactly
607/// when they are the same path.
608#[derive(PartialEq, Eq, Clone)]
609struct SetPath {
610    base: VariableId,
611    steps: Vec<Step>,
612}
613
614/// What each local `storage` reference names at the point being analyzed, the latest entry
615/// winning; `None` marks a reference no straight-line reading gives one answer for.
616type Bindings = [(VariableId, Option<SetPath>)];
617
618/// The path a set expression names, or `None` when it cannot be read: an index that varies, a
619/// call result, a reference without one straight-line binding, anything the analysis would have
620/// to evaluate.
621fn 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            // A local `storage` reference is another name for the set its last binding gave it.
639            // The walk records what stands where the loop is judged, so a rebinding past the
640            // loop does not reach back into it; a reference declared inside the analyzed loop
641            // has no entry and is bound by its initializer anew each turn, read under the same
642            // bindings.
643            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                // Bound neither by the walk nor by an initializer, a tuple-destructured
653                // component: nothing says which set it names, so it may name any of them.
654                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
674/// Whether two set operands can name the same set. Two paths that can be read name the same set
675/// exactly when they are equal; a path that cannot be read may be either.
676fn 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
683/// The argument at position `arg` of a positional call, or the one a named call binds to the
684/// callee's parameter at position `parameter`. Named arguments come in source order, which is
685/// neither the parameter order nor, in the method form, the argument order.
686fn 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
706/// Whether `receiver` is a value of one of the set struct types declared in a library (or
707/// contract) named `EnumerableSet` (`AddressSet` / `UintSet` / `Bytes32Set`), which tells the
708/// bound method form apart from the library-qualified form.
709fn 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}