Auto merge of #45753 - sinkuu:mir_copyprop_arg, r=arielb1

Fix MIR CopyPropagation errneously propagating assignments to function args

Compiling this code with MIR CopyPropagation activated will result in printing `5`,
because CopyProp errneously propagates the assignment of `5` to all `x`:

```rust
fn bar(mut x: u8) {
    println!("{}", x);
    x = 5;
}

fn main() {
    bar(123);
}

```

If a local is propagated, it will result in an ICE at trans due to an use-before-def:

```rust
fn dummy(x: u8) -> u8 { x }

fn foo(mut x: u8) {
    x = dummy(x); // this will assign a local to `x`
}
```
Currently CopyProp conservatively gives up if there are multiple assignments to a local,
but it is not took into account that arguments are already assigned from the beginning.
This PR fixes the problem by preventing propagation of assignments to function arguments.

src/librustc_mir/transform/copy_prop.rs

@@ -69,10 +69,14 @@ fn run_pass<'a, 'tcx>(&self,
             return;
         }
 
+        let mut def_use_analysis = DefUseAnalysis::new(mir);
         loop {
-            let mut def_use_analysis = DefUseAnalysis::new(mir);
             def_use_analysis.analyze(mir);
 
+            if eliminate_self_assignments(mir, &def_use_analysis) {
+                def_use_analysis.analyze(mir);
+            }
+
             let mut changed = false;
             for dest_local in mir.local_decls.indices() {
                 debug!("Considering destination local: {:?}", dest_local);
@@ -99,9 +103,14 @@ fn run_pass<'a, 'tcx>(&self,
                                dest_local);
                         continue
                     }
-                    let dest_lvalue_def = dest_use_info.defs_and_uses.iter().filter(|lvalue_def| {
-                        lvalue_def.context.is_mutating_use() && !lvalue_def.context.is_drop()
-                    }).next().unwrap();
+                    // Conservatively gives up if the dest is an argument,
+                    // because there may be uses of the original argument value.
+                    if mir.local_kind(dest_local) == LocalKind::Arg {
+                        debug!("  Can't copy-propagate local: dest {:?} (argument)",
+                            dest_local);
+                        continue;
+                    }
+                    let dest_lvalue_def = dest_use_info.defs_not_including_drop().next().unwrap();
                     location = dest_lvalue_def.location;
 
                     let basic_block = &mir[location.block];
@@ -151,6 +160,39 @@ fn run_pass<'a, 'tcx>(&self,
     }
 }
 
+fn eliminate_self_assignments<'tcx>(
+    mir: &mut Mir<'tcx>,
+    def_use_analysis: &DefUseAnalysis<'tcx>,
+) -> bool {
+    let mut changed = false;
+
+    for dest_local in mir.local_decls.indices() {
+        let dest_use_info = def_use_analysis.local_info(dest_local);
+
+        for def in dest_use_info.defs_not_including_drop() {
+            let location = def.location;
+            if let Some(stmt) = mir[location.block].statements.get(location.statement_index) {
+                match stmt.kind {
+                    StatementKind::Assign(
+                        Lvalue::Local(local),
+                        Rvalue::Use(Operand::Consume(Lvalue::Local(src_local))),
+                    ) if local == dest_local && dest_local == src_local => {}
+                    _ => {
+                        continue;
+                    }
+                }
+            } else {
+                continue;
+            }
+            debug!("Deleting a self-assignment for {:?}", dest_local);
+            mir.make_statement_nop(location);
+            changed = true;
+        }
+    }
+
+    changed
+}
+
 enum Action<'tcx> {
     PropagateLocalCopy(Local),
     PropagateConstant(Constant<'tcx>),

src/librustc_mir/util/def_use.rs

@@ -15,6 +15,8 @@
 use rustc_data_structures::indexed_vec::IndexVec;
 use std::marker::PhantomData;
 use std::mem;
+use std::slice;
+use std::iter;
 
 pub struct DefUseAnalysis<'tcx> {
     info: IndexVec<Local, Info<'tcx>>,
@@ -39,6 +41,8 @@ pub fn new(mir: &Mir<'tcx>) -> DefUseAnalysis<'tcx> {
     }
 
     pub fn analyze(&mut self, mir: &Mir<'tcx>) {
+        self.clear();
+
         let mut finder = DefUseFinder {
             info: mem::replace(&mut self.info, IndexVec::new()),
         };
@@ -46,6 +50,12 @@ pub fn analyze(&mut self, mir: &Mir<'tcx>) {
         self.info = finder.info
     }
 
+    fn clear(&mut self) {
+        for info in &mut self.info {
+            info.clear();
+        }
+    }
+
     pub fn local_info(&self, local: Local) -> &Info<'tcx> {
         &self.info[local]
     }
@@ -93,14 +103,24 @@ fn new() -> Info<'tcx> {
         }
     }
 
