From e22cc993fbf9918784ec38a25e04564e7b1973ee Mon Sep 17 00:00:00 2001
From: Jonas Schievink <jonasschievink@gmail.com>
Date: Sat, 9 May 2020 16:08:04 +0200
Subject: [PATCH] Visit move out of `_0` when visiting `return`

---
 src/librustc_middle/mir/visit.rs        | 18 +++++++++++++++++-
 src/librustc_mir/transform/copy_prop.rs |  7 ++++++-
 src/librustc_mir/transform/generator.rs | 10 ++++++++++
 src/librustc_mir/transform/inline.rs    |  6 +++++-
 4 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/librustc_middle/mir/visit.rs b/src/librustc_middle/mir/visit.rs
index 2f3d89dc029..1f097f24942 100644
--- a/src/librustc_middle/mir/visit.rs
+++ b/src/librustc_middle/mir/visit.rs
@@ -427,13 +427,29 @@ fn super_terminator_kind(&mut self,
                     TerminatorKind::Goto { .. } |
                     TerminatorKind::Resume |
                     TerminatorKind::Abort |
-                    TerminatorKind::Return |
                     TerminatorKind::GeneratorDrop |
                     TerminatorKind::Unreachable |
                     TerminatorKind::FalseEdges { .. } |
                     TerminatorKind::FalseUnwind { .. } => {
                     }
 
+                    TerminatorKind::Return => {
+                        // `return` logically moves from the return place `_0`. Note that the place
+                        // cannot be changed by any visitor, though.
+                        let $($mutability)? local = RETURN_PLACE;
+                        self.visit_local(
+                            & $($mutability)? local,
+                            PlaceContext::NonMutatingUse(NonMutatingUseContext::Move),
+                            source_location,
+                        );
+
+                        assert_eq!(
+                            local,
+                            RETURN_PLACE,
+                            "`MutVisitor` tried to mutate return place of `return` terminator"
+                        );
+                    }
+
                     TerminatorKind::SwitchInt {
                         discr,
                         switch_ty,
diff --git a/src/librustc_mir/transform/copy_prop.rs b/src/librustc_mir/transform/copy_prop.rs
index b9eb58f800e..ba406c72df8 100644
--- a/src/librustc_mir/transform/copy_prop.rs
+++ b/src/librustc_mir/transform/copy_prop.rs
@@ -73,7 +73,12 @@ fn run_pass(&self, tcx: TyCtxt<'tcx>, _source: MirSource<'tcx>, body: &mut Body<
                     }
                     // Conservatively gives up if the dest is an argument,
                     // because there may be uses of the original argument value.
-                    if body.local_kind(dest_local) == LocalKind::Arg {
+                    // Also gives up on the return place, as we cannot propagate into its implicit
+                    // use by `return`.
+                    if matches!(
+                        body.local_kind(dest_local),
+                        LocalKind::Arg | LocalKind::ReturnPointer
+                    ) {
                         debug!("  Can't copy-propagate local: dest {:?} (argument)", dest_local);
                         continue;
                     }
diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index d334006d7b5..bfc872be653 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -91,6 +91,16 @@ fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
             *local = self.to;
         }
     }
+
+    fn visit_terminator_kind(&mut self, kind: &mut TerminatorKind<'tcx>, location: Location) {
+        match kind {
+            TerminatorKind::Return => {
+                // Do not replace the implicit `_0` access here, as that's not possible. The
+                // transform already handles `return` correctly.
+            }
+            _ => self.super_terminator_kind(kind, location),
+        }
+    }
 }
 
 struct DerefArgVisitor<'tcx> {
diff --git a/src/librustc_mir/transform/inline.rs b/src/librustc_mir/transform/inline.rs
index a8e949ecb31..632408fde74 100644
--- a/src/librustc_mir/transform/inline.rs
+++ b/src/librustc_mir/transform/inline.rs
@@ -732,7 +732,11 @@ fn visit_retag(&mut self, kind: &mut RetagKind, place: &mut Place<'tcx>, loc: Lo
     }
 
     fn visit_terminator_kind(&mut self, kind: &mut TerminatorKind<'tcx>, loc: Location) {
-        self.super_terminator_kind(kind, loc);
+        // Don't try to modify the implicit `_0` access on return (`return` terminators are
+        // replaced down below anyways).
+        if !matches!(kind, TerminatorKind::Return) {
+            self.super_terminator_kind(kind, loc);
+        }
 
         match *kind {
             TerminatorKind::GeneratorDrop | TerminatorKind::Yield { .. } => bug!(),
-- 
2.44.0