From cce706ec3c4ed38ecaffc61dac3b835038cdc50b Mon Sep 17 00:00:00 2001
From: Andrew Zhogin <andrew.zhogin@gmail.com>
Date: Tue, 29 Apr 2025 20:41:36 +0700
Subject: [PATCH] Fix for async drop ice with partly dropped tuple

---
 .../rustc_mir_transform/src/elaborate_drop.rs |  2 +-
 compiler/rustc_mir_transform/src/patch.rs     | 22 ++++++++++++-------
 .../async-drop/partly-dropped-tuple.rs        | 11 ++++++++++
 3 files changed, 26 insertions(+), 9 deletions(-)
 create mode 100644 tests/ui/async-await/async-drop/partly-dropped-tuple.rs

diff --git a/compiler/rustc_mir_transform/src/elaborate_drop.rs b/compiler/rustc_mir_transform/src/elaborate_drop.rs
index 6f867f8105d50..73a58160a6aac 100644
--- a/compiler/rustc_mir_transform/src/elaborate_drop.rs
+++ b/compiler/rustc_mir_transform/src/elaborate_drop.rs
@@ -376,7 +376,7 @@ where
         if self.tcx().features().async_drop()
             && self.elaborator.body().coroutine.is_some()
             && self.elaborator.allow_async_drops()
-            && !self.elaborator.body()[bb].is_cleanup
+            && !self.elaborator.patch_ref().block(self.elaborator.body(), bb).is_cleanup
             && drop_ty.needs_async_drop(self.tcx(), self.elaborator.typing_env())
         {
             self.build_async_drop(
diff --git a/compiler/rustc_mir_transform/src/patch.rs b/compiler/rustc_mir_transform/src/patch.rs
index c7eb2a921c735..a872eae15f185 100644
--- a/compiler/rustc_mir_transform/src/patch.rs
+++ b/compiler/rustc_mir_transform/src/patch.rs
@@ -148,11 +148,20 @@ impl<'tcx> MirPatch<'tcx> {
         self.term_patch_map[bb].is_some()
     }
 
+    /// Universal getter for block data, either it is in 'old' blocks or in patched ones
+    pub(crate) fn block<'a>(
+        &'a self,
+        body: &'a Body<'tcx>,
+        bb: BasicBlock,
+    ) -> &'a BasicBlockData<'tcx> {
+        match bb.index().checked_sub(body.basic_blocks.len()) {
+            Some(new) => &self.new_blocks[new],
+            None => &body[bb],
+        }
+    }
+
     pub(crate) fn terminator_loc(&self, body: &Body<'tcx>, bb: BasicBlock) -> Location {
-        let offset = match bb.index().checked_sub(body.basic_blocks.len()) {
-            Some(index) => self.new_blocks[index].statements.len(),
-            None => body[bb].statements.len(),
-        };
+        let offset = self.block(body, bb).statements.len();
         Location { block: bb, statement_index: offset }
     }
 
@@ -284,10 +293,7 @@ impl<'tcx> MirPatch<'tcx> {
     }
 
     pub(crate) fn source_info_for_location(&self, body: &Body<'tcx>, loc: Location) -> SourceInfo {
-        let data = match loc.block.index().checked_sub(body.basic_blocks.len()) {
-            Some(new) => &self.new_blocks[new],
-            None => &body[loc.block],
-        };
+        let data = self.block(body, loc.block);
         Self::source_info_for_index(data, loc)
     }
 }
diff --git a/tests/ui/async-await/async-drop/partly-dropped-tuple.rs b/tests/ui/async-await/async-drop/partly-dropped-tuple.rs
new file mode 100644
index 0000000000000..147caaf4cfd1d
--- /dev/null
+++ b/tests/ui/async-await/async-drop/partly-dropped-tuple.rs
@@ -0,0 +1,11 @@
+//@ edition: 2024
+//@ build-pass
+#![crate_type = "lib"]
+#![allow(incomplete_features)]
+#![feature(async_drop)]
+async fn move_part_await_return_rest_tuple() -> Vec<usize> {
+    let x = (vec![3], vec![4, 4]);
+    drop(x.1);
+
+    x.0
+}