1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
#![allow(clippy::similar_names)] // `expr` and `expn`

use crate::visitors::{for_each_expr, Descend};

use arrayvec::ArrayVec;
use rustc_ast::{FormatArgs, FormatArgument, FormatPlaceholder};
use rustc_data_structures::fx::FxHashMap;
use rustc_hir::{self as hir, Expr, ExprKind, HirId, Node, QPath};
use rustc_lint::LateContext;
use rustc_span::def_id::DefId;
use rustc_span::hygiene::{self, MacroKind, SyntaxContext};
use rustc_span::{sym, BytePos, ExpnData, ExpnId, ExpnKind, Span, SpanData, Symbol};
use std::cell::OnceCell;
use std::ops::ControlFlow;
use std::rc::Rc;
use std::sync::atomic::{AtomicBool, Ordering};

const FORMAT_MACRO_DIAG_ITEMS: &[Symbol] = &[
    sym::assert_eq_macro,
    sym::assert_macro,
    sym::assert_ne_macro,
    sym::debug_assert_eq_macro,
    sym::debug_assert_macro,
    sym::debug_assert_ne_macro,
    sym::eprint_macro,
    sym::eprintln_macro,
    sym::format_args_macro,
    sym::format_macro,
    sym::print_macro,
    sym::println_macro,
    sym::std_panic_macro,
    sym::write_macro,
    sym::writeln_macro,
];

/// Returns true if a given Macro `DefId` is a format macro (e.g. `println!`)
pub fn is_format_macro(cx: &LateContext<'_>, macro_def_id: DefId) -> bool {
    if let Some(name) = cx.tcx.get_diagnostic_name(macro_def_id) {
        FORMAT_MACRO_DIAG_ITEMS.contains(&name)
    } else {
        false
    }
}

/// A macro call, like `vec![1, 2, 3]`.
///
/// Use `tcx.item_name(macro_call.def_id)` to get the macro name.
/// Even better is to check if it is a diagnostic item.
///
/// This structure is similar to `ExpnData` but it precludes desugaring expansions.
#[derive(Debug)]
pub struct MacroCall {
    /// Macro `DefId`
    pub def_id: DefId,
    /// Kind of macro
    pub kind: MacroKind,
    /// The expansion produced by the macro call
    pub expn: ExpnId,
    /// Span of the macro call site
    pub span: Span,
}

impl MacroCall {
    pub fn is_local(&self) -> bool {
        span_is_local(self.span)
    }
}

/// Returns an iterator of expansions that created the given span
pub fn expn_backtrace(mut span: Span) -> impl Iterator<Item = (ExpnId, ExpnData)> {
    std::iter::from_fn(move || {
        let ctxt = span.ctxt();
        if ctxt == SyntaxContext::root() {
            return None;
        }
        let expn = ctxt.outer_expn();
        let data = expn.expn_data();
        span = data.call_site;
        Some((expn, data))
    })
}

/// Checks whether the span is from the root expansion or a locally defined macro
pub fn span_is_local(span: Span) -> bool {
    !span.from_expansion() || expn_is_local(span.ctxt().outer_expn())
}

/// Checks whether the expansion is the root expansion or a locally defined macro
pub fn expn_is_local(expn: ExpnId) -> bool {
    if expn == ExpnId::root() {
        return true;
    }
    let data = expn.expn_data();
    let backtrace = expn_backtrace(data.call_site);
    std::iter::once((expn, data))
        .chain(backtrace)
        .find_map(|(_, data)| data.macro_def_id)
        .map_or(true, DefId::is_local)
}

/// Returns an iterator of macro expansions that created the given span.
/// Note that desugaring expansions are skipped.
pub fn macro_backtrace(span: Span) -> impl Iterator<Item = MacroCall> {
    expn_backtrace(span).filter_map(|(expn, data)| match data {
        ExpnData {
            kind: ExpnKind::Macro(kind, _),
            macro_def_id: Some(def_id),
            call_site: span,
            ..
        } => Some(MacroCall {
            def_id,
            kind,
            expn,
            span,
        }),
        _ => None,
    })
}

