Fix drop handling for `if let` expressions

MIR lowering for `if let` expressions is now more complicated now that
`if let` exists in HIR. This PR adds a scope for the variables bound in
an `if let` expression and then uses an approach similar to how we
handle loops to ensure that we reliably drop the correct variables.

5 files changed, 183 insertions, 110 deletions

diff --git a/compiler/rustc_mir_build/src/build/expr/into.rs b/compiler/rustc_mir_build/src/build/expr/into.rs
index 21c5a64f823..1a7bb602392 100644
--- a/compiler/rustc_mir_build/src/build/expr/into.rs
+++ b/compiler/rustc_mir_build/src/build/expr/into.rs
@@ -52,11 +52,33 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             ExprKind::Match { scrutinee, ref arms } => {
                 this.match_expr(destination, expr_span, block, &this.thir[scrutinee], arms)
             }
-            ExprKind::If { cond, then, else_opt } => {
-                let local_scope = this.local_scope();
-                let (mut then_blk, mut else_blk) =
-                    this.then_else_blocks(block, &this.thir[cond], local_scope, source_info);
-                unpack!(then_blk = this.expr_into_dest(destination, then_blk, &this.thir[then]));
+            ExprKind::If { cond, then, else_opt, if_then_scope } => {
+                let then_blk;
+                let then_expr = &this.thir[then];
+                let then_source_info = this.source_info(then_expr.span);
+                let condition_scope = this.local_scope();
+
+                let mut else_blk = unpack!(
+                    then_blk = this.in_scope(
+                        (if_then_scope, then_source_info),
+                        LintLevel::Inherited,
+                        |this| {
+                            let (then_block, else_block) =
+                                this.in_if_then_scope(condition_scope, |this| {
+                                    let then_blk = unpack!(this.then_else_break(
+                                        block,
+                                        &this.thir[cond],
+                                        condition_scope,
+                                        condition_scope,
+                                        then_expr.span,
+                                    ));
+                                    this.expr_into_dest(destination, then_blk, then_expr)
+                                });
+                            then_block.and(else_block)
+                        },
+                    )
+                );
+
                 else_blk = if let Some(else_opt) = else_opt {
                     unpack!(this.expr_into_dest(destination, else_blk, &this.thir[else_opt]))
                 } else {
@@ -81,9 +103,11 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
                 join_block.unit()
             }
-            ExprKind::Let { ref pat, expr } => {
-                let (true_block, false_block) =
-                    this.lower_let(block, &this.thir[expr], pat, expr_span);
+            ExprKind::Let { expr, ref pat } => {
+                let scope = this.local_scope();
+                let (true_block, false_block) = this.in_if_then_scope(scope, |this| {
+                    this.lower_let_else(block, &this.thir[expr], pat, scope, expr_span)
+                });
 
                 let join_block = this.cfg.start_new_block();
 
diff --git a/compiler/rustc_mir_build/src/build/matches/mod.rs b/compiler/rustc_mir_build/src/build/matches/mod.rs
index 7c4bdf1066a..d13c4f5a5a0 100644
--- a/compiler/rustc_mir_build/src/build/matches/mod.rs
+++ b/compiler/rustc_mir_build/src/build/matches/mod.rs
@@ -35,42 +35,49 @@ use std::convert::TryFrom;
 use std::mem;
 
 impl<'a, 'tcx> Builder<'a, 'tcx> {
-    pub(crate) fn then_else_blocks(
+    pub(crate) fn then_else_break(
         &mut self,
         mut block: BasicBlock,
         expr: &Expr<'tcx>,
-        scope: region::Scope,
-        source_info: SourceInfo,
-    ) -> (BasicBlock, BasicBlock) {
+        temp_scope: region::Scope,
+        break_scope: region::Scope,
+        variable_scope_span: Span,
+    ) -> BlockAnd<()> {
         let this = self;
         let expr_span = expr.span;
 
         match expr.kind {
             ExprKind::Scope { region_scope, lint_level, value } => {
-                let region_scope = (region_scope, source_info);
-                let then_block;
-                let else_block = unpack!(
-                    then_block = this.in_scope(region_scope, lint_level, |this| {
-                        let (then_block, else_block) =
-                            this.then_else_blocks(block, &this.thir[value], scope, source_info);
-                        then_block.and(else_block)
-                    })
-                );
-                (then_block, else_block)
+                let region_scope = (region_scope, this.source_info(expr_span));
+                this.in_scope(region_scope, lint_level, |this| {
+                    this.then_else_break(
+                        block,
+                        &this.thir[value],
+                        temp_scope,
+                        break_scope,
+                        variable_scope_span,
+                    )
+                })
             }
             ExprKind::Let { expr, ref pat } => {
-                // FIXME: Use correct span.
-                this.lower_let(block, &this.thir[expr], pat, expr_span)
+                this.lower_let_else(block, &this.thir[expr], pat, break_scope, variable_scope_span)
             }
             _ => {
+                // TODO `as_temp`?
                 let mutability = Mutability::Mut;
-                let place = unpack!(block = this.as_temp(block, Some(scope), expr, mutability));
+                let place =
+                    unpack!(block = this.as_temp(block, Some(temp_scope), expr, mutability));
                 let operand = Operand::Move(Place::from(place));
+
                 let then_block = this.cfg.start_new_block();
                 let else_block = this.cfg.start_new_block();
                 let term = TerminatorKind::if_(this.tcx, operand, then_block, else_block);
+
+                let source_info = this.source_info(expr_span);
                 this.cfg.terminate(block, source_info, term);
-                (then_block, else_block)
+                this.break_for_else(else_block, break_scope, source_info);
+
+                then_block.unit()
             }
         }
     }
@@ -302,6 +309,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
                 let arm_source_info = self.source_info(arm.span);
                 let arm_scope = (arm.scope, arm_source_info);
+                let match_scope = self.local_scope();
                 self.in_scope(arm_scope, arm.lint_level, |this| {
                     // `try_upvars_resolved` may fail if it is unable to resolve the given
                     // `PlaceBuilder` inside a closure. In this case, we don't want to include
@@ -340,6 +348,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         scrutinee_span,
                         Some(arm.span),
                         Some(arm.scope),
+                        Some(match_scope),
                     );
 
                     if let Some(source_scope) = scope {
@@ -384,6 +393,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         scrutinee_span: Span,
         arm_span: Option<Span>,
         arm_scope: Option<region::Scope>,
+        match_scope: Option<region::Scope>,
     ) -> BasicBlock {
         if candidate.subcandidates.is_empty() {
             // Avoid generating another `BasicBlock` when we only have one
@@ -395,6 +405,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 fake_borrow_temps,
                 scrutinee_span,
                 arm_span,
+                match_scope,
                 true,
             )
         } else {
@@ -431,6 +442,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         &fake_borrow_temps,
                         scrutinee_span,
                         arm_span,
+                        match_scope,
                         schedule_drops,
                     );
                     if arm_scope.is_none() {
@@ -616,6 +628,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             irrefutable_pat.span,
             None,
             None,
+            None,
         )
         .unit()
     }
@@ -1742,13 +1755,14 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 // Pat binding - used for `let` and function parameters as well.
 
 impl<'a, 'tcx> Builder<'a, 'tcx> {
-    pub fn lower_let(
+    crate fn lower_let_else(
         &mut self,
         mut block: BasicBlock,
         expr: &Expr<'tcx>,
         pat: &Pat<'tcx>,
+        else_target: region::Scope,
         span: Span,
-    ) -> (BasicBlock, BasicBlock) {
+    ) -> BlockAnd<()> {
         let expr_span = expr.span;
         let expr_place_builder = unpack!(block = self.lower_scrutinee(block, expr, expr_span));
         let mut guard_candidate = Candidate::new(expr_place_builder.clone(), &pat, false);
@@ -1769,6 +1783,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             expr_place = expr_builder.into_place(self.tcx, self.typeck_results);
             opt_expr_place = Some((Some(&expr_place), expr_span));
         }
+        let otherwise_post_guard_block = otherwise_candidate.pre_binding_block.unwrap();
+        self.break_for_else(otherwise_post_guard_block, else_target, self.source_info(expr_span));
+
         self.declare_bindings(None, pat.span.to(span), pat, ArmHasGuard(false), opt_expr_place);
         let post_guard_block = self.bind_pattern(
             self.source_info(pat.span),
@@ -1778,9 +1795,10 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             expr.span,
             None,
             None,
+            None,
         );
-        let otherwise_post_guard_block = otherwise_candidate.pre_binding_block.unwrap();
-        (post_guard_block, otherwise_post_guard_block)
+
+        post_guard_block.unit()
     }
 
     /// Initializes each of the bindings from the candidate by
@@ -1799,6 +1817,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         fake_borrows: &Vec<(Place<'tcx>, Local)>,
         scrutinee_span: Span,
         arm_span: Option<Span>,
+        match_scope: Option<region::Scope>,
         schedule_drops: bool,
     ) -> BasicBlock {
         debug!("bind_and_guard_matched_candidate(candidate={:?})", candidate);
@@ -1929,17 +1948,25 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 self.cfg.push_assign(block, scrutinee_source_info, Place::from(temp), borrow);
             }
 
-            let (guard_span, (post_guard_block, otherwise_post_guard_block)) = match *guard {
-                Guard::If(e) => {
-                    let e = &self.thir[e];
-                    let source_info = self.source_info(e.span);
-                    (e.span, self.test_bool(block, e, source_info))
-                }
-                Guard::IfLet(ref pat, scrutinee) => {
-                    let s = &self.thir[scrutinee];
-                    (s.span, self.lower_let(block, s, pat, arm_span.unwrap()))
-                }
-            };
+            let arm_span = arm_span.unwrap();
+            let arm_scope = self.local_scope();
+            let match_scope = match_scope.unwrap();
+            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 {
+                    Guard::If(e) => {
+                        let e = &this.thir[e];
+                        guard_span = e.span;
+                        this.then_else_break(block, e, arm_scope, match_scope, arm_span)
+                    }
+                    Guard::IfLet(ref pat, scrutinee) => {
+                        let s = &this.thir[scrutinee];
+                        guard_span = s.span;
+                        this.lower_let_else(block, s, pat, match_scope, arm_span)
+                    }
+                });
+
             let source_info = self.source_info(guard_span);
             let guard_end = self.source_info(tcx.sess.source_map().end_point(guard_span));
             let guard_frame = self.guard_context.pop().unwrap();
@@ -1955,10 +1982,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 self.cfg.terminate(unreachable, source_info, TerminatorKind::Unreachable);
                 unreachable
             });
-            let outside_scope = self.cfg.start_new_block();
-            self.exit_top_scope(otherwise_post_guard_block, outside_scope, source_info);
             self.false_edges(
-                outside_scope,
+                otherwise_post_guard_block,
                 otherwise_block,
                 candidate.next_candidate_pre_binding_block,
                 source_info,
diff --git a/compiler/rustc_mir_build/src/build/scope.rs b/compiler/rustc_mir_build/src/build/scope.rs
index 496db58758c..bd8d14fcd01 100644
--- a/compiler/rustc_mir_build/src/build/scope.rs
+++ b/compiler/rustc_mir_build/src/build/scope.rs
@@ -81,6 +81,8 @@ that contains only loops and breakable blocks. It tracks where a `break`,
 
 */
 
+use std::mem;
+
 use crate::build::{BlockAnd, BlockAndExtension, BlockFrame, Builder, CFG};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_index::vec::IndexVec;
@@ -93,9 +95,13 @@ use rustc_span::{Span, DUMMY_SP};
 #[derive(Debug)]
 pub struct Scopes<'tcx> {
     scopes: Vec<Scope>,
+
     /// The current set of breakable scopes. See module comment for more details.
     breakable_scopes: Vec<BreakableScope<'tcx>>,
 
+    /// The scope of the innermost if-then currently being lowered.
+    if_then_scope: Option<IfThenScope>,
+
     /// Drops that need to be done on unwind paths. See the comment on
     /// [DropTree] for more details.
     unwind_drops: DropTree,
@@ -164,6 +170,14 @@ struct BreakableScope<'tcx> {
     continue_drops: Option<DropTree>,
 }
 
+#[derive(Debug)]
+struct IfThenScope {
+    /// The if-then scope or arm scope
+    region_scope: region::Scope,
+    /// Drops that happen on the `else` path.
+    else_drops: DropTree,
+}
+
 /// The target of an expression that breaks out of a scope
 #[derive(Clone, Copy, Debug)]
 crate enum BreakableTarget {
@@ -183,6 +197,7 @@ const ROOT_NODE: DropIdx = DropIdx::from_u32(0);
 /// * Drops on unwind paths
 /// * Drops on generator drop paths (when a suspended generator is dropped)
 /// * Drops on return and loop exit paths
+/// * Drops on the else path in an `if let` chain
 ///
 /// Once no more nodes could be added to the tree, we lower it to MIR in one go
 /// in `build_mir`.
@@ -394,6 +409,7 @@ impl<'tcx> Scopes<'tcx> {
         Self {
             scopes: Vec::new(),
             breakable_scopes: Vec::new(),
+            if_then_scope: None,
             unwind_drops: DropTree::new(),
             generator_drops: DropTree::new(),
         }
@@ -483,6 +499,45 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         }
     }
 
+    /// Start an if-then scope which tracks drop for `if` expressions and `if`
+    /// guards.
+    ///
+    /// For an if-let chain:
+    ///
+    /// if let Some(x) = a && let Some(y) = b && let Some(z) = c { ... }
+    ///
+    /// there are three possible ways the condition can be false and we may have
+    /// to drop `x`, `x` and `y`, or neither depending on which binding fails.
+    /// To handle this correctly we use a `DropTree` in a similar way to a
+    /// `loop` expression and 'break' out on all of the 'else' paths.
+    ///
+    /// Notes:
+    /// - We don't need to keep a stack of scopes in the `Builder` because the
+    ///   'else' paths will only leave the innermost scope.
+    /// - This is also used for match guards.
+    crate fn in_if_then_scope<F>(
+        &mut self,
+        region_scope: region::Scope,
+        f: F,
+    ) -> (BasicBlock, BasicBlock)
+    where
+        F: FnOnce(&mut Builder<'a, 'tcx>) -> BlockAnd<()>,
+    {
+        let scope = IfThenScope { region_scope, else_drops: DropTree::new() };
+        let previous_scope = mem::replace(&mut self.scopes.if_then_scope, Some(scope));
+
+        let then_block = unpack!(f(self));
+
+        let if_then_scope = mem::replace(&mut self.scopes.if_then_scope, previous_scope).unwrap();
+        assert!(if_then_scope.region_scope == region_scope);
+
+        let else_block = self
+            .build_exit_tree(if_then_scope.else_drops, None)
+            .map_or_else(|| self.cfg.start_new_block(), |else_block_and| unpack!(else_block_and));
+
+        (then_block, else_block)
+    }
+
     crate fn in_opt_scope<F, R>(
         &mut self,
         opt_scope: Option<(region::Scope, SourceInfo)>,
@@ -651,6 +706,36 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         self.cfg.start_new_block().unit()
     }
 
+    crate fn break_for_else(
+        &mut self,
+        block: BasicBlock,
+        target: region::Scope,
+        source_info: SourceInfo,
+    ) {
+        let scope_index = self.scopes.scope_index(target, source_info.span);
+        let if_then_scope = self
+            .scopes
+            .if_then_scope
+            .as_mut()
+            .unwrap_or_else(|| span_bug!(source_info.span, "no if-then scope found"));
+
+        assert_eq!(if_then_scope.region_scope, target, "breaking to incorrect scope");
+
+        let mut drop_idx = ROOT_NODE;
+        let drops = &mut if_then_scope.else_drops;
+        for scope in &self.scopes.scopes[scope_index + 1..] {
+            for drop in &scope.drops {
+                drop_idx = drops.add_drop(*drop, drop_idx);
+            }
+        }
+        drops.add_entry(block, drop_idx);
+
+        // `build_drop_tree` doesn't have access to our source_info, so we
+        // create a dummy terminator now. `TerminatorKind::Resume` is used
+        // because MIR type checking will panic if it hasn't been overwritten.
+        self.cfg.terminate(block, source_info, TerminatorKind::Resume);
+    }
+
     // Add a dummy `Assign` statement to the CFG, with the span for the source code's `continue`
     // statement.
     fn add_dummy_assignment(&mut self, span: &Span, block: BasicBlock, source_info: SourceInfo) {
@@ -659,16 +744,6 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         self.cfg.push_assign_unit(block, source_info, temp_place, self.tcx);
     }
 
-    crate fn exit_top_scope(
-        &mut self,
-        mut block: BasicBlock,
-        target: BasicBlock,
-        source_info: SourceInfo,
-    ) {
-        block = self.leave_top_scope(block);
-        self.cfg.terminate(block, source_info, TerminatorKind::Goto { target });
-    }
-
     fn leave_top_scope(&mut self, block: BasicBlock) -> BasicBlock {
         // If we are emitting a `drop` statement, we need to have the cached
         // diverge cleanup pads ready in case that drop panics.
@@ -927,61 +1002,6 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
     // Other
     // =====
-    /// Branch based on a boolean condition.
-    ///
-    /// This is a special case because the temporary for the condition needs to
-    /// be dropped on both the true and the false arm.
-    crate fn test_bool(
-        &mut self,
-        mut block: BasicBlock,
-        condition: &Expr<'tcx>,
-        source_info: SourceInfo,
-    ) -> (BasicBlock, BasicBlock) {
-        let cond = unpack!(block = self.as_local_operand(block, condition));
-        let true_block = self.cfg.start_new_block();
-        let false_block = self.cfg.start_new_block();
-        let term = TerminatorKind::if_(self.tcx, cond.clone(), true_block, false_block);
-        self.cfg.terminate(block, source_info, term);
-
-        match cond {
-            // Don't try to drop a constant
-            Operand::Constant(_) => (),
-            Operand::Copy(place) | Operand::Move(place) => {
-                if let Some(cond_temp) = place.as_local() {
-                    // Manually drop the condition on both branches.
-                    let top_scope = self.scopes.scopes.last_mut().unwrap();
-                    let top_drop_data = top_scope.drops.pop().unwrap();
-                    if self.generator_kind.is_some() {
-                        top_scope.invalidate_cache();
-                    }
-
-                    match top_drop_data.kind {
-                        DropKind::Value { .. } => {
-                            bug!("Drop scheduled on top of condition variable")
-                        }
-                        DropKind::Storage => {
-                            let source_info = top_drop_data.source_info;
-                            let local = top_drop_data.local;
-                            assert_eq!(local, cond_temp, "Drop scheduled on top of condition");
-                            self.cfg.push(
-                                true_block,
-                                Statement { source_info, kind: StatementKind::StorageDead(local) },
-                            );
-                            self.cfg.push(
-                                false_block,
-                                Statement { source_info, kind: StatementKind::StorageDead(local) },
-                            );
-                        }
-                    }
-                } else {
-                    bug!("Expected as_local_operand to produce a temporary");
-                }
-            }
-        }
-
-        (true_block, false_block)
-    }
-
     /// 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 {
diff --git a/compiler/rustc_mir_build/src/thir/cx/expr.rs b/compiler/rustc_mir_build/src/thir/cx/expr.rs
index 9106f1d2342..66005be05df 100644
--- a/compiler/rustc_mir_build/src/thir/cx/expr.rs
+++ b/compiler/rustc_mir_build/src/thir/cx/expr.rs
@@ -594,6 +594,10 @@ impl<'tcx> Cx<'tcx> {
                 ExprKind::Let { expr: self.mirror_expr(expr), pat: self.pattern_from_hir(pat) }
             }
             hir::ExprKind::If(cond, then, else_opt) => ExprKind::If {
+                if_then_scope: region::Scope {
+                    id: then.hir_id.local_id,
+                    data: region::ScopeData::IfThen,
+                },
                 cond: self.mirror_expr(cond),
                 then: self.mirror_expr(then),
                 else_opt: else_opt.map(|el| self.mirror_expr(el)),
diff --git a/compiler/rustc_mir_build/src/thir/visit.rs b/compiler/rustc_mir_build/src/thir/visit.rs
index dda6c74d4de..51c371b8720 100644
--- a/compiler/rustc_mir_build/src/thir/visit.rs
+++ b/compiler/rustc_mir_build/src/thir/visit.rs
@@ -34,7 +34,7 @@ pub fn walk_expr<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, expr: &Exp
             visitor.visit_expr(&visitor.thir()[value])
         }
         Box { value } => visitor.visit_expr(&visitor.thir()[value]),
-        If { cond, then, else_opt } => {
+        If { cond, then, else_opt, if_then_scope: _ } => {
             visitor.visit_expr(&visitor.thir()[cond]);
             visitor.visit_expr(&visitor.thir()[then]);
             if let Some(else_expr) = else_opt {