Skip to main content

forge_lint/sol/analysis/
stmts.rs

1//! Statement-shape probes over Solar HIR.
2
3use super::is_exit_call;
4use solar::{
5    ast::FunctionKind,
6    sema::{
7        Gcx,
8        hir::{self, Expr, FunctionId, LoopSource, Stmt, StmtKind, Visit},
9    },
10};
11use std::ops::ControlFlow;
12
13/// Runs `f` on every statement (pre-order, nested ones included) until it breaks.
14struct StmtVisitor<'gcx, F> {
15    hir: &'gcx hir::Hir<'gcx>,
16    f: F,
17}
18
19impl<'gcx, F: FnMut(&'gcx Stmt<'gcx>) -> ControlFlow<()>> Visit<'gcx> for StmtVisitor<'gcx, F> {
20    type BreakValue = ();
21
22    fn hir(&self) -> &'gcx hir::Hir<'gcx> {
23        self.hir
24    }
25
26    fn visit_stmt(&mut self, stmt: &'gcx Stmt<'gcx>) -> ControlFlow<()> {
27        (self.f)(stmt)?;
28        self.walk_stmt(stmt)
29    }
30}
31
32/// Runs `f` on every statement of `stmts` and their nested statements (pre-order) until it breaks.
33pub fn visit_stmts<'gcx>(
34    hir: &'gcx hir::Hir<'gcx>,
35    stmts: impl IntoIterator<Item = &'gcx Stmt<'gcx>>,
36    f: impl FnMut(&'gcx Stmt<'gcx>) -> ControlFlow<()>,
37) -> ControlFlow<()> {
38    let mut visitor = StmtVisitor { hir, f };
39    stmts.into_iter().try_for_each(|stmt| visitor.visit_stmt(stmt))
40}
41
42/// The expression directly owned by `stmt` (nested statements excluded).
43pub fn stmt_expr<'gcx>(
44    hir: &'gcx hir::Hir<'gcx>,
45    stmt: &'gcx Stmt<'gcx>,
46) -> Option<&'gcx Expr<'gcx>> {
47    match stmt.kind {
48        StmtKind::DeclSingle(var_id) => hir.variable(var_id).initializer,
49        StmtKind::DeclMulti(_, expr)
50        | StmtKind::Expr(expr)
51        | StmtKind::Emit(expr)
52        | StmtKind::Revert(expr)
53        | StmtKind::Return(Some(expr))
54        | StmtKind::If(expr, ..) => Some(expr),
55        StmtKind::Try(try_stmt) => Some(&try_stmt.expr),
56        _ => None,
57    }
58}
59
60/// True when executing `stmt` provably prevents control from continuing past it: `return`,
61/// `revert`, `selfdestruct`, `require(false, ..)` / `assert(false)`, a block containing any such
62/// statement, an `if` whose both arms exit, a `try` whose every clause exits, or a `do-while`
63/// whose body exits without `break`/`continue`.
64pub fn branch_always_exits(gcx: Gcx<'_>, stmt: &Stmt<'_>) -> bool {
65    match &stmt.kind {
66        StmtKind::Return(_) | StmtKind::Revert(_) => true,
67        StmtKind::Expr(expr) => is_exit_call(gcx, expr),
68        StmtKind::Block(b) | StmtKind::UncheckedBlock(b) => {
69            b.stmts.iter().any(|expr| branch_always_exits(gcx, expr))
70        }
71        StmtKind::If(_, t, Some(e)) => branch_always_exits(gcx, t) && branch_always_exits(gcx, e),
72        StmtKind::Loop(block, LoopSource::DoWhile) => {
73            let user = do_while_user_stmts(block.stmts);
74            !stmts_break_or_continue(user) && user.iter().any(|expr| branch_always_exits(gcx, expr))
75        }
76        StmtKind::Try(t) => {
77            !t.clauses.is_empty()
78                && t.clauses
79                    .iter()
80                    .all(|c| c.block.stmts.iter().any(|expr| branch_always_exits(gcx, expr)))
81        }
82        _ => false,
83    }
84}
85
86/// The `for` update statement of a loop, which runs after every iteration.
87pub const fn loop_update<'gcx>(source: LoopSource<'gcx>) -> Option<&'gcx Stmt<'gcx>> {
88    match source {
89        LoopSource::For { update } => update,
90        LoopSource::While | LoopSource::DoWhile => None,
91    }
92}
93
94/// The statements of one loop iteration: the body followed by the `for` update, if any.
95pub fn loop_stmts<'gcx>(
96    block: hir::Block<'gcx>,
97    source: LoopSource<'gcx>,
98) -> impl Iterator<Item = &'gcx Stmt<'gcx>> + Clone {
99    block.stmts.iter().chain(loop_update(source))
100}
101
102/// Number of `_` placeholders in `stmts`, recursing into nested control flow.
103pub fn count_placeholders(stmts: &[Stmt<'_>]) -> usize {
104    stmts.iter().map(count_placeholders_in_stmt).sum()
105}
106
107fn count_placeholders_in_stmt(stmt: &Stmt<'_>) -> usize {
108    match &stmt.kind {
109        StmtKind::Placeholder => 1,
110        StmtKind::Block(b) | StmtKind::UncheckedBlock(b) => count_placeholders(b.stmts),
111        StmtKind::Loop(b, source) => loop_stmts(*b, *source).map(count_placeholders_in_stmt).sum(),
112        StmtKind::If(_, t, e) => {
113            count_placeholders_in_stmt(t) + e.as_ref().map_or(0, |e| count_placeholders_in_stmt(e))
114        }
115        StmtKind::Try(t) => t.clauses.iter().map(|c| count_placeholders(c.block.stmts)).sum(),
116        _ => 0,
117    }
118}
119
120/// Collects the statements executed before the first placeholder of a modifier body, following
121/// nested blocks. Returns `None` when the placeholder is not reached unconditionally (e.g. it is
122/// inside an `if`, loop or `try`).
123pub fn stmts_before_placeholder<'a, 'gcx>(
124    stmts: &'a [Stmt<'gcx>],
125    out: &mut Vec<&'a Stmt<'gcx>>,
126) -> Option<()> {
127    for (i, stmt) in stmts.iter().enumerate() {
128        match &stmt.kind {
129            StmtKind::Placeholder => {
130                out.extend(&stmts[..i]);
131                return Some(());
132            }
133            StmtKind::Block(b) | StmtKind::UncheckedBlock(b) if count_placeholders(b.stmts) > 0 => {
134                out.extend(&stmts[..i]);
135                return stmts_before_placeholder(b.stmts, out);
136            }
137            _ if count_placeholders_in_stmt(stmt) > 0 => return None,
138            _ => {}
139        }
140    }
141    None
142}
143
144/// Strips the trailing `if (cond) break;` that lowers `do { ... } while (cond);`.
145pub fn do_while_user_stmts<'a, 'gcx>(stmts: &'a [Stmt<'gcx>]) -> &'a [Stmt<'gcx>] {
146    match stmts.split_last() {
147        Some((last, rest)) if is_loop_termination_if(last) => rest,
148        _ => stmts,
149    }
150}
151
152/// `if (...) break;` as synthesized by the `do-while` lowering.
153pub fn is_loop_termination_if(stmt: &Stmt<'_>) -> bool {
154    let StmtKind::If(_, t, e) = &stmt.kind else { return false };
155    is_break_stmt(t) || e.as_ref().is_some_and(|e| is_break_stmt(e))
156}
157
158/// `break`, possibly wrapped in single-statement blocks.
159pub fn is_break_stmt(stmt: &Stmt<'_>) -> bool {
160    match &stmt.kind {
161        StmtKind::Break => true,
162        StmtKind::Block(b) | StmtKind::UncheckedBlock(b) => {
163            b.stmts.len() == 1 && is_break_stmt(&b.stmts[0])
164        }
165        _ => false,
166    }
167}
168
169/// `break`/`continue` targeting the current loop (nested loops shadow them).
170pub fn stmts_break_or_continue(stmts: &[Stmt<'_>]) -> bool {
171    stmts.iter().any(|stmt| match &stmt.kind {
172        StmtKind::Break | StmtKind::Continue => true,
173        StmtKind::Block(b) | StmtKind::UncheckedBlock(b) => stmts_break_or_continue(b.stmts),
174        StmtKind::If(_, t, e) => {
175            stmts_break_or_continue(std::slice::from_ref(*t))
176                || e.is_some_and(|e| stmts_break_or_continue(std::slice::from_ref(e)))
177        }
178        StmtKind::Try(t) => t.clauses.iter().any(|c| stmts_break_or_continue(c.block.stmts)),
179        _ => false,
180    })
181}
182
183/// The statements a modifier runs before its unique `_;`, when that placeholder is reached
184/// unconditionally. `None` for non-modifiers, bodiless modifiers and conditional placeholders.
185pub fn modifier_prefix<'gcx>(
186    hir: &'gcx hir::Hir<'gcx>,
187    fid: FunctionId,
188) -> Option<Vec<&'gcx Stmt<'gcx>>> {
189    let modifier = hir.function(fid);
190    let body = modifier.body.filter(|_| modifier.kind == FunctionKind::Modifier)?;
191    if count_placeholders(body.stmts) != 1 {
192        return None;
193    }
194    let mut prefix = Vec::new();
195    stmts_before_placeholder(body.stmts, &mut prefix)?;
196    Some(prefix)
197}