Return a value from find_format_args instead of using a callback

Alexendoo · Alexendoo · commit d3cb51c9f183 · 2023-09-01T23:53:45.000Z
diff --git a/clippy_lints/src/explicit_write.rs b/clippy_lints/src/explicit_write.rs
@@ -57,54 +57,52 @@ impl<'tcx> LateLintPass<'tcx> for ExplicitWrite {
             } else {
                 None
             }
+            && let Some(format_args) = find_format_args(cx, write_arg, ExpnId::root())
         {
-            find_format_args(cx, write_arg, ExpnId::root(), |format_args| {
-                let calling_macro =
-                    // ordering is important here, since `writeln!` uses `write!` internally
-                    if is_expn_of(write_call.span, "writeln").is_some() {
-                        Some("writeln")
-                    } else if is_expn_of(write_call.span, "write").is_some() {
-                        Some("write")
-                    } else {
-                        None
-                    };
-                let prefix = if dest_name == "stderr" {
-                    "e"
-                } else {
-                    ""
-                };
+            // ordering is important here, since `writeln!` uses `write!` internally
+            let calling_macro = if is_expn_of(write_call.span, "writeln").is_some() {
+                Some("writeln")
+            } else if is_expn_of(write_call.span, "write").is_some() {
+                Some("write")
+            } else {
+                None
+            };
+            let prefix = if dest_name == "stderr" {
+                "e"
+            } else {
+                ""
+            };
 
-                // We need to remove the last trailing newline from the string because the
-                // underlying `fmt::write` function doesn't know whether `println!` or `print!` was
-                // used.
-                let (used, sugg_mac) = if let Some(macro_name) = calling_macro {
-                    (
-                        format!("{macro_name}!({dest_name}(), ...)"),
-                        macro_name.replace("write", "print"),
-                    )
-                } else {
-                    (
-                        format!("{dest_name}().write_fmt(...)"),
-                        "print".into(),
-                    )
-                };
-                let mut applicability = Applicability::MachineApplicable;
-                let inputs_snippet = snippet_with_applicability(
-                    cx,
-                    format_args_inputs_span(format_args),
-                    "..",
-                    &mut applicability,
-                );
-                span_lint_and_sugg(
-                    cx,
-                    EXPLICIT_WRITE,
-                    expr.span,
-                    &format!("use of `{used}.unwrap()`"),
-                    "try",
-                    format!("{prefix}{sugg_mac}!({inputs_snippet})"),
-                    applicability,
-                );
-            });
+            // We need to remove the last trailing newline from the string because the
+            // underlying `fmt::write` function doesn't know whether `println!` or `print!` was
+            // used.
+            let (used, sugg_mac) = if let Some(macro_name) = calling_macro {
+                (
+                    format!("{macro_name}!({dest_name}(), ...)"),
+                    macro_name.replace("write", "print"),
+                )
+            } else {
+                (
+                    format!("{dest_name}().write_fmt(...)"),
+                    "print".into(),
+                )
+            };
+            let mut applicability = Applicability::MachineApplicable;
+            let inputs_snippet = snippet_with_applicability(
+                cx,
+                format_args_inputs_span(format_args),
+                "..",
+                &mut applicability,
+            );
+            span_lint_and_sugg(
+                cx,
+                EXPLICIT_WRITE,
+                expr.span,
+                &format!("use of `{used}.unwrap()`"),
+                "try",
+                format!("{prefix}{sugg_mac}!({inputs_snippet})"),
+                applicability,
+            );
         }
     }
 }
diff --git a/clippy_lints/src/format.rs b/clippy_lints/src/format.rs
@@ -43,14 +43,10 @@ declare_lint_pass!(UselessFormat => [USELESS_FORMAT]);
 
 impl<'tcx> LateLintPass<'tcx> for UselessFormat {
     fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) {
-        let Some(macro_call) = root_macro_call_first_node(cx, expr) else {
-            return;
-        };
-        if !cx.tcx.is_diagnostic_item(sym::format_macro, macro_call.def_id) {
-            return;
-        }
-
-        find_format_args(cx, expr, macro_call.expn, |format_args| {
+        if let Some(macro_call) = root_macro_call_first_node(cx, expr)
+            && cx.tcx.is_diagnostic_item(sym::format_macro, macro_call.def_id)
+            && let Some(format_args) = find_format_args(cx, expr, macro_call.expn)
+        {
             let mut applicability = Applicability::MachineApplicable;
             let call_site = macro_call.span;
 
@@ -91,7 +87,7 @@ impl<'tcx> LateLintPass<'tcx> for UselessFormat {
                 },
                 _ => {},
             }
-        });
+        }
     }
 }
 
