Handle more cases with conditionally initialized/dropped values

Author: eholk

compiler/rustc_typeck/src/check/generator_interior.rs

@@ -871,7 +871,8 @@ impl<'tcx> Visitor<'tcx> for DropRangeVisitor<'tcx> {
                         self.visit_expr(if_true);
                         true_ranges = self.swap_drop_ranges(true_ranges);
 
-                        false_ranges = self.swap_drop_ranges(false_ranges);
+                        false_ranges =
+                            self.swap_drop_ranges(trim_drop_ranges(&false_ranges, self.expr_count));
                         self.visit_expr(if_false);
                         false_ranges = self.swap_drop_ranges(false_ranges);
 
@@ -908,6 +909,31 @@ impl<'tcx> Visitor<'tcx> for DropRangeVisitor<'tcx> {
                 let body_drop_ranges = self.swap_drop_ranges(old_drop_ranges);
                 self.merge_drop_ranges_at(body_drop_ranges, join_point);
             }
+            ExprKind::Match(scrutinee, arms, ..) => {
+                self.visit_expr(scrutinee);
+
+                let forked_ranges = self.fork_drop_ranges();
+                let arm_drops = arms
+                    .iter()
+                    .map(|Arm { hir_id, pat, body, guard, .. }| {
+                        debug!("match arm {:?} starts at {}", hir_id, self.expr_count);
+                        let old_ranges = self
+                            .swap_drop_ranges(trim_drop_ranges(&forked_ranges, self.expr_count));
+                        self.visit_pat(pat);
+                        match guard {
+                            Some(Guard::If(expr)) => self.visit_expr(expr),
+                            Some(Guard::IfLet(pat, expr)) => {
+                                self.visit_pat(pat);
+                                self.visit_expr(expr);
+                            }
+                            None => (),
+                        }
+                        self.visit_expr(body);
+                        self.swap_drop_ranges(old_ranges)
+                    })
+                    .collect::<Vec<_>>();
+                arm_drops.into_iter().for_each(|drops| self.merge_drop_ranges(drops));
+            }
             _ => intravisit::walk_expr(self, expr),
         }
 
@@ -926,6 +952,10 @@ impl<'tcx> Visitor<'tcx> for DropRangeVisitor<'tcx> {
     }
 }
 
+fn trim_drop_ranges(drop_ranges: &HirIdMap<DropRange>, trim_from: usize) -> HirIdMap<DropRange> {
+    drop_ranges.iter().map(|(k, v)| (*k, v.trimmed(trim_from))).collect()
+}
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 enum Event {
     Drop(usize),
@@ -1058,4 +1088,20 @@ impl DropRange {
             }],
         }
     }
+
+    fn trimmed(&self, trim_from: usize) -> Self {
+        let start = if self.is_dropped_at(trim_from) {
+            Event::Drop(trim_from)
+        } else {
+            Event::Reinit(trim_from)
+        };
+
+        Self {
+            events: [start]
+                .iter()
+                .chain(self.events.iter().skip_while(|event| event.location() <= trim_from))
+                .cloned()
+                .collect(),
+        }
+    }
 }

src/test/ui/generator/drop-control-flow.rs

@@ -5,15 +5,14 @@
 // included in the generator type.
 
 #![feature(generators, negative_impls)]
-
 #![allow(unused_assignments, dead_code)]
 
 struct Ptr;
 impl<'a> Drop for Ptr {
     fn drop(&mut self) {}
 }
 
-struct NonSend {}
+struct NonSend;
 impl !Send for NonSend {}
 
 fn assert_send<T: Send>(_: T) {}
@@ -51,6 +50,29 @@ fn if_let(arg: Option<i32>) {
     };
 }
 
+fn init_in_if(arg: bool) {
+    assert_send(|| {
+        let mut x = NonSend;
+        drop(x);
+        if arg {
+            x = NonSend;
+        } else {
+            yield;
+        }
+    })
+}
+
+fn init_in_match_arm(arg: Option<i32>) {
+    assert_send(|| {
+        let mut x = NonSend;
+        drop(x);
+        match arg {
+            Some(_) => x = NonSend,
+            None => yield,
+        }
+    })
+}
+
 fn reinit() {
     let _ = || {
         let mut arr = [Ptr];
@@ -73,9 +95,27 @@ fn loop_uninit() {
     };
 }
 
+fn nested_loop() {
+    let _ = || {
+        let mut arr = [Ptr];
+        let mut count = 0;
+        drop(arr);
+        while count < 3 {
+            for _ in 0..3 {
+                yield;
+            }
+            arr = [Ptr];
+            count += 1;
+        }
+    };
+}
+
 fn main() {
     one_armed_if(true);
     if_let(Some(41));
+    init_in_if(true);
+    init_in_match_arm(Some(41));
     reinit();
     loop_uninit();
+    nested_loop();
 }