/// If the macro backtrace of `span` has a macro call at the root expansion
/// (i.e. not a nested macro call), returns `Some` with the `MacroCall`
pub fn root_macro_call(span: Span) -> Option<MacroCall> {
    macro_backtrace(span).last()
}

/// Like [`root_macro_call`], but only returns `Some` if `node` is the "first node"
/// produced by the macro call, as in [`first_node_in_macro`].
pub fn root_macro_call_first_node(cx: &LateContext<'_>, node: &impl HirNode) -> Option<MacroCall> {
    if first_node_in_macro(cx, node) != Some(ExpnId::root()) {
        return None;
    }
    root_macro_call(node.span())
}

/// Like [`macro_backtrace`], but only returns macro calls where `node` is the "first node" of the
/// macro call, as in [`first_node_in_macro`].
pub fn first_node_macro_backtrace(cx: &LateContext<'_>, node: &impl HirNode) -> impl Iterator<Item = MacroCall> {
    let span = node.span();
    first_node_in_macro(cx, node)
        .into_iter()
        .flat_map(move |expn| macro_backtrace(span).take_while(move |macro_call| macro_call.expn != expn))
}

/// If `node` is the "first node" in a macro expansion, returns `Some` with the `ExpnId` of the
/// macro call site (i.e. the parent of the macro expansion). This generally means that `node`
/// is the outermost node of an entire macro expansion, but there are some caveats noted below.
/// This is useful for finding macro calls while visiting the HIR without processing the macro call
/// at every node within its expansion.
///
/// If you already have immediate access to the parent node, it is simpler to
/// just check the context of that span directly (e.g. `parent.span.from_expansion()`).
///
/// If a macro call is in statement position, it expands to one or more statements.
/// In that case, each statement *and* their immediate descendants will all yield `Some`
/// with the `ExpnId` of the containing block.
///
/// A node may be the "first node" of multiple macro calls in a macro backtrace.
/// The expansion of the outermost macro call site is returned in such cases.
pub fn first_node_in_macro(cx: &LateContext<'_>, node: &impl HirNode) -> Option<ExpnId> {
    // get the macro expansion or return `None` if not found
    // `macro_backtrace` importantly ignores desugaring expansions
    let expn = macro_backtrace(node.span()).next()?.expn;

    // get the parent node, possibly skipping over a statement
    // if the parent is not found, it is sensible to return `Some(root)`
    let hir = cx.tcx.hir();
    let mut parent_iter = hir.parent_iter(node.hir_id());
    let (parent_id, _) = match parent_iter.next() {
        None => return Some(ExpnId::root()),
        Some((_, Node::Stmt(_))) => match parent_iter.next() {
            None => return Some(ExpnId::root()),
            Some(next) => next,
        },
        Some(next) => next,
    };

    // get the macro expansion of the parent node
    let parent_span = hir.span(parent_id);
    let Some(parent_macro_call) = macro_backtrace(parent_span).next() else {
        // the parent node is not in a macro
        return Some(ExpnId::root());
    };

    if parent_macro_call.expn.is_descendant_of(expn) {
        // `node` is input to a macro call
        return None;
    }

    Some(parent_macro_call.expn)
}

/* Specific Macro Utils */

/// Is `def_id` of `std::panic`, `core::panic` or any inner implementation macros
pub fn is_panic(cx: &LateContext<'_>, def_id: DefId) -> bool {
    let Some(name) = cx.tcx.get_diagnostic_name(def_id) else {
        return false;
    };
    matches!(
        name,
        sym::core_panic_macro
            | sym::std_panic_macro
            | sym::core_panic_2015_macro
            | sym::std_panic_2015_macro
            | sym::core_panic_2021_macro
    )
}

/// Is `def_id` of `assert!` or `debug_assert!`
pub fn is_assert_macro(cx: &LateContext<'_>, def_id: DefId) -> bool {
    let Some(name) = cx.tcx.get_diagnostic_name(def_id) else {
        return false;
    };
    matches!(name, sym::assert_macro | sym::debug_assert_macro)
}