diff --git a/clippy_lints/src/format_args.rs b/clippy_lints/src/format_args.rs
@@ -186,15 +186,10 @@ impl FormatArgs {
 
 impl<'tcx> LateLintPass<'tcx> for FormatArgs {
     fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'tcx>) {
-        let Some(macro_call) = root_macro_call_first_node(cx, expr) else {
-            return;
-        };
-        if !is_format_macro(cx, macro_call.def_id) {
-            return;
-        }
-        let name = cx.tcx.item_name(macro_call.def_id);
-
-        find_format_args(cx, expr, macro_call.expn, |format_args| {
+        if let Some(macro_call) = root_macro_call_first_node(cx, expr)
+            && is_format_macro(cx, macro_call.def_id)
+            && let Some(format_args) = find_format_args(cx, expr, macro_call.expn)
+        {
             for piece in &format_args.template {
                 if let FormatArgsPiece::Placeholder(placeholder) = piece
                     && let Ok(index) = placeholder.argument.index
@@ -212,6 +207,7 @@ impl<'tcx> LateLintPass<'tcx> for FormatArgs {
                     }
 
                     if let Ok(arg_hir_expr) = arg_expr {
+                        let name = cx.tcx.item_name(macro_call.def_id);
                         check_format_in_format_args(cx, macro_call.span, name, arg_hir_expr);
                         check_to_string_in_format_args(cx, name, arg_hir_expr);
                     }
@@ -221,7 +217,7 @@ impl<'tcx> LateLintPass<'tcx> for FormatArgs {
             if self.msrv.meets(msrvs::FORMAT_ARGS_CAPTURE) {
                 check_uninlined_args(cx, format_args, macro_call.span, macro_call.def_id, self.ignore_mixed);
             }
-        });
+        }
     }
 
     extract_msrv_attr!(LateContext);
diff --git a/clippy_lints/src/format_impl.rs b/clippy_lints/src/format_impl.rs
@@ -170,30 +170,29 @@ fn check_self_in_format_args<'tcx>(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>,
     if let Some(outer_macro) = root_macro_call_first_node(cx, expr)
         && let macro_def_id = outer_macro.def_id
         && is_format_macro(cx, macro_def_id)
+        && let Some(format_args) = find_format_args(cx, expr, outer_macro.expn)
     {
-        find_format_args(cx, expr, outer_macro.expn, |format_args| {
-            for piece in &format_args.template {
-                if let FormatArgsPiece::Placeholder(placeholder) = piece
-                    && let trait_name = match placeholder.format_trait {
-                        FormatTrait::Display => sym::Display,
-                        FormatTrait::Debug => sym::Debug,
-                        FormatTrait::LowerExp => sym!(LowerExp),
-                        FormatTrait::UpperExp => sym!(UpperExp),
-                        FormatTrait::Octal => sym!(Octal),
-                        FormatTrait::Pointer => sym::Pointer,
-                        FormatTrait::Binary => sym!(Binary),
-                        FormatTrait::LowerHex => sym!(LowerHex),
-                        FormatTrait::UpperHex => sym!(UpperHex),
-                    }
-                    && trait_name == impl_trait.name
-                    && let Ok(index) = placeholder.argument.index
-                    && let Some(arg) = format_args.arguments.all_args().get(index)
-                    && let Ok(arg_expr) = find_format_arg_expr(expr, arg)
-                {
-                    check_format_arg_self(cx, expr.span, arg_expr, impl_trait);
+        for piece in &format_args.template {
+            if let FormatArgsPiece::Placeholder(placeholder) = piece
+                && let trait_name = match placeholder.format_trait {
+                    FormatTrait::Display => sym::Display,
+                    FormatTrait::Debug => sym::Debug,
+                    FormatTrait::LowerExp => sym!(LowerExp),
+                    FormatTrait::UpperExp => sym!(UpperExp),
+                    FormatTrait::Octal => sym!(Octal),
+                    FormatTrait::Pointer => sym::Pointer,
+                    FormatTrait::Binary => sym!(Binary),
+                    FormatTrait::LowerHex => sym!(LowerHex),
+                    FormatTrait::UpperHex => sym!(UpperHex),
                 }
+                && trait_name == impl_trait.name
+                && let Ok(index) = placeholder.argument.index
+                && let Some(arg) = format_args.arguments.all_args().get(index)
+                && let Ok(arg_expr) = find_format_arg_expr(expr, arg)
+            {
+                check_format_arg_self(cx, expr.span, arg_expr, impl_trait);
             }
-        });
+        }
     }
 }
 
diff --git a/clippy_lints/src/lib.rs b/clippy_lints/src/lib.rs
@@ -607,7 +607,7 @@ pub fn register_plugins(store: &mut rustc_lint::LintStore, sess: &Session, conf:
                 .collect(),
         ))
     });
-    store.register_early_pass(|| Box::new(utils::format_args_collector::FormatArgsCollector));
+    store.register_early_pass(|| Box::<utils::format_args_collector::FormatArgsCollector>::default());
     store.register_late_pass(|_| Box::new(utils::dump_hir::DumpHir));
     store.register_late_pass(|_| Box::new(utils::author::Author));
     let await_holding_invalid_types = conf.await_holding_invalid_types.clone();
diff --git a/clippy_lints/src/methods/expect_fun_call.rs b/clippy_lints/src/methods/expect_fun_call.rs
@@ -131,12 +131,11 @@ pub(super) fn check<'tcx>(
 
     let mut applicability = Applicability::MachineApplicable;
 
-    //Special handling for `format!` as arg_root
+    // Special handling for `format!` as arg_root
     if let Some(macro_call) = root_macro_call_first_node(cx, arg_root) {
-        if !cx.tcx.is_diagnostic_item(sym::format_macro, macro_call.def_id) {
-            return;
-        }
-        find_format_args(cx, arg_root, macro_call.expn, |format_args| {
+        if cx.tcx.is_diagnostic_item(sym::format_macro, macro_call.def_id)
+            && let Some(format_args) = find_format_args(cx, arg_root, macro_call.expn)
+        {
             let span = format_args_inputs_span(format_args);
             let sugg = snippet_with_applicability(cx, span, "..", &mut applicability);
             span_lint_and_sugg(
@@ -148,7 +147,7 @@ pub(super) fn check<'tcx>(
                 format!("unwrap_or_else({closure_args} panic!({sugg}))"),
                 applicability,
             );
-        });
+        }
         return;
     }
 
diff --git a/clippy_lints/src/utils/format_args_collector.rs b/clippy_lints/src/utils/format_args_collector.rs
@@ -1,12 +1,14 @@
 use clippy_utils::macros::collect_ast_format_args;
 use clippy_utils::source::snippet_opt;
 use itertools::Itertools;
-use rustc_ast::{Expr, ExprKind, FormatArgs};
+use rustc_ast::{Crate, Expr, ExprKind, FormatArgs};
+use rustc_data_structures::fx::FxHashMap;
 use rustc_lexer::{tokenize, TokenKind};
 use rustc_lint::{EarlyContext, EarlyLintPass};
-use rustc_session::{declare_lint_pass, declare_tool_lint};
-use rustc_span::hygiene;
+use rustc_session::{declare_tool_lint, impl_lint_pass};
+use rustc_span::{hygiene, Span};
 use std::iter::once;
+use std::mem;
 
 declare_clippy_lint! {
     /// ### What it does
@@ -17,7 +19,12 @@ declare_clippy_lint! {
     "collects `format_args` AST nodes for use in later lints"
 }
 
-declare_lint_pass!(FormatArgsCollector => [FORMAT_ARGS_COLLECTOR]);
+#[derive(Default)]
+pub struct FormatArgsCollector {
+    format_args: FxHashMap<Span, FormatArgs>,
+}
+
+impl_lint_pass!(FormatArgsCollector => [FORMAT_ARGS_COLLECTOR]);
 
 impl EarlyLintPass for FormatArgsCollector {
     fn check_expr(&mut self, cx: &EarlyContext<'_>, expr: &Expr) {
@@ -26,9 +33,13 @@ impl EarlyLintPass for FormatArgsCollector {
                 return;
             }
 
-            collect_ast_format_args(expr.span, args);
+            self.format_args.insert(expr.span, (**args).clone());
         }
     }
+
+    fn check_crate_post(&mut self, _: &EarlyContext<'_>, _: &Crate) {
+        collect_ast_format_args(mem::take(&mut self.format_args));
+    }
 }
 
 /// Detects if the format string or an argument has its span set by a proc macro to something inside
diff --git a/clippy_lints/src/write.rs b/clippy_lints/src/write.rs
@@ -304,7 +304,7 @@ impl<'tcx> LateLintPass<'tcx> for Write {
             _ => return,
         }
 
-        find_format_args(cx, expr, macro_call.expn, |format_args| {
+        if let Some(format_args) = find_format_args(cx, expr, macro_call.expn) {
             // ignore `writeln!(w)` and `write!(v, some_macro!())`
             if format_args.span.from_expansion() {
                 return;
@@ -334,7 +334,7 @@ impl<'tcx> LateLintPass<'tcx> for Write {
                     }
                 }
             }
-        });
+        }
     }
 }
 
