Rollup merge of #70367 - nikomatsakis:issue-69307, r=Aaron1011

save/restore `pessimistic_yield` when entering bodies

This flag is used to make the execution order around `+=` operators
pessimistic. Failure to save/restore the flag was causing independent
async blocks to effect one another, leading to strange ICEs and failed
assumptions.

Fixes #69307

r? @Zoxc

3 files changed, 62 insertions, 0 deletions

diff --git a/src/librustc_passes/region.rs b/src/librustc_passes/region.rs
index 485480d020f..4c02127c106 100644
--- a/src/librustc_passes/region.rs
+++ b/src/librustc_passes/region.rs
@@ -717,9 +717,17 @@ impl<'tcx> Visitor<'tcx> for RegionResolutionVisitor<'tcx> {
             self.cx.parent
         );
 
+        // Save all state that is specific to the outer function
+        // body. These will be restored once down below, once we've
+        // visited the body.
         let outer_ec = mem::replace(&mut self.expr_and_pat_count, 0);
         let outer_cx = self.cx;
         let outer_ts = mem::take(&mut self.terminating_scopes);
+        // The 'pessimistic yield' flag is set to true when we are
+        // processing a `+=` statement and have to make pessimistic
+        // control flow assumptions. This doesn't apply to nested
+        // bodies within the `+=` statements. See #69307.
+        let outer_pessimistic_yield = mem::replace(&mut self.pessimistic_yield, false);
         self.terminating_scopes.insert(body.value.hir_id.local_id);
 
         if let Some(root_id) = self.cx.root_id {
@@ -771,6 +779,7 @@ impl<'tcx> Visitor<'tcx> for RegionResolutionVisitor<'tcx> {
         self.expr_and_pat_count = outer_ec;
         self.cx = outer_cx;
         self.terminating_scopes = outer_ts;
+        self.pessimistic_yield = outer_pessimistic_yield;
     }
 
     fn visit_arm(&mut self, a: &'tcx Arm<'tcx>) {
diff --git a/src/test/ui/async-await/issues/issue-69307-nested.rs b/src/test/ui/async-await/issues/issue-69307-nested.rs
new file mode 100644
index 00000000000..b7cdf3987f1
--- /dev/null
+++ b/src/test/ui/async-await/issues/issue-69307-nested.rs
@@ -0,0 +1,30 @@
+// Regression test for #69307
+//
+// Having a `async { .. foo.await .. }` block appear inside of a `+=`
+// expression was causing an ICE due to a failure to save/restore
+// state in the AST numbering pass when entering a nested body.
+//
+// check-pass
+// edition:2018
+
+fn block_on<F>(_: F) -> usize {
+    0
+}
+
+fn main() {}
+
+async fn bar() {
+    let mut sum = 0;
+    sum += {
+        block_on(async {
+            baz().await;
+            let mut inner = 1;
+            inner += block_on(async {
+                baz().await;
+                0
+            })
+        })
+    };
+}
+
+async fn baz() {}
diff --git a/src/test/ui/async-await/issues/issue-69307.rs b/src/test/ui/async-await/issues/issue-69307.rs
new file mode 100644
index 00000000000..4dae96ec8a6
--- /dev/null
+++ b/src/test/ui/async-await/issues/issue-69307.rs
@@ -0,0 +1,23 @@
+// Regression test for #69307
+//
+// Having a `async { .. foo.await .. }` block appear inside of a `+=`
+// expression was causing an ICE due to a failure to save/restore
+// state in the AST numbering pass when entering a nested body.
+//
+// check-pass
+// edition:2018
+
+fn block_on<F>(_: F) -> usize {
+    0
+}
+
+fn main() {}
+
+async fn bar() {
+    let mut sum = 0;
+    sum += block_on(async {
+        baz().await;
+    });
+}
+
+async fn baz() {}