From 298ca2f6799201152e2e75a781a7aabb29424aea Mon Sep 17 00:00:00 2001
From: Eric Holk <ericholk@microsoft.com>
Date: Thu, 4 Nov 2021 10:56:07 -0700
Subject: [PATCH] Basic loop support

---
 .../src/check/generator_interior.rs           | 36 +++++++++++++++----
 src/test/ui/generator/drop-control-flow.rs    |  2 +-
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/compiler/rustc_typeck/src/check/generator_interior.rs b/compiler/rustc_typeck/src/check/generator_interior.rs
index 65644a54c4f..d7ad84684d1 100644
--- a/compiler/rustc_typeck/src/check/generator_interior.rs
+++ b/compiler/rustc_typeck/src/check/generator_interior.rs
@@ -726,15 +726,19 @@ fn intersect_drop_ranges(&mut self, drops: HirIdMap<DropRange>) {
         })
     }
 
-    fn merge_drop_ranges(&mut self, drops: HirIdMap<DropRange>) {
+    fn merge_drop_ranges_at(&mut self, drops: HirIdMap<DropRange>, join_point: usize) {
         drops.into_iter().for_each(|(k, v)| {
             if !self.drop_ranges.contains_key(&k) {
                 self.drop_ranges.insert(k, DropRange { events: vec![] });
             }
-            self.drop_ranges.get_mut(&k).unwrap().merge_with(&v, self.expr_count);
+            self.drop_ranges.get_mut(&k).unwrap().merge_with(&v, join_point);
         });
     }
 
+    fn merge_drop_ranges(&mut self, drops: HirIdMap<DropRange>) {
+        self.merge_drop_ranges_at(drops, self.expr_count);
+    }
+
     /// ExprUseVisitor's consume callback doesn't go deep enough for our purposes in all
     /// expressions. This method consumes a little deeper into the expression when needed.
     fn consume_expr(&mut self, expr: &hir::Expr<'_>) {
@@ -893,6 +897,17 @@ fn visit_expr(&mut self, expr: &'tcx Expr<'tcx>) {
 
                 reinit = Some(lhs);
             }
+            ExprKind::Loop(body, ..) => {
+                let body_drop_ranges = self.fork_drop_ranges();
+                let old_drop_ranges = self.swap_drop_ranges(body_drop_ranges);
+
+                let join_point = self.expr_count;
+
+                self.visit_block(body);
+
+                let body_drop_ranges = self.swap_drop_ranges(old_drop_ranges);
+                self.merge_drop_ranges_at(body_drop_ranges, join_point);
+            }
             _ => intravisit::walk_expr(self, expr),
         }
 
@@ -1007,11 +1022,20 @@ fn reinit(&mut self, location: usize) {
     ///
     /// After merging, the value will be dead at the end of the range only if it was dead
     /// at the end of both self and other.
-    ///
-    /// Assumes that all locations in each range are less than joinpoint
     fn merge_with(&mut self, other: &DropRange, join_point: usize) {
-        let mut events: Vec<_> =
-            self.events.iter().merge(other.events.iter()).dedup().cloned().collect();
+        let join_event = if self.is_dropped_at(join_point) && other.is_dropped_at(join_point) {
+            Event::Drop(join_point)
+        } else {
+            Event::Reinit(join_point)
+        };
+        let mut events: Vec<_> = self
+            .events
+            .iter()
+            .merge([join_event].iter())
+            .merge(other.events.iter())
+            .dedup()
+            .cloned()
+            .collect();
 
         events.push(if self.is_dropped_at(join_point) && other.is_dropped_at(join_point) {
             Event::Drop(join_point)
diff --git a/src/test/ui/generator/drop-control-flow.rs b/src/test/ui/generator/drop-control-flow.rs
index b180a61b104..6587e54df60 100644
--- a/src/test/ui/generator/drop-control-flow.rs
+++ b/src/test/ui/generator/drop-control-flow.rs
@@ -77,5 +77,5 @@ fn main() {
     one_armed_if(true);
     if_let(Some(41));
     reinit();
-    // loop_uninit();
+    loop_uninit();
 }