diff --git a/clippy_utils/src/macros.rs b/clippy_utils/src/macros.rs
@@ -10,7 +10,7 @@ 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::RefCell;
+use std::cell::OnceCell;
 use std::ops::ControlFlow;
 use std::sync::atomic::{AtomicBool, Ordering};
 
@@ -374,28 +374,31 @@ thread_local! {
     /// 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.
-    static AST_FORMAT_ARGS: RefCell<FxHashMap<Span, FormatArgs>> = {
+    static AST_FORMAT_ARGS: OnceCell<FxHashMap<Span, 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",
         );
 
-        RefCell::default()
+        OnceCell::new()
     };
 }
 
 /// Record [`rustc_ast::FormatArgs`] for use in late lint passes, this should only be called by
 /// `FormatArgsCollector`
-pub fn collect_ast_format_args(span: Span, format_args: &FormatArgs) {
+#[allow(clippy::implicit_hasher)]
+#[doc(hidden)]
+pub fn collect_ast_format_args(format_args: FxHashMap<Span, FormatArgs>) {
     AST_FORMAT_ARGS.with(|ast_format_args| {
-        ast_format_args.borrow_mut().insert(span, format_args.clone());
+        let result = ast_format_args.set(format_args);
+        debug_assert!(result.is_ok(), "`collect_ast_format_args` should only be called once");
     });
 }
 
-/// Calls `callback` with 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, callback: impl FnOnce(&FormatArgs)) {
+/// Returns an AST [`FormatArgs`] node if a `format_args` expansion is found as a descendant of
+/// `expn_id`
+pub fn find_format_args<'a>(cx: &'a LateContext<'_>, start: &Expr<'_>, expn_id: ExpnId) -> Option<&'a FormatArgs> {
     let format_args_expr = for_each_expr(start, |expr| {
         let ctxt = expr.span.ctxt();
         if ctxt.outer_expn().is_descendant_of(expn_id) {
@@ -410,13 +413,18 @@ pub fn find_format_args(cx: &LateContext<'_>, start: &Expr<'_>, expn_id: ExpnId,
         } else {
             ControlFlow::Continue(Descend::No)
         }
-    });
+    })?;
 