#[derive(Debug)]
pub enum PanicExpn<'a> {
    /// No arguments - `panic!()`
    Empty,
    /// A string literal or any `&str` - `panic!("message")` or `panic!(message)`
    Str(&'a Expr<'a>),
    /// A single argument that implements `Display` - `panic!("{}", object)`
    Display(&'a Expr<'a>),
    /// Anything else - `panic!("error {}: {}", a, b)`
    Format(&'a Expr<'a>),
}

impl<'a> PanicExpn<'a> {
    pub fn parse(expr: &'a Expr<'a>) -> Option<Self> {
        let ExprKind::Call(callee, args) = &expr.kind else {
            return None;
        };
        let ExprKind::Path(QPath::Resolved(_, path)) = &callee.kind else {
            return None;
        };
        let name = path.segments.last().unwrap().ident.as_str();

        // This has no argument
        if name == "panic_cold_explicit" {
            return Some(Self::Empty);
        };

        let [arg, rest @ ..] = args else {
            return None;
        };
        let result = match name {
            "panic" if arg.span.ctxt() == expr.span.ctxt() => Self::Empty,
            "panic" | "panic_str" => Self::Str(arg),
            "panic_display" | "panic_cold_display" => {
                let ExprKind::AddrOf(_, _, e) = &arg.kind else {
                    return None;
                };
                Self::Display(e)
            },
            "panic_fmt" => Self::Format(arg),
            // Since Rust 1.52, `assert_{eq,ne}` macros expand to use:
            // `core::panicking::assert_failed(.., left_val, right_val, None | Some(format_args!(..)));`
            "assert_failed" => {
                // It should have 4 arguments in total (we already matched with the first argument,
                // so we're just checking for 3)
                if rest.len() != 3 {
                    return None;
                }
                // `msg_arg` is either `None` (no custom message) or `Some(format_args!(..))` (custom message)
                let msg_arg = &rest[2];
                match msg_arg.kind {
                    ExprKind::Call(_, [fmt_arg]) => Self::Format(fmt_arg),
                    _ => Self::Empty,
                }
            },
            _ => return None,
        };
        Some(result)
    }
}

/// Finds the arguments of an `assert!` or `debug_assert!` macro call within the macro expansion
pub fn find_assert_args<'a>(
    cx: &LateContext<'_>,
    expr: &'a Expr<'a>,
    expn: ExpnId,
) -> Option<(&'a Expr<'a>, PanicExpn<'a>)> {
    find_assert_args_inner(cx, expr, expn).map(|([e], mut p)| {
        // `assert!(..)` expands to `core::panicking::panic("assertion failed: ...")` (which we map to
        // `PanicExpn::Str(..)`) and `assert!(.., "..")` expands to
        // `core::panicking::panic_fmt(format_args!(".."))` (which we map to `PanicExpn::Format(..)`).
        // So even we got `PanicExpn::Str(..)` that means there is no custom message provided
        if let PanicExpn::Str(_) = p {
            p = PanicExpn::Empty;
        }

        (e, p)
    })
}

/// Finds the arguments of an `assert_eq!` or `debug_assert_eq!` macro call within the macro
/// expansion
pub fn find_assert_eq_args<'a>(
    cx: &LateContext<'_>,
    expr: &'a Expr<'a>,
    expn: ExpnId,
) -> Option<(&'a Expr<'a>, &'a Expr<'a>, PanicExpn<'a>)> {
    find_assert_args_inner(cx, expr, expn).map(|([a, b], p)| (a, b, p))
}

