From e0a726ca4bc08d1009803a6c819f701173810761 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Tue, 27 Feb 2024 03:14:55 +0000
Subject: [PATCH] Adjust error yield/await lowering

---
 compiler/rustc_ast_lowering/src/expr.rs       | 52 +++++++++++++++----
 .../async-outside-of-await-issue-121096.rs    |  9 ++++
 ...async-outside-of-await-issue-121096.stderr | 12 +++++
 3 files changed, 64 insertions(+), 9 deletions(-)
 create mode 100644 tests/ui/async-await/async-outside-of-await-issue-121096.rs
 create mode 100644 tests/ui/async-await/async-outside-of-await-issue-121096.stderr

diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index 9950db4784b9c..599b7ca0af671 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -764,10 +764,28 @@ impl<'hir> LoweringContext<'_, 'hir> {
             Some(hir::CoroutineKind::Coroutine(_))
             | Some(hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::Gen, _))
             | None => {
-                return hir::ExprKind::Err(self.dcx().emit_err(AwaitOnlyInAsyncFnAndBlocks {
-                    await_kw_span,
-                    item_span: self.current_item,
-                }));
+                // Lower to a block `{ EXPR; <error> }` so that the awaited expr
+                // is not accidentally orphaned.
+                let stmt_id = self.next_id();
+                let expr_err = self.expr(
+                    expr.span,
+                    hir::ExprKind::Err(self.dcx().emit_err(AwaitOnlyInAsyncFnAndBlocks {
+                        await_kw_span,
+                        item_span: self.current_item,
+                    })),
+                );
+                return hir::ExprKind::Block(
+                    self.block_all(
+                        expr.span,
+                        arena_vec![self; hir::Stmt {
+                            hir_id: stmt_id,
+                            kind: hir::StmtKind::Semi(expr),
+                            span: expr.span,
+                        }],
+                        Some(self.arena.alloc(expr_err)),
+                    ),
+                    None,
+                );
             }
         };
 
@@ -1500,12 +1518,31 @@ impl<'hir> LoweringContext<'_, 'hir> {
     }
 
     fn lower_expr_yield(&mut self, span: Span, opt_expr: Option<&Expr>) -> hir::ExprKind<'hir> {
+        let yielded =
+            opt_expr.as_ref().map(|x| self.lower_expr(x)).unwrap_or_else(|| self.expr_unit(span));
+
         let is_async_gen = match self.coroutine_kind {
             Some(hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::Gen, _)) => false,
             Some(hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::AsyncGen, _)) => true,
             Some(hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::Async, _)) => {
-                return hir::ExprKind::Err(
-                    self.dcx().emit_err(AsyncCoroutinesNotSupported { span }),
+                // Lower to a block `{ EXPR; <error> }` so that the awaited expr
+                // is not accidentally orphaned.
+                let stmt_id = self.next_id();
+                let expr_err = self.expr(
+                    yielded.span,
+                    hir::ExprKind::Err(self.dcx().emit_err(AsyncCoroutinesNotSupported { span })),
+                );
+                return hir::ExprKind::Block(
+                    self.block_all(
+                        yielded.span,
+                        arena_vec![self; hir::Stmt {
+                            hir_id: stmt_id,
+                            kind: hir::StmtKind::Semi(yielded),
+                            span: yielded.span,
+                        }],
+                        Some(self.arena.alloc(expr_err)),
+                    ),
+                    None,
                 );
             }
             Some(hir::CoroutineKind::Coroutine(_)) => {
@@ -1535,9 +1572,6 @@ impl<'hir> LoweringContext<'_, 'hir> {
             }
         };
 
-        let yielded =
-            opt_expr.as_ref().map(|x| self.lower_expr(x)).unwrap_or_else(|| self.expr_unit(span));
-
         if is_async_gen {
             // `yield $expr` is transformed into `task_context = yield async_gen_ready($expr)`.
             // This ensures that we store our resumed `ResumeContext` correctly, and also that
diff --git a/tests/ui/async-await/async-outside-of-await-issue-121096.rs b/tests/ui/async-await/async-outside-of-await-issue-121096.rs
new file mode 100644
index 0000000000000..e3999035ef91f
--- /dev/null
+++ b/tests/ui/async-await/async-outside-of-await-issue-121096.rs
@@ -0,0 +1,9 @@
+//@ edition:2021
+
+fn main() {
+    async {
+        use std::ops::Add;
+        let _ = 1.add(3);
+    }.await
+    //~^ ERROR `await` is only allowed inside `async` functions and blocks
+}
diff --git a/tests/ui/async-await/async-outside-of-await-issue-121096.stderr b/tests/ui/async-await/async-outside-of-await-issue-121096.stderr
new file mode 100644
index 0000000000000..b0677a83864e0
--- /dev/null
+++ b/tests/ui/async-await/async-outside-of-await-issue-121096.stderr
@@ -0,0 +1,12 @@
+error[E0728]: `await` is only allowed inside `async` functions and blocks
+  --> $DIR/async-outside-of-await-issue-121096.rs:7:7
+   |
+LL | fn main() {
+   |    ---- this is not `async`
+...
+LL |     }.await
+   |       ^^^^^ only allowed inside `async` functions and blocks
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0728`.