From 0babbf11e6b1b4aa95f30e01373d3f1488eca9c4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= <john.kare.alsaker@gmail.com>
Date: Tue, 31 Jul 2018 21:54:30 +0200
Subject: [PATCH] Don't count MIR locals as borrowed after StorageDead when
 finding locals live across a yield terminator

---
 .../dataflow/impls/borrowed_locals.rs         | 16 +++++++++---
 src/librustc_mir/transform/generator.rs       |  3 ++-
 .../generator/yield-in-initializer.rs         | 25 +++++++++++++++++++
 3 files changed, 39 insertions(+), 5 deletions(-)
 create mode 100644 src/test/run-pass/generator/yield-in-initializer.rs

diff --git a/src/librustc_mir/dataflow/impls/borrowed_locals.rs b/src/librustc_mir/dataflow/impls/borrowed_locals.rs
index 96f4c6b60f5..c8c41c13b0f 100644
--- a/src/librustc_mir/dataflow/impls/borrowed_locals.rs
+++ b/src/librustc_mir/dataflow/impls/borrowed_locals.rs
@@ -15,9 +15,9 @@ use rustc::mir::visit::Visitor;
 use dataflow::BitDenotation;
 
 /// This calculates if any part of a MIR local could have previously been borrowed.
-/// This means that once a local has been borrowed, its bit will always be set
-/// from that point and onwards, even if the borrow ends. You could also think of this
-/// as computing the lifetimes of infinite borrows.
+/// This means that once a local has been borrowed, its bit will be set
+/// from that point and onwards, until we see a StorageDead statement for the local,
+/// at which points there is no memory associated with the local, so it cannot be borrowed.
 /// This is used to compute which locals are live during a yield expression for
 /// immovable generators.
 #[derive(Copy, Clone)]
@@ -50,9 +50,17 @@ impl<'a, 'tcx> BitDenotation for HaveBeenBorrowedLocals<'a, 'tcx> {
     fn statement_effect(&self,
                         sets: &mut BlockSets<Local>,
                         loc: Location) {
+        let stmt = &self.mir[loc.block].statements[loc.statement_index];
+
         BorrowedLocalsVisitor {
             sets,
-        }.visit_statement(loc.block, &self.mir[loc.block].statements[loc.statement_index], loc);
+        }.visit_statement(loc.block, stmt, loc);
+
+        // StorageDead invalidates all borrows and raw pointers to a local
+        match stmt.kind {
+            StatementKind::StorageDead(l) => sets.kill(&l),
+            _ => (),
+        }
     }
 
     fn terminator_effect(&self,
diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index 6a9258fe2c9..a3647edd155 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -433,7 +433,8 @@ fn locals_live_across_suspend_points<'a, 'tcx,>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                 // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
                 // This is correct for movable generators since borrows cannot live across
                 // suspension points. However for immovable generators we need to account for
-                // borrows, so we conseratively assume that all borrowed locals live forever.
+                // borrows, so we conseratively assume that all borrowed locals are live until
+                // we find a StorageDead statement referencing the locals.
                 // To do this we just union our `liveness` result with `borrowed_locals`, which
                 // contains all the locals which has been borrowed before this suspension point.
                 // If a borrow is converted to a raw reference, we must also assume that it lives
diff --git a/src/test/run-pass/generator/yield-in-initializer.rs b/src/test/run-pass/generator/yield-in-initializer.rs
new file mode 100644
index 00000000000..3042061226b
--- /dev/null
+++ b/src/test/run-pass/generator/yield-in-initializer.rs
@@ -0,0 +1,25 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#![feature(generators)]
+
+fn main() {
+    static || {
+        loop {
+            // Test that `opt` is not live across the yield, even when borrowed in a loop
+            // See https://github.com/rust-lang/rust/issues/52792
+            let opt = {
+                yield;
+                true
+            };
+            &opt;
+        }
+    };
+}