fn find_assert_args_inner<'a, const N: usize>(
    cx: &LateContext<'_>,
    expr: &'a Expr<'a>,
    expn: ExpnId,
) -> Option<([&'a Expr<'a>; N], PanicExpn<'a>)> {
    let macro_id = expn.expn_data().macro_def_id?;
    let (expr, expn) = match cx.tcx.item_name(macro_id).as_str().strip_prefix("debug_") {
        None => (expr, expn),
        Some(inner_name) => find_assert_within_debug_assert(cx, expr, expn, Symbol::intern(inner_name))?,
    };
    let mut args = ArrayVec::new();
    let panic_expn = for_each_expr(expr, |e| {
        if args.is_full() {
            match PanicExpn::parse(e) {
                Some(expn) => ControlFlow::Break(expn),
                None => ControlFlow::Continue(Descend::Yes),
            }
        } else if is_assert_arg(cx, e, expn) {
            args.push(e);
            ControlFlow::Continue(Descend::No)
        } else {
            ControlFlow::Continue(Descend::Yes)
        }
    });
    let args = args.into_inner().ok()?;
    // if no `panic!(..)` is found, use `PanicExpn::Empty`
    // to indicate that the default assertion message is used
    let panic_expn = panic_expn.unwrap_or(PanicExpn::Empty);
    Some((args, panic_expn))
}

fn find_assert_within_debug_assert<'a>(
    cx: &LateContext<'_>,
    expr: &'a Expr<'a>,
    expn: ExpnId,
    assert_name: Symbol,
) -> Option<(&'a Expr<'a>, ExpnId)> {
    for_each_expr(expr, |e| {
        if !e.span.from_expansion() {
            return ControlFlow::Continue(Descend::No);
        }
        let e_expn = e.span.ctxt().outer_expn();
        if e_expn == expn {
            ControlFlow::Continue(Descend::Yes)
        } else if e_expn.expn_data().macro_def_id.map(|id| cx.tcx.item_name(id)) == Some(assert_name) {
            ControlFlow::Break((e, e_expn))
        } else {
            ControlFlow::Continue(Descend::No)
        }
    })
}

fn is_assert_arg(cx: &LateContext<'_>, expr: &Expr<'_>, assert_expn: ExpnId) -> bool {
    if !expr.span.from_expansion() {
        return true;
    }
    let result = macro_backtrace(expr.span).try_for_each(|macro_call| {
        if macro_call.expn == assert_expn {
            ControlFlow::Break(false)
        } else {
            match cx.tcx.item_name(macro_call.def_id) {
                // `cfg!(debug_assertions)` in `debug_assert!`
                sym::cfg => ControlFlow::Continue(()),
                // assert!(other_macro!(..))
                _ => ControlFlow::Break(true),
            }
        }
    });
    match result {
        ControlFlow::Break(is_assert_arg) => is_assert_arg,
        ControlFlow::Continue(()) => true,
    }
}

thread_local! {
    /// We preserve the [`FormatArgs`] structs from the early pass for use in the late pass to be
    /// able to access the many features of a [`LateContext`].
    ///
    /// A thread local is used because [`FormatArgs`] is `!Send` and `!Sync`, we are making an
    /// assumption that the early pass that populates the map and the later late passes will all be
    /// running on the same thread.
    #[doc(hidden)]
    pub static AST_FORMAT_ARGS: OnceCell<FxHashMap<Span, Rc<FormatArgs>>> = {
        static CALLED: AtomicBool = AtomicBool::new(false);
        debug_assert!(
            !CALLED.swap(true, Ordering::SeqCst),
            "incorrect assumption: `AST_FORMAT_ARGS` should only be accessed by a single thread",
        );

        OnceCell::new()
    };
}

/// Returns an AST [`FormatArgs`] node if a `format_args` expansion is found as a descendant of
/// `expn_id`
pub fn find_format_args(cx: &LateContext<'_>, start: &Expr<'_>, expn_id: ExpnId) -> Option<Rc<FormatArgs>> {
    let format_args_expr = for_each_expr(start, |expr| {
        let ctxt = expr.span.ctxt();
        if ctxt.outer_expn().is_descendant_of(expn_id) {
            if macro_backtrace(expr.span)
                .map(|macro_call| cx.tcx.item_name(macro_call.def_id))
                .any(|name| matches!(name, sym::const_format_args | sym::format_args | sym::format_args_nl))
            {
                ControlFlow::Break(expr)
            } else {
                ControlFlow::Continue(Descend::Yes)
            }
        } else {
            ControlFlow::Continue(Descend::No)
        }
    })?;

    AST_FORMAT_ARGS.with(|ast_format_args| {
        ast_format_args
            .get()?
            .get(&format_args_expr.span.with_parent(None))
            .map(Rc::clone)
    })
}

