fix: break inside async closure has incorrect span for enclosing closure

linyihai · linyihai · commit 1ecf1a2823b6 · 2024-05-27T14:19:41.000+08:00
diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
@@ -203,17 +203,19 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     fn_decl_span,
                     fn_arg_span,
                 }) => match coroutine_kind {
-                    Some(coroutine_kind) => self.lower_expr_coroutine_closure(
-                        binder,
-                        *capture_clause,
-                        e.id,
-                        hir_id,
-                        *coroutine_kind,
-                        fn_decl,
-                        body,
-                        *fn_decl_span,
-                        *fn_arg_span,
-                    ),
+                    Some(coroutine_kind) => self.with_new_scopes(*fn_decl_span, |this| {
+                        this.lower_expr_coroutine_closure(
+                            binder,
+                            *capture_clause,
+                            e.id,
+                            hir_id,
+                            *coroutine_kind,
+                            fn_decl,
+                            body,
+                            *fn_decl_span,
+                            *fn_arg_span,
+                        )
+                    }),
                     None => self.lower_expr_closure(
                         binder,
                         *capture_clause,
@@ -684,6 +686,12 @@ impl<'hir> LoweringContext<'_, 'hir> {
             (params, res)
         });
 
+        let fn_decl_span = if matches!(coroutine_source, hir::CoroutineSource::Closure) {
+            self.current_item.map_or(span, |fn_decl_span| span.with_lo(fn_decl_span.lo()))
+        } else {
+            span
+        };
+
         // `static |<_task_context?>| -> <return_ty> { <body> }`:
         hir::ExprKind::Closure(self.arena.alloc(hir::Closure {
             def_id: self.local_def_id(closure_node_id),
@@ -692,7 +700,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             bound_generic_params: &[],
             fn_decl,
             body,
-            fn_decl_span: self.lower_span(span),
+            fn_decl_span: self.lower_span(fn_decl_span),
             fn_arg_span: None,
             kind: hir::ClosureKind::Coroutine(coroutine_kind),
             constness: hir::Constness::NotConst,
diff --git a/tests/ui/async-await/async-closures/wrong-fn-kind.stderr b/tests/ui/async-await/async-closures/wrong-fn-kind.stderr
@@ -20,15 +20,16 @@ LL | fn needs_async_fn(_: impl async Fn()) {}
    |                           ^^^^^^^^^^ required by this bound in `needs_async_fn`
 
 error[E0596]: cannot borrow `x` as mutable, as it is a captured variable in a `Fn` closure
-  --> $DIR/wrong-fn-kind.rs:9:29
+  --> $DIR/wrong-fn-kind.rs:9:20
    |
 LL |   fn needs_async_fn(_: impl async Fn()) {}
    |                        --------------- change this to accept `FnMut` instead of `Fn`
 ...
 LL |       needs_async_fn(async || {
-   |  _____--------------_--------_^
-   | |     |              |
-   | |     |              in this closure
+   |       -------------- ^-------
+   |       |              |
+   |  _____|______________in this closure
+   | |     |
    | |     expects `Fn` instead of `FnMut`
 LL | |
 LL | |         x += 1;
diff --git a/tests/ui/coroutine/break-inside-coroutine-issue-124495.rs b/tests/ui/coroutine/break-inside-coroutine-issue-124495.rs
@@ -18,6 +18,7 @@ async gen fn async_gen_fn() {
 
 fn main() {
     let _ = async { break; }; //~ ERROR `break` inside `async` block
+
     let _ = async || { break; }; //~ ERROR `break` inside `async` closure
 
     let _ = gen { break; }; //~ ERROR `break` inside `gen` block
diff --git a/tests/ui/coroutine/break-inside-coroutine-issue-124495.stderr b/tests/ui/coroutine/break-inside-coroutine-issue-124495.stderr
@@ -38,16 +38,16 @@ LL |     let _ = async { break; };
    |             enclosing `async` block
 
 error[E0267]: `break` inside `async` closure
-  --> $DIR/break-inside-coroutine-issue-124495.rs:21:24
+  --> $DIR/break-inside-coroutine-issue-124495.rs:22:24
    |
 LL |     let _ = async || { break; };
-   |                      --^^^^^---
-   |                      | |
-   |                      | cannot `break` inside `async` closure
-   |                      enclosing `async` closure
+   |             -----------^^^^^---
+   |             |          |
+   |             |          cannot `break` inside `async` closure
+   |             enclosing `async` closure
 
 error[E0267]: `break` inside `gen` block
-  --> $DIR/break-inside-coroutine-issue-124495.rs:23:19
+  --> $DIR/break-inside-coroutine-issue-124495.rs:24:19
    |
 LL |     let _ = gen { break; };
    |             ------^^^^^---
@@ -56,7 +56,7 @@ LL |     let _ = gen { break; };
    |             enclosing `gen` block
 
 error[E0267]: `break` inside `async gen` block
-  --> $DIR/break-inside-coroutine-issue-124495.rs:25:25
+  --> $DIR/break-inside-coroutine-issue-124495.rs:26:25
    |
 LL |     let _ = async gen { break; };
    |             ------------^^^^^---
