From 2bc04bdac26fbb49e3b4ef3673f6a5143df307e7 Mon Sep 17 00:00:00 2001
From: koka <koka.code@gmail.com>
Date: Fri, 11 Nov 2022 00:14:18 +0900
Subject: [PATCH] fix: cognitive_complexity for async fn

---
 clippy_lints/src/cognitive_complexity.rs | 21 +++++++++++++++------
 tests/ui/cognitive_complexity.rs         |  8 ++++++++
 tests/ui/cognitive_complexity.stderr     | 10 +++++++++-
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/clippy_lints/src/cognitive_complexity.rs b/clippy_lints/src/cognitive_complexity.rs
index 77af3b53d633..1c3a89a97824 100644
--- a/clippy_lints/src/cognitive_complexity.rs
+++ b/clippy_lints/src/cognitive_complexity.rs
@@ -4,11 +4,11 @@ use clippy_utils::diagnostics::span_lint_and_help;
 use clippy_utils::source::snippet_opt;
 use clippy_utils::ty::is_type_diagnostic_item;
 use clippy_utils::visitors::for_each_expr;
-use clippy_utils::LimitStack;
+use clippy_utils::{get_async_fn_body, is_async_fn, LimitStack};
 use core::ops::ControlFlow;
 use rustc_ast::ast::Attribute;
 use rustc_hir::intravisit::FnKind;
-use rustc_hir::{Body, ExprKind, FnDecl, HirId};
+use rustc_hir::{Body, Expr, ExprKind, FnDecl, HirId};
 use rustc_lint::{LateContext, LateLintPass, LintContext};
 use rustc_session::{declare_tool_lint, impl_lint_pass};
 use rustc_span::source_map::Span;
@@ -56,15 +56,13 @@ impl CognitiveComplexity {
         cx: &LateContext<'tcx>,
         kind: FnKind<'tcx>,
         decl: &'tcx FnDecl<'_>,
-        body: &'tcx Body<'_>,
+        expr: &'tcx Expr<'_>,
         body_span: Span,
     ) {
         if body_span.from_expansion() {
             return;
         }
 
-        let expr = body.value;
-
         let mut cc = 1u64;
         let mut returns = 0u64;
         let _: Option<!> = for_each_expr(expr, |e| {
@@ -146,7 +144,18 @@ impl<'tcx> LateLintPass<'tcx> for CognitiveComplexity {
     ) {
         let def_id = cx.tcx.hir().local_def_id(hir_id);
         if !cx.tcx.has_attr(def_id.to_def_id(), sym::test) {
-            self.check(cx, kind, decl, body, span);
+            let expr = if is_async_fn(kind) {
+                match get_async_fn_body(cx.tcx, body) {
+                    Some(b) => b,
+                    None => {
+                        return;
+                    },
+                }
+            } else {
+                body.value
+            };
+
+            self.check(cx, kind, decl, expr, span);
         }
     }
 
diff --git a/tests/ui/cognitive_complexity.rs b/tests/ui/cognitive_complexity.rs
index 912e6788afdd..498a9f0644f5 100644
--- a/tests/ui/cognitive_complexity.rs
+++ b/tests/ui/cognitive_complexity.rs
@@ -393,3 +393,11 @@ impl Moo {
         }
     }
 }
+
+#[clippy::cognitive_complexity = "1"]
+mod issue9300 {
+    async fn a() {
+        let a = 0;
+        if a == 0 {}
+    }
+}
diff --git a/tests/ui/cognitive_complexity.stderr b/tests/ui/cognitive_complexity.stderr
index d7f2f24e52f2..208f4d7b3474 100644
--- a/tests/ui/cognitive_complexity.stderr
+++ b/tests/ui/cognitive_complexity.stderr
@@ -135,5 +135,13 @@ LL |     fn moo(&self) {
    |
    = help: you could split it up into multiple smaller functions
 
-error: aborting due to 17 previous errors
+error: the function has a cognitive complexity of (2/1)
+  --> $DIR/cognitive_complexity.rs:399:14
+   |
+LL |     async fn a() {
+   |              ^
+   |
+   = help: you could split it up into multiple smaller functions
+
+error: aborting due to 18 previous errors