+    fn clear(&mut self) {
+        self.defs_and_uses.clear();
+    }
+
     pub fn def_count(&self) -> usize {
         self.defs_and_uses.iter().filter(|lvalue_use| lvalue_use.context.is_mutating_use()).count()
     }
 
     pub fn def_count_not_including_drop(&self) -> usize {
+        self.defs_not_including_drop().count()
+    }
+
+    pub fn defs_not_including_drop(
+        &self,
+    ) -> iter::Filter<slice::Iter<Use<'tcx>>, fn(&&Use<'tcx>) -> bool> {
         self.defs_and_uses.iter().filter(|lvalue_use| {
             lvalue_use.context.is_mutating_use() && !lvalue_use.context.is_drop()
-        }).count()
+        })
     }
 
     pub fn use_count(&self) -> usize {

src/test/mir-opt/copy_propagation_arg.rs

@@ -0,0 +1,116 @@
+// Copyright 2017 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.
+
+// Check that CopyPropagation does not propagate an assignment to a function argument
+// (doing so can break usages of the original argument value)
+
+fn dummy(x: u8) -> u8 {
+    x
+}
+
+fn foo(mut x: u8) {
+    // calling `dummy` to make an use of `x` that copyprop cannot eliminate
+    x = dummy(x); // this will assign a local to `x`
+}
+
+fn bar(mut x: u8) {
+    dummy(x);
+    x = 5;
+}
+
+fn baz(mut x: i32) {
+    // self-assignment to a function argument should be eliminated
+    x = x;
+}
+
+fn main() {
+    // Make sure the function actually gets instantiated.
+    foo(0);
+    bar(0);
+    baz(0);
+}
+
+// END RUST SOURCE
+// START rustc.foo.CopyPropagation.before.mir
+// bb0: {
+//     StorageLive(_2);
+//     StorageLive(_3);
+//     _3 = _1;
+//     _2 = const dummy(_3) -> bb1;
+// }
+// bb1: {
+//     StorageDead(_3);
+//     _1 = _2;
+//     StorageDead(_2);
+//     _0 = ();
+//     return;
+// }
+// END rustc.foo.CopyPropagation.before.mir
+// START rustc.foo.CopyPropagation.after.mir
+// bb0: {
+//     StorageLive(_2);
+//     nop;
+//     nop;
+//     _2 = const dummy(_1) -> bb1;
+// }
+// bb1: {
+//     nop;
+//     _1 = _2;
+//     StorageDead(_2);
+//     _0 = ();
+//     return;
+// }
+// END rustc.foo.CopyPropagation.after.mir
+// START rustc.bar.CopyPropagation.before.mir
+// bb0: {
+//     StorageLive(_3);
+//     _3 = _1;
+//     _2 = const dummy(_3) -> bb1;
+// }
+// bb1: {
+//     StorageDead(_3);
+//     _1 = const 5u8;
+//     _0 = ();
+//     return;
+// }
+// END rustc.bar.CopyPropagation.before.mir
+// START rustc.bar.CopyPropagation.after.mir
+// bb0: {
+//     nop;
+//     nop;
+//     _2 = const dummy(_1) -> bb1;
+// }
+// bb1: {
+//     nop;
+//     _1 = const 5u8;
+//     _0 = ();
+//     return;
+// }
+// END rustc.bar.CopyPropagation.after.mir
+// START rustc.baz.CopyPropagation.before.mir
+// bb0: {
+//     StorageLive(_2);
+//     _2 = _1;
+//     _1 = _2;
+//     StorageDead(_2);
+//     _0 = ();
+//     return;
+// }
+// END rustc.baz.CopyPropagation.before.mir
+// START rustc.baz.CopyPropagation.after.mir
+// bb0: {
+//     nop;
+//     nop;
+//     nop;
+//     nop;
+//     _0 = ();
+//     return;
+// }
+// END rustc.baz.CopyPropagation.after.mir