Use late lint pass for `missing_assert_message` lint

Co-authored-by: Weihang Lo <me@weihanglo.tw>

2 files changed, 27 insertions, 81 deletions

diff --git a/clippy_lints/src/lib.rs b/clippy_lints/src/lib.rs
index e3be798f30b..6d4ec527756 100644
--- a/clippy_lints/src/lib.rs
+++ b/clippy_lints/src/lib.rs
@@ -927,7 +927,7 @@ pub fn register_plugins(store: &mut rustc_lint::LintStore, sess: &Session, conf:
     });
     store.register_late_pass(|_| Box::new(no_mangle_with_rust_abi::NoMangleWithRustAbi));
     store.register_late_pass(|_| Box::new(collection_is_never_read::CollectionIsNeverRead));
-    store.register_pre_expansion_pass(|| Box::<missing_assert_message::MissingAssertMessage>::default());
+    store.register_late_pass(|_| Box::new(missing_assert_message::MissingAssertMessage));
     // add lints here, do not remove this comment, it's used in `new_lint`
 }
 
diff --git a/clippy_lints/src/missing_assert_message.rs b/clippy_lints/src/missing_assert_message.rs
index f499b3e50ca..2ff0cee2925 100644
--- a/clippy_lints/src/missing_assert_message.rs
+++ b/clippy_lints/src/missing_assert_message.rs
@@ -1,11 +1,9 @@
 use clippy_utils::diagnostics::span_lint_and_help;
-use rustc_ast::ast;
-use rustc_ast::{
-    token::{Token, TokenKind},
-    tokenstream::TokenTree,
-};
-use rustc_lint::{EarlyContext, EarlyLintPass};
-use rustc_session::{declare_tool_lint, impl_lint_pass};
+use clippy_utils::macros::{find_assert_args, find_assert_eq_args, root_macro_call_first_node, PanicExpn};
+use clippy_utils::{is_in_cfg_test, is_in_test_function};
+use rustc_hir::Expr;
+use rustc_lint::{LateContext, LateLintPass};
+use rustc_session::{declare_lint_pass, declare_tool_lint};
 use rustc_span::sym;
 
 declare_clippy_lint! {
@@ -37,93 +35,41 @@ declare_clippy_lint! {
     "checks assertions without a custom panic message"
 }
 
-#[derive(Default, Clone, Debug)]
-pub struct MissingAssertMessage {
-    // This field will be greater than zero if we are inside a `#[test]` or `#[cfg(test)]`
-    test_deepnes: usize,
-}
+declare_lint_pass!(MissingAssertMessage => [MISSING_ASSERT_MESSAGE]);
 
-impl_lint_pass!(MissingAssertMessage => [MISSING_ASSERT_MESSAGE]);
+impl<'tcx> LateLintPass<'tcx> for MissingAssertMessage {
+    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) {
+        let Some(macro_call) = root_macro_call_first_node(cx, expr) else { return };
+        let single_argument = match cx.tcx.get_diagnostic_name(macro_call.def_id) {
+            Some(sym::assert_macro | sym::debug_assert_macro) => true,
+            Some(
+                sym::assert_eq_macro | sym::assert_ne_macro | sym::debug_assert_eq_macro | sym::debug_assert_ne_macro,
+            ) => false,
+            _ => return,
+        };
 
-impl EarlyLintPass for MissingAssertMessage {
-    fn check_mac(&mut self, cx: &EarlyContext<'_>, mac_call: &ast::MacCall) {
-        if self.test_deepnes != 0 {
+        // This lint would be very noisy in tests, so just ignore if we're in test context
+        if is_in_test_function(cx.tcx, expr.hir_id) || is_in_cfg_test(cx.tcx, expr.hir_id) {
             return;
         }
 
-        let Some(last_segment) = mac_call.path.segments.last() else { return; };
-        let num_separators_needed = match last_segment.ident.as_str() {
-            "assert" | "debug_assert" => 1,
-            "assert_eq" | "assert_ne" | "debug_assert_eq" | "debug_assert_ne" => 2,
-            _ => return,
+        let panic_expn = if single_argument {
+            let Some((_, panic_expn)) = find_assert_args(cx, expr, macro_call.expn) else { return };
+            panic_expn
+        } else {
+            let Some((_, _, panic_expn)) = find_assert_eq_args(cx, expr, macro_call.expn) else { return };
+            panic_expn
         };
-        let num_separators = num_commas_on_arguments(mac_call);
 
-        if num_separators < num_separators_needed {
+        if let PanicExpn::Empty = panic_expn {
             span_lint_and_help(
                 cx,
                 MISSING_ASSERT_MESSAGE,
-                mac_call.span(),
+                macro_call.span,
                 "assert without any message",
                 None,
                 "consider describing why the failing assert is problematic",
             );
         }
     }
-
-    fn check_item(&mut self, _: &EarlyContext<'_>, item: &ast::Item) {
-        if item.attrs.iter().any(is_a_test_attribute) {
-            self.test_deepnes += 1;
-        }
-    }
-
-    fn check_item_post(&mut self, _: &EarlyContext<'_>, item: &ast::Item) {
-        if item.attrs.iter().any(is_a_test_attribute) {
-            self.test_deepnes -= 1;
-        }
-    }
-}
-
-// Returns number of commas (excluding trailing comma) from `MacCall`'s arguments.
-fn num_commas_on_arguments(mac_call: &ast::MacCall) -> usize {
-    let mut num_separators = 0;
-    let mut is_trailing = false;
-    for tt in mac_call.args.tokens.trees() {
-        match tt {
-            TokenTree::Token(
-                Token {
-                    kind: TokenKind::Comma,
-                    span: _,
-                },
-                _,
-            ) => {
-                num_separators += 1;
-                is_trailing = true;
-            },
-            _ => {
-                is_trailing = false;
-            },
-        }
-    }
-    if is_trailing {
-        num_separators -= 1;
-    }
-    num_separators
-}
-
-// Returns true if the attribute is either a `#[test]` or a `#[cfg(test)]`.
-fn is_a_test_attribute(attr: &ast::Attribute) -> bool {
-    if attr.has_name(sym::test) {
-        return true;
-    }
-
-    if attr.has_name(sym::cfg)
-        && let Some(items) = attr.meta_item_list()
-        && let [item] = &*items
-        && item.has_name(sym::test)
-    {
-        true
-    } else {
-        false
-    }
 }