Do not merge locals that have their address taken.

Author: cjgillot

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -3,6 +3,7 @@ use rustc_index::vec::IndexVec;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
+use rustc_mir_dataflow::impls::borrowed_locals;
 
 use crate::ssa::SsaLocals;
 use crate::MirPass;
@@ -33,7 +34,8 @@ impl<'tcx> MirPass<'tcx> for CopyProp {
 
 fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
-    let ssa = SsaLocals::new(tcx, param_env, body);
+    let borrowed_locals = borrowed_locals(body);
+    let ssa = SsaLocals::new(tcx, param_env, body, &borrowed_locals);
 
     let fully_moved = fully_moved_locals(&ssa, body);
     debug!(?fully_moved);
@@ -42,14 +44,19 @@ fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     for (local, &head) in ssa.copy_classes().iter_enumerated() {
         if local != head {
             storage_to_remove.insert(head);
-            storage_to_remove.insert(local);
         }
     }
 
     let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);
 
-    Replacer { tcx, copy_classes: &ssa.copy_classes(), fully_moved, storage_to_remove }
-        .visit_body_preserves_cfg(body);
+    Replacer {
+        tcx,
+        copy_classes: &ssa.copy_classes(),
+        fully_moved,
+        borrowed_locals,
+        storage_to_remove,
+    }
+    .visit_body_preserves_cfg(body);
 
     if any_replacement {
         crate::simplify::remove_unused_definitions(body);
@@ -94,6 +101,7 @@ struct Replacer<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     fully_moved: BitSet<Local>,
     storage_to_remove: BitSet<Local>,
+    borrowed_locals: BitSet<Local>,
     copy_classes: &'a IndexVec<Local, Local>,
 }
 
@@ -102,8 +110,45 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
         self.tcx
     }
 
-    fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
-        *local = self.copy_classes[*local];
+    fn visit_local(&mut self, local: &mut Local, ctxt: PlaceContext, _: Location) {
+        let new_local = self.copy_classes[*local];
+        match ctxt {
+            // Do not modify the local in storage statements.
+            PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead) => {}
+            // The local should have been marked as non-SSA.
+            PlaceContext::MutatingUse(_) => assert_eq!(*local, new_local),
+            // We access the value.
+            _ => *local = new_local,
+        }
+    }
+
+    fn visit_place(&mut self, place: &mut Place<'tcx>, ctxt: PlaceContext, loc: Location) {
+        if let Some(new_projection) = self.process_projection(&place.projection, loc) {
+            place.projection = self.tcx().intern_place_elems(&new_projection);
+        }
+
+        let observes_address = match ctxt {
+            PlaceContext::NonMutatingUse(
+                NonMutatingUseContext::SharedBorrow
+                | NonMutatingUseContext::ShallowBorrow
+                | NonMutatingUseContext::UniqueBorrow
+                | NonMutatingUseContext::AddressOf,
+            ) => true,
+            // For debuginfo, merging locals is ok.
+            PlaceContext::NonUse(NonUseContext::VarDebugInfo) => {
+                self.borrowed_locals.contains(place.local)
+            }
+            _ => false,
+        };
+        if observes_address && !place.is_indirect() {
+            // We observe the address of `place.local`. Do not replace it.
+        } else {
+            self.visit_local(
+                &mut place.local,
+                PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy),
+                loc,
+            )
+        }
     }
 
     fn visit_operand(&mut self, operand: &mut Operand<'tcx>, loc: Location) {
@@ -117,17 +162,17 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
     }
 
     fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
-        if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = stmt.kind
+        if let StatementKind::StorageDead(l) = stmt.kind
             && self.storage_to_remove.contains(l)
         {
             stmt.make_nop();
-        }
-        if let StatementKind::Assign(box (ref place, _)) = stmt.kind
-            && let Some(l) = place.as_local()
-            && self.copy_classes[l] != l
+        } else if let StatementKind::Assign(box (ref place, ref mut rvalue)) = stmt.kind
+            && place.as_local().is_some()
         {
-            stmt.make_nop();
+            // Do not replace assignments.
+            self.visit_rvalue(rvalue, loc)
+        } else {
+            self.super_statement(stmt, loc);
         }
-        self.super_statement(stmt, loc);
     }
 }

compiler/rustc_mir_transform/src/ssa.rs

