fix unwind drop glue for if-then scopes

Author: dingxiangfei2009

Diff for: 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(

Diff for: 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)
                 });
 

Diff for: 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);

Diff for: 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)
@@ -1021,6 +1023,38 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         cached_drop
     }
 
+    /// This is similar to [diverge_cleanup_target] 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()
+            .find_map(|(scope_idx, scope)| {
+                scope.cached_unwind_block.map(|cached_block| (scope_idx + 1, cached_block))
+            })
+            .unwrap_or((0, ROOT_NODE));
+
+        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);
+                }
+            }
+            scope.cached_unwind_block = Some(cached_drop);
+        }
+
+        cached_drop
+    }
+
     /// Prepares to create a path that performs all required cleanup for a
     /// terminator that can unwind at the given basic block.
     ///
@@ -1222,21 +1256,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

Diff for: src/test/ui/let-else/issue-102317.rs

@@ -0,0 +1,20 @@
+// issue #102317
+// run-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() {}

Diff for: 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() {}