Auto merge of #102394 - dingxiangfei2009:issue-102317, r=oli-obk

Fix unwind drop glue for if-then scopes

cc `@est31`

Fix #102317
Fix #99852

This PR fixes the drop glue for unwinding from a panic originated in a drop while breaking out for the else block in an `if-then` scope.
MIR validation does not fail for the synchronous versions of the test program, because `StorageDead` statements are skipped over in the unwinding process. It is only becoming a problem when it is inside a generator where `StorageDead` must be kept around.

compiler/rustc_mir_build/src/build/block.rs

@@ -245,11 +245,6 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                                                 OutsideGuard,
                                                 true,
                                             );
-                                            this.schedule_drop_for_binding(
-                                                node,
-                                                span,
-                                                OutsideGuard,
-                                            );
                                         },
                                     );
                                     this.ast_let_else(

compiler/rustc_mir_build/src/build/expr/into.rs

@@ -74,7 +74,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                                 this.source_info(then_expr.span)
                             };
                             let (then_block, else_block) =
-                                this.in_if_then_scope(condition_scope, |this| {
+                                this.in_if_then_scope(condition_scope, then_expr.span, |this| {
                                     let then_blk = unpack!(this.then_else_break(
                                         block,
                                         &this.thir[cond],
@@ -107,7 +107,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             }
             ExprKind::Let { expr, ref pat } => {
                 let scope = this.local_scope();
-                let (true_block, false_block) = this.in_if_then_scope(scope, |this| {
+                let (true_block, false_block) = this.in_if_then_scope(scope, expr_span, |this| {
                     this.lower_let_expr(block, &this.thir[expr], pat, scope, None, expr_span)
                 });
 

compiler/rustc_mir_build/src/build/matches/mod.rs

@@ -1986,7 +1986,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             let mut guard_span = rustc_span::DUMMY_SP;
 
             let (post_guard_block, otherwise_post_guard_block) =
-                self.in_if_then_scope(match_scope, |this| match *guard {
+                self.in_if_then_scope(match_scope, guard_span, |this| match *guard {
                     Guard::If(e) => {
                         let e = &this.thir[e];
                         guard_span = e.span;
@@ -2301,7 +2301,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         pattern: &Pat<'tcx>,
     ) -> BlockAnd<BasicBlock> {
         let else_block_span = self.thir[else_block].span;
-        let (matching, failure) = self.in_if_then_scope(*let_else_scope, |this| {
+        let (matching, failure) = self.in_if_then_scope(*let_else_scope, else_block_span, |this| {
             let scrutinee = unpack!(block = this.lower_scrutinee(block, init, initializer_span));
             let pat = Pat { ty: init.ty, span: else_block_span, kind: PatKind::Wild };
             let mut wildcard = Candidate::new(scrutinee.clone(), &pat, false, this);

compiler/rustc_mir_build/src/build/scope.rs

@@ -466,9 +466,10 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         let normal_exit_block = f(self);
         let breakable_scope = self.scopes.breakable_scopes.pop().unwrap();
         assert!(breakable_scope.region_scope == region_scope);
-        let break_block = self.build_exit_tree(breakable_scope.break_drops, None);
+        let break_block =
+            self.build_exit_tree(breakable_scope.break_drops, region_scope, span, None);
         if let Some(drops) = breakable_scope.continue_drops {
-            self.build_exit_tree(drops, loop_block);
+            self.build_exit_tree(drops, region_scope, span, loop_block);
         }
         match (normal_exit_block, break_block) {
             (Some(block), None) | (None, Some(block)) => block,
@@ -510,6 +511,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     pub(crate) fn in_if_then_scope<F>(
         &mut self,
         region_scope: region::Scope,
+        span: Span,
         f: F,
     ) -> (BasicBlock, BasicBlock)
     where
@@ -524,7 +526,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         assert!(if_then_scope.region_scope == region_scope);
 
         let else_block = self
-            .build_exit_tree(if_then_scope.else_drops, None)
+            .build_exit_tree(if_then_scope.else_drops, region_scope, span, None)
             .map_or_else(|| self.cfg.start_new_block(), |else_block_and| unpack!(else_block_and));
 
         (then_block, else_block)
@@ -997,10 +999,18 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     /// Returns the [DropIdx] for the innermost drop if the function unwound at
     /// this point. The `DropIdx` will be created if it doesn't already exist.
     fn diverge_cleanup(&mut self) -> DropIdx {
-        let is_generator = self.generator_kind.is_some();
-        let (uncached_scope, mut cached_drop) = self
-            .scopes
-            .scopes
+        // It is okay to use dummy span because the getting scope index on the topmost scope
+        // must always succeed.
+        self.diverge_cleanup_target(self.scopes.topmost(), DUMMY_SP)
+    }
+
+    /// This is similar to [diverge_cleanup](Self::diverge_cleanup) except its target is set to
+    /// some ancestor scope instead of the current scope.
+    /// It is possible to unwind to some ancestor scope if some drop panics as
+    /// the program breaks out of a if-then scope.
+    fn diverge_cleanup_target(&mut self, target_scope: region::Scope, span: Span) -> DropIdx {
+        let target = self.scopes.scope_index(target_scope, span);
+        let (uncached_scope, mut cached_drop) = self.scopes.scopes[..=target]
             .iter()
             .enumerate()
             .rev()
@@ -1009,7 +1019,12 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             })
             .unwrap_or((0, ROOT_NODE));
 
-        for scope in &mut self.scopes.scopes[uncached_scope..] {
+        if uncached_scope > target {
+            return cached_drop;
+        }
+
+        let is_generator = self.generator_kind.is_some();
+        for scope in &mut self.scopes.scopes[uncached_scope..=target] {
             for drop in &scope.drops {
                 if is_generator || drop.kind == DropKind::Value {
                     cached_drop = self.scopes.unwind_drops.add_drop(*drop, cached_drop);
@@ -1222,21 +1237,24 @@ impl<'a, 'tcx: 'a> Builder<'a, 'tcx> {
     fn build_exit_tree(
         &mut self,
         mut drops: DropTree,
+        else_scope: region::Scope,
+        span: Span,
         continue_block: Option<BasicBlock>,
     ) -> Option<BlockAnd<()>> {
         let mut blocks = IndexVec::from_elem(None, &drops.drops);
         blocks[ROOT_NODE] = continue_block;
 
         drops.build_mir::<ExitScopes>(&mut self.cfg, &mut blocks);
+        let is_generator = self.generator_kind.is_some();
 
         // Link the exit drop tree to unwind drop tree.
         if drops.drops.iter().any(|(drop, _)| drop.kind == DropKind::Value) {
-            let unwind_target = self.diverge_cleanup();
+            let unwind_target = self.diverge_cleanup_target(else_scope, span);
             let mut unwind_indices = IndexVec::from_elem_n(unwind_target, 1);
             for (drop_idx, drop_data) in drops.drops.iter_enumerated().skip(1) {
                 match drop_data.0.kind {
                     DropKind::Storage => {
-                        if self.generator_kind.is_some() {
+                        if is_generator {
                             let unwind_drop = self
                                 .scopes
                                 .unwind_drops

src/test/ui/let-else/issue-102317.rs

@@ -0,0 +1,20 @@
+// issue #102317
+// build-pass
+// compile-flags: --edition 2021 -C opt-level=3 -Zvalidate-mir
+
+struct SegmentJob;
+
+impl Drop for SegmentJob {
+    fn drop(&mut self) {}
+}
+
+pub async fn run() -> Result<(), ()> {
+    let jobs = Vec::<SegmentJob>::new();
+    let Some(_job) = jobs.into_iter().next() else {
+        return Ok(())
+    };
+
+    Ok(())
+}
+
+fn main() {}

src/test/ui/mir/issue-99852.rs

@@ -0,0 +1,24 @@
+// check-pass
+// compile-flags: -Z validate-mir
+#![feature(let_chains)]
+
+fn lambda<T, U>() -> U
+where
+    T: Default,
+    U: Default,
+{
+    let foo: Result<T, ()> = Ok(T::default());
+    let baz: U = U::default();
+
+    if let Ok(foo) = foo && let Ok(bar) = transform(foo) {
+        bar
+    } else {
+        baz
+    }
+}
+
+fn transform<T, U>(input: T) -> Result<U, ()> {
+    todo!()
+}
+
+fn main() {}