From d526adad2580d20c6f01c4c285e699492dbd45f7 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Thu, 27 Jun 2024 15:32:29 -0400
Subject: [PATCH] Basic lint detecting closure-returning-async-block

---
 compiler/rustc_lint/messages.ftl          |   3 +
 compiler/rustc_lint/src/async_closures.rs | 110 ++++++++++++++++++++++
 compiler/rustc_lint/src/lib.rs            |   3 +
 3 files changed, 116 insertions(+)
 create mode 100644 compiler/rustc_lint/src/async_closures.rs

diff --git a/compiler/rustc_lint/messages.ftl b/compiler/rustc_lint/messages.ftl
index 46cf87d1e3c..8c4cfe9b87e 100644
--- a/compiler/rustc_lint/messages.ftl
+++ b/compiler/rustc_lint/messages.ftl
@@ -187,6 +187,9 @@ lint_cfg_attr_no_attributes =
 
 lint_check_name_unknown_tool = unknown lint tool: `{$tool_name}`
 
+lint_closure_returning_async_block = closure returning async block can be made into an async closure
+    .label = this async block can be removed, and the closure can be turned into an async closure
+
 lint_command_line_source = `forbid` lint level was set on command line
 
 lint_confusable_identifier_pair = found both `{$existing_sym}` and `{$sym}` as identifiers, which look alike
diff --git a/compiler/rustc_lint/src/async_closures.rs b/compiler/rustc_lint/src/async_closures.rs
new file mode 100644
index 00000000000..8a72b1d153b
--- /dev/null
+++ b/compiler/rustc_lint/src/async_closures.rs
@@ -0,0 +1,110 @@
+use rustc_hir as hir;
+use rustc_macros::LintDiagnostic;
+use rustc_session::{declare_lint, declare_lint_pass};
+use rustc_span::Span;
+
+use crate::{LateContext, LateLintPass};
+
+declare_lint! {
+    /// The `closure_returning_async_block` lint detects cases where users
+    /// write a closure that returns an async block.
+    ///
+    /// ### Example
+    ///
+    /// ```rust
+    /// #![warn(closure_returning_async_block)]
+    /// let c = |x: &str| async {};
+    /// ```
+    ///
+    /// {{produces}}
+    ///
+    /// ### Explanation
+    ///
+    /// Using an async closure is preferable over a closure that returns an
+    /// async block, since async closures are less restrictive in how its
+    /// captures are allowed to be used.
+    ///
+    /// For example, this code does not work with a closure returning an async
+    /// block:
+    ///
+    /// ```rust,compile_fail
+    /// async fn callback(x: &str) {}
+    ///
+    /// let captured_str = String::new();
+    /// let c = move || async {
+    ///     callback(&captured_str).await;
+    /// };
+    /// ```
+    ///
+    /// But it does work with async closures:
+    ///
+    /// ```rust
+    /// #![feature(async_closure)]
+    ///
+    /// async fn callback(x: &str) {}
+    ///
+    /// let captured_str = String::new();
+    /// let c = async move || {
+    ///     callback(&captured_str).await;
+    /// };
+    /// ```
+    pub CLOSURE_RETURNING_ASYNC_BLOCK,
+    Allow,
+    "closure that returns `async {}` could be rewritten as an async closure",
+    @feature_gate = async_closure;
+}
+
+declare_lint_pass!(
+    /// Lint for potential usages of async closures and async fn trait bounds.
+    AsyncClosureUsage => [CLOSURE_RETURNING_ASYNC_BLOCK]
+);
+
+impl<'tcx> LateLintPass<'tcx> for AsyncClosureUsage {
+    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx hir::Expr<'tcx>) {
+        let hir::ExprKind::Closure(&hir::Closure {
+            body,
+            kind: hir::ClosureKind::Closure,
+            fn_decl_span,
+            ..
+        }) = expr.kind
+        else {
+            return;
+        };
+
+        let mut body = cx.tcx.hir().body(body).value;
+
+        // Only peel blocks that have no expressions.
+        while let hir::ExprKind::Block(&hir::Block { stmts: [], expr: Some(tail), .. }, None) =
+            body.kind
+        {
+            body = tail;
+        }
+
+        let hir::ExprKind::Closure(&hir::Closure {
+            kind:
+                hir::ClosureKind::Coroutine(hir::CoroutineKind::Desugared(
+                    hir::CoroutineDesugaring::Async,
+                    hir::CoroutineSource::Block,
+                )),
+            fn_decl_span: async_decl_span,
+            ..
+        }) = body.kind
+        else {
+            return;
+        };
+
+        cx.tcx.emit_node_span_lint(
+            CLOSURE_RETURNING_ASYNC_BLOCK,
+            expr.hir_id,
+            fn_decl_span,
+            ClosureReturningAsyncBlock { async_decl_span },
+        );
+    }
+}
+
+#[derive(LintDiagnostic)]
+#[diag(lint_closure_returning_async_block)]
+struct ClosureReturningAsyncBlock {
+    #[label]
+    async_decl_span: Span,
+}
diff --git a/compiler/rustc_lint/src/lib.rs b/compiler/rustc_lint/src/lib.rs
index 7dae2de7bfb..43e6e6498b3 100644
--- a/compiler/rustc_lint/src/lib.rs
+++ b/compiler/rustc_lint/src/lib.rs
@@ -41,6 +41,7 @@
 #![feature(trait_upcasting)]
 // tidy-alphabetical-end
 
+mod async_closures;
 mod async_fn_in_trait;
 pub mod builtin;
 mod context;
@@ -86,6 +87,7 @@ use rustc_hir::def_id::LocalModDefId;
 use rustc_middle::query::Providers;
 use rustc_middle::ty::TyCtxt;
 
+use async_closures::AsyncClosureUsage;
 use async_fn_in_trait::AsyncFnInTrait;
 use builtin::*;
 use deref_into_dyn_supertrait::*;
@@ -227,6 +229,7 @@ late_lint_methods!(
             MapUnitFn: MapUnitFn,
             MissingDebugImplementations: MissingDebugImplementations,
             MissingDoc: MissingDoc,
+            AsyncClosureUsage: AsyncClosureUsage,
             AsyncFnInTrait: AsyncFnInTrait,
             NonLocalDefinitions: NonLocalDefinitions::default(),
             ImplTraitOvercaptures: ImplTraitOvercaptures,