@@ -6,7 +6,6 @@ use rustc_middle::middle::resolve_lifetime::Set1;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{ParamEnv, TyCtxt};
-use rustc_mir_dataflow::impls::borrowed_locals;
 
 #[derive(Debug)]
 pub struct SsaLocals {
@@ -21,19 +20,23 @@ pub struct SsaLocals {
 }
 
 impl SsaLocals {
-    pub fn new<'tcx>(tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>, body: &Body<'tcx>) -> SsaLocals {
+    pub fn new<'tcx>(
+        tcx: TyCtxt<'tcx>,
+        param_env: ParamEnv<'tcx>,
+        body: &Body<'tcx>,
+        borrowed_locals: &BitSet<Local>,
+    ) -> SsaLocals {
         let assignment_order = Vec::new();
 
         let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls);
         let dominators = body.basic_blocks.dominators();
         let mut visitor = SsaVisitor { assignments, assignment_order, dominators };
 
-        let borrowed = borrowed_locals(body);
         for (local, decl) in body.local_decls.iter_enumerated() {
             if matches!(body.local_kind(local), LocalKind::Arg) {
                 visitor.assignments[local] = Set1::One(LocationExtended::Arg);
             }
-            if borrowed.contains(local) && !decl.ty.is_freeze(tcx, param_env) {
+            if borrowed_locals.contains(local) && !decl.ty.is_freeze(tcx, param_env) {
                 visitor.assignments[local] = Set1::Many;
             }
         }

tests/mir-opt/const_debuginfo.main.ConstDebugInfo.diff

@@ -56,8 +56,11 @@
       }
 
       bb0: {
+          StorageLive(_1);                 // scope 0 at $DIR/const_debuginfo.rs:+1:9: +1:10
           _1 = const 1_u8;                 // scope 0 at $DIR/const_debuginfo.rs:+1:13: +1:16
+          StorageLive(_2);                 // scope 1 at $DIR/const_debuginfo.rs:+2:9: +2:10
           _2 = const 2_u8;                 // scope 1 at $DIR/const_debuginfo.rs:+2:13: +2:16
+          StorageLive(_3);                 // scope 2 at $DIR/const_debuginfo.rs:+3:9: +3:10
           _3 = const 3_u8;                 // scope 2 at $DIR/const_debuginfo.rs:+3:13: +3:16
           StorageLive(_4);                 // scope 3 at $DIR/const_debuginfo.rs:+4:9: +4:12
           StorageLive(_5);                 // scope 3 at $DIR/const_debuginfo.rs:+4:15: +4:20

tests/mir-opt/const_prop/bad_op_mod_by_zero.main.ConstProp.diff

@@ -18,6 +18,7 @@
       }
 
       bb0: {
+          StorageLive(_1);                 // scope 0 at $DIR/bad_op_mod_by_zero.rs:+1:9: +1:10
           _1 = const 0_i32;                // scope 0 at $DIR/bad_op_mod_by_zero.rs:+1:13: +1:14
           StorageLive(_2);                 // scope 1 at $DIR/bad_op_mod_by_zero.rs:+2:9: +2:11
 -         _4 = Eq(_1, const 0_i32);        // scope 1 at $DIR/bad_op_mod_by_zero.rs:+2:14: +2:19

tests/mir-opt/const_prop/scalar_literal_propagation.main.ConstProp.diff

@@ -11,6 +11,7 @@
       }
 
       bb0: {
+          StorageLive(_1);                 // scope 0 at $DIR/scalar_literal_propagation.rs:+1:9: +1:10
           _1 = const 1_u32;                // scope 0 at $DIR/scalar_literal_propagation.rs:+1:13: +1:14
           StorageLive(_2);                 // scope 1 at $DIR/scalar_literal_propagation.rs:+2:5: +2:15
 -         _2 = consume(_1) -> bb1;         // scope 1 at $DIR/scalar_literal_propagation.rs:+2:5: +2:15

tests/mir-opt/copy-prop/borrowed_local.f.CopyProp.diff

@@ -0,0 +1,34 @@
+- // MIR for `f` before CopyProp
++ // MIR for `f` after CopyProp
+  
+  fn f() -> bool {
+      let mut _0: bool;                    // return place in scope 0 at $DIR/borrowed_local.rs:+0:11: +0:15
+      let mut _1: u8;                      // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+      let mut _2: &u8;                     // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+      let mut _3: u8;                      // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+      let mut _4: &u8;                     // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+  
+      bb0: {
+          _1 = const 5_u8;                 // scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+          _2 = &_1;                        // scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+          _3 = _1;                         // scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+          _4 = &_3;                        // scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+          _0 = cmp_ref(_2, _4) -> bb1;     // scope 0 at $DIR/borrowed_local.rs:+8:13: +8:45
+                                           // mir::Constant
+                                           // + span: $DIR/borrowed_local.rs:23:29: 23:36
+                                           // + literal: Const { ty: for<'a, 'b> fn(&'a u8, &'b u8) -> bool {cmp_ref}, val: Value(<ZST>) }
+      }
+  
+      bb1: {
+-         _0 = opaque::<u8>(_3) -> bb2;    // scope 0 at $DIR/borrowed_local.rs:+12:13: +12:38
++         _0 = opaque::<u8>(_1) -> bb2;    // scope 0 at $DIR/borrowed_local.rs:+12:13: +12:38
+                                           // mir::Constant
+                                           // + span: $DIR/borrowed_local.rs:27:28: 27:34
+                                           // + literal: Const { ty: fn(u8) -> bool {opaque::<u8>}, val: Value(<ZST>) }
+      }
+  
+      bb2: {
+          return;                          // scope 0 at $DIR/borrowed_local.rs:+15:13: +15:21
+      }
+  }
+  

tests/mir-opt/copy-prop/borrowed_local.rs

@@ -0,0 +1,39 @@
+// unit-test: CopyProp
+
+#![feature(custom_mir, core_intrinsics)]
+#![allow(unused_assignments)]
+extern crate core;
+use core::intrinsics::mir::*;
+
+fn opaque(_: impl Sized) -> bool { true }
+
+fn cmp_ref(a: &u8, b: &u8) -> bool {
+    std::ptr::eq(a as *const u8, b as *const u8)
+}
+
+#[custom_mir(dialect = "analysis", phase = "post-cleanup")]
+fn f() -> bool {
+    mir!(
+        {
+            let a = 5_u8;
+            let r1 = &a;
+            let b = a;
+            // We cannot propagate the place `a`.
+            let r2 = &b;
+            Call(RET, next, cmp_ref(r1, r2))
+        }
+        next = {
+            // But we can propagate the value `a`.
+            Call(RET, ret, opaque(b))
+        }
+        ret = {
+            Return()
+        }
+    )
+}
+
+fn main() {
+    assert!(!f());
+}
+
+// EMIT_MIR borrowed_local.f.CopyProp.diff

tests/mir-opt/copy-prop/cycle.main.CopyProp.diff

@@ -29,7 +29,7 @@
       }
 
       bb1: {
--         StorageLive(_2);                 // scope 1 at $DIR/cycle.rs:+2:9: +2:10
+          StorageLive(_2);                 // scope 1 at $DIR/cycle.rs:+2:9: +2:10
           _2 = _1;                         // scope 1 at $DIR/cycle.rs:+2:13: +2:14
 -         StorageLive(_3);                 // scope 2 at $DIR/cycle.rs:+3:9: +3:10
 -         _3 = _2;                         // scope 2 at $DIR/cycle.rs:+3:13: +3:14

tests/mir-opt/copy-prop/dead_stores_79191.f.CopyProp.after.mir

@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
     }
 
     bb0: {
+        StorageLive(_2);                 // scope 0 at $DIR/dead_stores_79191.rs:+1:9: +1:10
         _2 = _1;                         // scope 0 at $DIR/dead_stores_79191.rs:+1:13: +1:14
         _1 = const 5_usize;              // scope 1 at $DIR/dead_stores_79191.rs:+2:5: +2:10
         _1 = _2;                         // scope 1 at $DIR/dead_stores_79191.rs:+3:5: +3:10

tests/mir-opt/copy-prop/dead_stores_better.f.CopyProp.after.mir

@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
     }
 
     bb0: {
+        StorageLive(_2);                 // scope 0 at $DIR/dead_stores_better.rs:+1:9: +1:10
         _2 = _1;                         // scope 0 at $DIR/dead_stores_better.rs:+1:13: +1:14
         _1 = const 5_usize;              // scope 1 at $DIR/dead_stores_better.rs:+2:5: +2:10
         _1 = _2;                         // scope 1 at $DIR/dead_stores_better.rs:+3:5: +3:10