Fix wrong suggestion for async gen block and add regression ui test for #139839

4 files changed, 126 insertions, 4 deletions

diff --git a/compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs b/compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs
index fe6dff7ff1b..959cf9fa513 100644
--- a/compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs
+++ b/compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs
@@ -3376,10 +3376,15 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
 
         let (sugg_span, suggestion) = match tcx.sess.source_map().span_to_snippet(args_span) {
             Ok(string) => {
-                let coro_prefix = if string.starts_with("async") {
-                    // `async` is 5 chars long. Not using `.len()` to avoid the cast from `usize`
-                    // to `u32`.
-                    Some(5)
+                let coro_prefix = if let Some(sub) = string.strip_prefix("async") {
+                    let trimmed_sub = sub.trim_end();
+                    if trimmed_sub.ends_with("gen") {
+                        // `async` is 5 chars long.
+                        Some((trimmed_sub.len() + 5) as _)
+                    } else {
+                        // `async` is 5 chars long.
+                        Some(5)
+                    }
                 } else if string.starts_with("gen") {
                     // `gen` is 3 chars long
                     Some(3)
diff --git a/tests/ui/async-await/async-gen-move-suggestion.fixed b/tests/ui/async-await/async-gen-move-suggestion.fixed
new file mode 100644
index 00000000000..d8020765528
--- /dev/null
+++ b/tests/ui/async-await/async-gen-move-suggestion.fixed
@@ -0,0 +1,35 @@
+// This is a regression test for <https://github.com/rust-lang/rust/issues/139839>.
+// It ensures that the "add `move` keyword" suggestion is valid.
+
+//@ run-rustfix
+//@ edition:2024
+
+#![feature(coroutines)]
+#![feature(gen_blocks)]
+#![feature(async_iterator)]
+
+use std::async_iter::AsyncIterator;
+
+#[allow(dead_code)]
+fn moved() -> impl AsyncIterator<Item = u32> {
+    let mut x = "foo".to_string();
+
+    async gen move { //~ ERROR
+        x.clear();
+        for x in 3..6 { yield x }
+    }
+}
+
+#[allow(dead_code)]
+fn check_with_whitespace_chars() -> impl AsyncIterator<Item = u32> {
+    let mut x = "foo".to_string();
+
+    async // Just to check that whitespace characters are correctly handled
+    gen move { //~^ ERROR
+        x.clear();
+        for x in 3..6 { yield x }
+    }
+}
+
+fn main() {
+}
diff --git a/tests/ui/async-await/async-gen-move-suggestion.rs b/tests/ui/async-await/async-gen-move-suggestion.rs
new file mode 100644
index 00000000000..825fb0fd189
--- /dev/null
+++ b/tests/ui/async-await/async-gen-move-suggestion.rs
@@ -0,0 +1,35 @@
+// This is a regression test for <https://github.com/rust-lang/rust/issues/139839>.
+// It ensures that the "add `move` keyword" suggestion is valid.
+
+//@ run-rustfix
+//@ edition:2024
+
+#![feature(coroutines)]
+#![feature(gen_blocks)]
+#![feature(async_iterator)]
+
+use std::async_iter::AsyncIterator;
+
+#[allow(dead_code)]
+fn moved() -> impl AsyncIterator<Item = u32> {
+    let mut x = "foo".to_string();
+
+    async gen { //~ ERROR
+        x.clear();
+        for x in 3..6 { yield x }
+    }
+}
+
+#[allow(dead_code)]
+fn check_with_whitespace_chars() -> impl AsyncIterator<Item = u32> {
+    let mut x = "foo".to_string();
+
+    async // Just to check that whitespace characters are correctly handled
+    gen { //~^ ERROR
+        x.clear();
+        for x in 3..6 { yield x }
+    }
+}
+
+fn main() {
+}
diff --git a/tests/ui/async-await/async-gen-move-suggestion.stderr b/tests/ui/async-await/async-gen-move-suggestion.stderr
new file mode 100644
index 00000000000..b8cdb8be7a4
--- /dev/null
+++ b/tests/ui/async-await/async-gen-move-suggestion.stderr
@@ -0,0 +1,47 @@
+error[E0373]: async gen block may outlive the current function, but it borrows `x`, which is owned by the current function
+  --> $DIR/async-gen-move-suggestion.rs:17:5
+   |
+LL |     async gen {
+   |     ^^^^^^^^^ may outlive borrowed value `x`
+LL |         x.clear();
+   |         - `x` is borrowed here
+   |
+note: async gen block is returned here
+  --> $DIR/async-gen-move-suggestion.rs:17:5
+   |
+LL | /     async gen {
+LL | |         x.clear();
+LL | |         for x in 3..6 { yield x }
+LL | |     }
+   | |_____^
+help: to force the async gen block to take ownership of `x` (and any other referenced variables), use the `move` keyword
+   |
+LL |     async gen move {
+   |               ++++
+
+error[E0373]: async gen block may outlive the current function, but it borrows `x`, which is owned by the current function
+  --> $DIR/async-gen-move-suggestion.rs:27:5
+   |
+LL | /     async // Just to check that whitespace characters are correctly handled
+LL | |     gen {
+   | |_______^ may outlive borrowed value `x`
+LL |           x.clear();
+   |           - `x` is borrowed here
+   |
+note: async gen block is returned here
+  --> $DIR/async-gen-move-suggestion.rs:27:5
+   |
+LL | /     async // Just to check that whitespace characters are correctly handled
+LL | |     gen {
+LL | |         x.clear();
+LL | |         for x in 3..6 { yield x }
+LL | |     }
+   | |_____^
+help: to force the async gen block to take ownership of `x` (and any other referenced variables), use the `move` keyword
+   |
+LL |     gen move {
+   |         ++++
+
+error: aborting due to 2 previous errors
+
+For more information about this error, try `rustc --explain E0373`.