From 209087b8fa26f28c847a5a98afa0bc52ed4f769b Mon Sep 17 00:00:00 2001
From: Dylan MacKenzie <ecstaticmorse@gmail.com>
Date: Thu, 9 Apr 2020 13:02:23 -0700
Subject: [PATCH] Use `Visitor` for `AlwaysLiveLocals`

---
 src/librustc_mir/util/storage.rs | 33 ++++++++++++++++++--------------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/librustc_mir/util/storage.rs b/src/librustc_mir/util/storage.rs
index 2ce9bed794d..0b7b1c29537 100644
--- a/src/librustc_mir/util/storage.rs
+++ b/src/librustc_mir/util/storage.rs
@@ -1,5 +1,6 @@
 use rustc_index::bit_set::BitSet;
-use rustc_middle::mir::{self, Local};
+use rustc_middle::mir::visit::Visitor;
+use rustc_middle::mir::{self, Local, Location};
 
 /// The set of locals in a MIR body that do not have `StorageLive`/`StorageDead` annotations.
 ///
@@ -12,20 +13,12 @@
 
 impl AlwaysLiveLocals {
     pub fn new(body: &mir::Body<'tcx>) -> Self {
-        let mut locals = BitSet::new_filled(body.local_decls.len());
-
-        // FIXME: Use a visitor for this when `visit_body` can take a plain `Body`.
-        for block in body.basic_blocks().iter() {
-            for stmt in &block.statements {
-                if let mir::StatementKind::StorageLive(l) | mir::StatementKind::StorageDead(l) =
-                    stmt.kind
-                {
-                    locals.remove(l);
-                }
-            }
-        }
+        let mut ret = AlwaysLiveLocals(BitSet::new_filled(body.local_decls.len()));
+
+        let mut vis = StorageAnnotationVisitor(&mut ret);
+        vis.visit_body(body);
 
-        AlwaysLiveLocals(locals)
+        ret
     }
 
     pub fn into_inner(self) -> BitSet<Local> {
@@ -40,3 +33,15 @@ fn deref(&self) -> &Self::Target {
         &self.0
     }
 }
+
+/// Removes locals that have `Storage*` annotations from `AlwaysLiveLocals`.
+struct StorageAnnotationVisitor<'a>(&'a mut AlwaysLiveLocals);
+
+impl Visitor<'tcx> for StorageAnnotationVisitor<'_> {
+    fn visit_statement(&mut self, statement: &mir::Statement<'tcx>, _location: Location) {
+        use mir::StatementKind::{StorageDead, StorageLive};
+        if let StorageLive(l) | StorageDead(l) = statement.kind {
+            (self.0).0.remove(l);
+        }
+    }
+}
-- 
2.44.0