/// Attempt to find the [`rustc_hir::Expr`] that corresponds to the [`FormatArgument`]'s value, if
/// it cannot be found it will return the [`rustc_ast::Expr`].
pub fn find_format_arg_expr<'hir, 'ast>(
    start: &'hir Expr<'hir>,
    target: &'ast FormatArgument,
) -> Result<&'hir rustc_hir::Expr<'hir>, &'ast rustc_ast::Expr> {
    let SpanData {
        lo,
        hi,
        ctxt,
        parent: _,
    } = target.expr.span.data();

    for_each_expr(start, |expr| {
        // When incremental compilation is enabled spans gain a parent during AST to HIR lowering,
        // since we're comparing an AST span to a HIR one we need to ignore the parent field
        let data = expr.span.data();
        if data.lo == lo && data.hi == hi && data.ctxt == ctxt {
            ControlFlow::Break(expr)
        } else {
            ControlFlow::Continue(())
        }
    })
    .ok_or(&target.expr)
}

/// Span of the `:` and format specifiers
///
/// ```ignore
/// format!("{:.}"), format!("{foo:.}")
///           ^^                  ^^
/// ```
pub fn format_placeholder_format_span(placeholder: &FormatPlaceholder) -> Option<Span> {
    let base = placeholder.span?.data();

    // `base.hi` is `{...}|`, subtract 1 byte (the length of '}') so that it points before the closing
    // brace `{...|}`
    Some(Span::new(
        placeholder.argument.span?.hi(),
        base.hi - BytePos(1),
        base.ctxt,
        base.parent,
    ))
}

/// Span covering the format string and values
///
/// ```ignore
/// format("{}.{}", 10, 11)
/// //     ^^^^^^^^^^^^^^^
/// ```
pub fn format_args_inputs_span(format_args: &FormatArgs) -> Span {
    match format_args.arguments.explicit_args() {
        [] => format_args.span,
        [.., last] => format_args
            .span
            .to(hygiene::walk_chain(last.expr.span, format_args.span.ctxt())),
    }
}

/// Returns the [`Span`] of the value at `index` extended to the previous comma, e.g. for the value
/// `10`
///
/// ```ignore
/// format("{}.{}", 10, 11)
/// //            ^^^^
/// ```
pub fn format_arg_removal_span(format_args: &FormatArgs, index: usize) -> Option<Span> {
    let ctxt = format_args.span.ctxt();

    let current = hygiene::walk_chain(format_args.arguments.by_index(index)?.expr.span, ctxt);

    let prev = if index == 0 {
        format_args.span
    } else {
        hygiene::walk_chain(format_args.arguments.by_index(index - 1)?.expr.span, ctxt)
    };

    Some(current.with_lo(prev.hi()))
}

/// Where a format parameter is being used in the format string
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum FormatParamUsage {
    /// Appears as an argument, e.g. `format!("{}", foo)`
    Argument,
    /// Appears as a width, e.g. `format!("{:width$}", foo, width = 1)`
    Width,
    /// Appears as a precision, e.g. `format!("{:.precision$}", foo, precision = 1)`
    Precision,
}

/// A node with a `HirId` and a `Span`
pub trait HirNode {
    fn hir_id(&self) -> HirId;
    fn span(&self) -> Span;
}

macro_rules! impl_hir_node {
    ($($t:ident),*) => {
        $(impl HirNode for hir::$t<'_> {
            fn hir_id(&self) -> HirId {
                self.hir_id
            }
            fn span(&self) -> Span {
                self.span
            }
        })*
    };
}

impl_hir_node!(Expr, Pat);

impl HirNode for hir::Item<'_> {
    fn hir_id(&self) -> HirId {
        self.hir_id()
    }

    fn span(&self) -> Span {
        self.span
    }
}