-    if let Some(expr) = format_args_expr {
-        AST_FORMAT_ARGS.with(|ast_format_args| {
-            ast_format_args.borrow().get(&expr.span).map(callback);
-        });
-    }
+    AST_FORMAT_ARGS.with(|ast_format_args| {
+        let found = ast_format_args.get()?.get(&format_args_expr.span)?;
+
+        // SAFETY: the lifetime is 'thread which isn't nameable, so we instead use the lifetime of the
+        // reference to the `LateContext`
+        //
+        // It's possible for a `LateContext` that outlives the thread to exist e.g. with `Box::leak`, but
+        // only `rustc_lint` can construct a `LateContext` and it does not do that
+        unsafe { Some(std::mem::transmute::<&FormatArgs, &'a FormatArgs>(found)) }
+    })
 }
 
 /// Attempt to find the [`rustc_hir::Expr`] that corresponds to the [`FormatArgument`]'s value, if

Original file line number	Diff line number	Diff line change
`@@ -304,7 +304,7 @@ impl<'tcx> LateLintPass<'tcx> for Write {`
`304`	`304`	`_ => return,`
`305`	`305`	`}`
`306`	`306`
`307`		`- find_format_args(cx, expr, macro_call.expn, \|format_args\| {`
	`307`	`+ if let Some(format_args) = find_format_args(cx, expr, macro_call.expn) {`
`308`	`308`	// ignore `writeln!(w)` and `write!(v, some_macro!())`
`309`	`309`	`if format_args.span.from_expansion() {`
`310`	`310`	`return;`
`@@ -334,7 +334,7 @@ impl<'tcx> LateLintPass<'tcx> for Write {`
`334`	`334`	`}`
`335`	`335`	`}`
`336`	`336`	`}`
`337`		`- });`
	`337`	`+ }`
`338`	`338`	`}`
`339`	`339`	`}`
`340`	`340`