Auto merge of #76683 - simonvandel:inst-combine-deref, r=oli-obk

Add optimization to avoid load of address

Look for the sequence
```rust
_2 = &_1;
...
_5 = (*_2)
```

in which we can replace the last statement with `_5 = _1` to avoid the load of _2

compiler/rustc_mir/src/transform/instcombine.rs

@@ -4,9 +4,14 @@ use crate::transform::{MirPass, MirSource};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_hir::Mutability;
 use rustc_index::vec::Idx;
-use rustc_middle::mir::visit::{MutVisitor, Visitor};
 use rustc_middle::mir::{
-    BinOp, Body, Constant, Local, Location, Operand, Place, PlaceRef, ProjectionElem, Rvalue,
+    visit::PlaceContext,
+    visit::{MutVisitor, Visitor},
+    Statement,
+};
+use rustc_middle::mir::{
+    BinOp, Body, BorrowKind, Constant, Local, Location, Operand, Place, PlaceRef, ProjectionElem,
+    Rvalue,
 };
 use rustc_middle::ty::{self, TyCtxt};
 use std::mem;
@@ -71,10 +76,36 @@ impl<'tcx> MutVisitor<'tcx> for InstCombineVisitor<'tcx> {
             *rvalue = Rvalue::Use(operand);
         }
 
+        if let Some(place) = self.optimizations.unneeded_deref.remove(&location) {
+            debug!("unneeded_deref: replacing {:?} with {:?}", rvalue, place);
+            *rvalue = Rvalue::Use(Operand::Copy(place));
+        }
+
         self.super_rvalue(rvalue, location)
     }
 }
 
+struct MutatingUseVisitor {
+    has_mutating_use: bool,
+    local_to_look_for: Local,
+}
+
+impl MutatingUseVisitor {
+    fn has_mutating_use_in_stmt(local: Local, stmt: &Statement<'tcx>, location: Location) -> bool {
+        let mut _self = Self { has_mutating_use: false, local_to_look_for: local };
+        _self.visit_statement(stmt, location);
+        _self.has_mutating_use
+    }
+}
+
+impl<'tcx> Visitor<'tcx> for MutatingUseVisitor {
+    fn visit_local(&mut self, local: &Local, context: PlaceContext, _: Location) {
+        if *local == self.local_to_look_for {
+            self.has_mutating_use |= context.is_mutating_use();
+        }
+    }
+}
+
 /// Finds optimization opportunities on the MIR.
 struct OptimizationFinder<'b, 'tcx> {
     body: &'b Body<'tcx>,
@@ -87,6 +118,85 @@ impl OptimizationFinder<'b, 'tcx> {
         OptimizationFinder { body, tcx, optimizations: OptimizationList::default() }
     }
 
+    fn find_deref_of_address(&mut self, rvalue: &Rvalue<'tcx>, location: Location) -> Option<()> {
+        // Look for the sequence
+        //
+        // _2 = &_1;
+        // ...
+        // _5 = (*_2);
+        //
+        // which we can replace the last statement with `_5 = _1;` to avoid the load of `_2`.
+        if let Rvalue::Use(op) = rvalue {
+            let local_being_derefed = match op.place()?.as_ref() {
+                PlaceRef { local, projection: [ProjectionElem::Deref] } => Some(local),
+                _ => None,
+            }?;
+
+            let stmt_index = location.statement_index;
+            // Look behind for statement that assigns the local from a address of operator.
+            // 6 is chosen as a heuristic determined by seeing the number of times
+            // the optimization kicked in compiling rust std.
+            let lower_index = stmt_index.saturating_sub(6);
+            let statements_to_look_in = self.body.basic_blocks()[location.block].statements
+                [lower_index..stmt_index]
+                .iter()
+                .rev();
+            for stmt in statements_to_look_in {
+                match &stmt.kind {
+                    // Exhaustive match on statements to detect conditions that warrant we bail out of the optimization.
+                    rustc_middle::mir::StatementKind::Assign(box (l, r))
+                        if l.local == local_being_derefed =>
+                    {
+                        match r {
+                            // Looking for immutable reference e.g _local_being_deref = &_1;
+                            Rvalue::Ref(
+                                _,
+                                // Only apply the optimization if it is an immutable borrow.
+                                BorrowKind::Shared,
+                                place_taken_address_of,
+                            ) => {
+                                self.optimizations
+                                    .unneeded_deref
+                                    .insert(location, *place_taken_address_of);
+                                return Some(());
+                            }
+
+                            // We found an assignment of `local_being_deref` that is not an immutable ref, e.g the following sequence
+                            // _2 = &_1;
+                            // _3 = &5
+                            // _2 = _3;  <-- this means it is no longer valid to replace the last statement with `_5 = _1;`
+                            // _5 = (*_2);
+                            _ => return None,
+                        }
+                    }
+
+                    // Inline asm can do anything, so bail out of the optimization.
+                    rustc_middle::mir::StatementKind::LlvmInlineAsm(_) => return None,
+
+                    // Check that `local_being_deref` is not being used in a mutating way which can cause misoptimization.
+                    rustc_middle::mir::StatementKind::Assign(box (_, _))
+                    | rustc_middle::mir::StatementKind::Coverage(_)
+                    | rustc_middle::mir::StatementKind::Nop
+                    | rustc_middle::mir::StatementKind::FakeRead(_, _)
+                    | rustc_middle::mir::StatementKind::StorageLive(_)
+                    | rustc_middle::mir::StatementKind::StorageDead(_)
+                    | rustc_middle::mir::StatementKind::Retag(_, _)
+                    | rustc_middle::mir::StatementKind::AscribeUserType(_, _)
+                    | rustc_middle::mir::StatementKind::SetDiscriminant { .. } => {
+                        if MutatingUseVisitor::has_mutating_use_in_stmt(
+                            local_being_derefed,
+                            stmt,
+                            location,
+                        ) {
+                            return None;
+                        }
+                    }
+                }
+            }
+        }
+        Some(())
+    }
+
     fn find_unneeded_equality_comparison(&mut self, rvalue: &Rvalue<'tcx>, location: Location) {
         // find Ne(_place, false) or Ne(false, _place)
         // or   Eq(_place, true) or Eq(true, _place)
@@ -153,6 +263,8 @@ impl Visitor<'tcx> for OptimizationFinder<'b, 'tcx> {
             }
         }
 
+        let _ = self.find_deref_of_address(rvalue, location);
+
         self.find_unneeded_equality_comparison(rvalue, location);
 
         self.super_rvalue(rvalue, location)
@@ -164,4 +276,5 @@ struct OptimizationList<'tcx> {
     and_stars: FxHashSet<Location>,
     arrays_lengths: FxHashMap<Location, Constant<'tcx>>,
     unneeded_equality_comparison: FxHashMap<Location, Operand<'tcx>>,
+    unneeded_deref: FxHashMap<Location, Place<'tcx>>,
 }

src/test/mir-opt/const_prop/ref_deref.main.ConstProp.diff

@@ -19,7 +19,7 @@
                                            // + span: $DIR/ref_deref.rs:5:6: 5:10
                                            // + literal: Const { ty: &i32, val: Unevaluated(WithOptConstParam { did: DefId(0:3 ~ ref_deref[317d]::main[0]), const_param_did: None }, [], Some(promoted[0])) }
           _2 = _4;                         // scope 0 at $DIR/ref_deref.rs:5:6: 5:10
--         _1 = (*_2);                      // scope 0 at $DIR/ref_deref.rs:5:5: 5:10
+-         _1 = (*_4);                      // scope 0 at $DIR/ref_deref.rs:5:5: 5:10
 +         _1 = const 4_i32;                // scope 0 at $DIR/ref_deref.rs:5:5: 5:10
           StorageDead(_2);                 // scope 0 at $DIR/ref_deref.rs:5:10: 5:11
           StorageDead(_1);                 // scope 0 at $DIR/ref_deref.rs:5:10: 5:11

src/test/mir-opt/const_prop/ref_deref_project.main.ConstProp.diff

@@ -19,7 +19,7 @@
                                            // + span: $DIR/ref_deref_project.rs:5:6: 5:17
                                            // + literal: Const { ty: &(i32, i32), val: Unevaluated(WithOptConstParam { did: DefId(0:3 ~ ref_deref_project[317d]::main[0]), const_param_did: None }, [], Some(promoted[0])) }
           _2 = &((*_4).1: i32);            // scope 0 at $DIR/ref_deref_project.rs:5:6: 5:17
-          _1 = (*_2);                      // scope 0 at $DIR/ref_deref_project.rs:5:5: 5:17
+          _1 = ((*_4).1: i32);             // scope 0 at $DIR/ref_deref_project.rs:5:5: 5:17
           StorageDead(_2);                 // scope 0 at $DIR/ref_deref_project.rs:5:17: 5:18
           StorageDead(_1);                 // scope 0 at $DIR/ref_deref_project.rs:5:17: 5:18
           _0 = const ();                   // scope 0 at $DIR/ref_deref_project.rs:4:11: 6:2

src/test/mir-opt/inline/inline_closure_borrows_arg.foo.Inline.after.mir

@@ -9,13 +9,14 @@ fn foo(_1: T, _2: &i32) -> i32 {
     let mut _5: (&i32, &i32);            // in scope 0 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
     let mut _6: &i32;                    // in scope 0 at $DIR/inline-closure-borrows-arg.rs:16:7: 16:8
     let mut _7: &i32;                    // in scope 0 at $DIR/inline-closure-borrows-arg.rs:16:10: 16:11
-    let mut _8: &i32;                    // in scope 0 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
     let mut _9: &i32;                    // in scope 0 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
+    let mut _10: &i32;                   // in scope 0 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
     scope 1 {
         debug x => _3;                   // in scope 1 at $DIR/inline-closure-borrows-arg.rs:12:9: 12:10
         scope 2 {
-            debug r => _8;               // in scope 2 at $DIR/inline-closure-borrows-arg.rs:12:14: 12:15
-            debug _s => _9;              // in scope 2 at $DIR/inline-closure-borrows-arg.rs:12:23: 12:25
+            debug r => _9;               // in scope 2 at $DIR/inline-closure-borrows-arg.rs:12:14: 12:15
+            debug _s => _10;             // in scope 2 at $DIR/inline-closure-borrows-arg.rs:12:23: 12:25
+            let _8: &i32;                // in scope 2 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
         }
     }
     scope 3 {
@@ -33,13 +34,16 @@ fn foo(_1: T, _2: &i32) -> i32 {
         _7 = &(*_2);                     // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:10: 16:11
         (_5.0: &i32) = move _6;          // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
         (_5.1: &i32) = move _7;          // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
-        StorageLive(_8);                 // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
-        _8 = move (_5.0: &i32);          // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
         StorageLive(_9);                 // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
-        _9 = move (_5.1: &i32);          // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
-        _0 = (*_8);                      // scope 3 at $DIR/inline-closure-borrows-arg.rs:14:9: 14:18
+        _9 = move (_5.0: &i32);          // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
+        StorageLive(_10);                // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
+        _10 = move (_5.1: &i32);         // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
+        StorageLive(_8);                 // scope 2 at $DIR/inline-closure-borrows-arg.rs:13:13: 13:21
+        _8 = _9;                         // scope 2 at $DIR/inline-closure-borrows-arg.rs:13:24: 13:27
+        _0 = (*_9);                      // scope 3 at $DIR/inline-closure-borrows-arg.rs:14:9: 14:18
+        StorageDead(_8);                 // scope 2 at $DIR/inline-closure-borrows-arg.rs:15:5: 15:6
+        StorageDead(_10);                // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
         StorageDead(_9);                 // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
-        StorageDead(_8);                 // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:5: 16:12
         StorageDead(_7);                 // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:11: 16:12
         StorageDead(_6);                 // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:11: 16:12
         StorageDead(_5);                 // scope 1 at $DIR/inline-closure-borrows-arg.rs:16:11: 16:12

src/test/mir-opt/inst_combine_deref.deep_opt.InstCombine.diff

@@ -0,0 +1,92 @@
+- // MIR for `deep_opt` before InstCombine
++ // MIR for `deep_opt` after InstCombine
+
+  fn deep_opt() -> (u64, u64, u64) {
+      let mut _0: (u64, u64, u64);         // return place in scope 0 at $DIR/inst_combine_deref.rs:11:18: 11:33
+      let _1: u64;                         // in scope 0 at $DIR/inst_combine_deref.rs:12:9: 12:11
+      let mut _10: u64;                    // in scope 0 at $DIR/inst_combine_deref.rs:21:6: 21:8
+      let mut _11: u64;                    // in scope 0 at $DIR/inst_combine_deref.rs:21:10: 21:12
+      let mut _12: u64;                    // in scope 0 at $DIR/inst_combine_deref.rs:21:14: 21:16
+      scope 1 {
+          debug x1 => _1;                  // in scope 1 at $DIR/inst_combine_deref.rs:12:9: 12:11
+          let _2: u64;                     // in scope 1 at $DIR/inst_combine_deref.rs:13:9: 13:11
+          scope 2 {
+              debug x2 => _2;              // in scope 2 at $DIR/inst_combine_deref.rs:13:9: 13:11
+              let _3: u64;                 // in scope 2 at $DIR/inst_combine_deref.rs:14:9: 14:11
+              scope 3 {
+                  debug x3 => _3;          // in scope 3 at $DIR/inst_combine_deref.rs:14:9: 14:11
+                  let _4: &u64;            // in scope 3 at $DIR/inst_combine_deref.rs:15:9: 15:11
+                  scope 4 {
+                      debug y1 => _4;      // in scope 4 at $DIR/inst_combine_deref.rs:15:9: 15:11
+                      let _5: &u64;        // in scope 4 at $DIR/inst_combine_deref.rs:16:9: 16:11
+                      scope 5 {
+                          debug y2 => _5;  // in scope 5 at $DIR/inst_combine_deref.rs:16:9: 16:11
+                          let _6: &u64;    // in scope 5 at $DIR/inst_combine_deref.rs:17:9: 17:11
+                          scope 6 {
+                              debug y3 => _6; // in scope 6 at $DIR/inst_combine_deref.rs:17:9: 17:11
+                              let _7: u64; // in scope 6 at $DIR/inst_combine_deref.rs:18:9: 18:11
+                              scope 7 {
+                                  debug z1 => _7; // in scope 7 at $DIR/inst_combine_deref.rs:18:9: 18:11
+                                  let _8: u64; // in scope 7 at $DIR/inst_combine_deref.rs:19:9: 19:11
+                                  scope 8 {
+                                      debug z2 => _8; // in scope 8 at $DIR/inst_combine_deref.rs:19:9: 19:11
+                                      let _9: u64; // in scope 8 at $DIR/inst_combine_deref.rs:20:9: 20:11
+                                      scope 9 {
+                                          debug z3 => _9; // in scope 9 at $DIR/inst_combine_deref.rs:20:9: 20:11
+                                      }
+                                  }
+                              }
+                          }
+                      }
+                  }
+              }
+          }
+      }
+
+      bb0: {
+          StorageLive(_1);                 // scope 0 at $DIR/inst_combine_deref.rs:12:9: 12:11
+          _1 = const 1_u64;                // scope 0 at $DIR/inst_combine_deref.rs:12:14: 12:15
+          StorageLive(_2);                 // scope 1 at $DIR/inst_combine_deref.rs:13:9: 13:11
+          _2 = const 2_u64;                // scope 1 at $DIR/inst_combine_deref.rs:13:14: 13:15
+          StorageLive(_3);                 // scope 2 at $DIR/inst_combine_deref.rs:14:9: 14:11
+          _3 = const 3_u64;                // scope 2 at $DIR/inst_combine_deref.rs:14:14: 14:15
+          StorageLive(_4);                 // scope 3 at $DIR/inst_combine_deref.rs:15:9: 15:11
+          _4 = &_1;                        // scope 3 at $DIR/inst_combine_deref.rs:15:14: 15:17
+          StorageLive(_5);                 // scope 4 at $DIR/inst_combine_deref.rs:16:9: 16:11
+          _5 = &_2;                        // scope 4 at $DIR/inst_combine_deref.rs:16:14: 16:17
+          StorageLive(_6);                 // scope 5 at $DIR/inst_combine_deref.rs:17:9: 17:11
+          _6 = &_3;                        // scope 5 at $DIR/inst_combine_deref.rs:17:14: 17:17
+          StorageLive(_7);                 // scope 6 at $DIR/inst_combine_deref.rs:18:9: 18:11
+-         _7 = (*_4);                      // scope 6 at $DIR/inst_combine_deref.rs:18:14: 18:17
++         _7 = _1;                         // scope 6 at $DIR/inst_combine_deref.rs:18:14: 18:17
+          StorageLive(_8);                 // scope 7 at $DIR/inst_combine_deref.rs:19:9: 19:11
+-         _8 = (*_5);                      // scope 7 at $DIR/inst_combine_deref.rs:19:14: 19:17
++         _8 = _2;                         // scope 7 at $DIR/inst_combine_deref.rs:19:14: 19:17
+          StorageLive(_9);                 // scope 8 at $DIR/inst_combine_deref.rs:20:9: 20:11
+-         _9 = (*_6);                      // scope 8 at $DIR/inst_combine_deref.rs:20:14: 20:17
++         _9 = _3;                         // scope 8 at $DIR/inst_combine_deref.rs:20:14: 20:17
+          StorageLive(_10);                // scope 9 at $DIR/inst_combine_deref.rs:21:6: 21:8
+          _10 = _7;                        // scope 9 at $DIR/inst_combine_deref.rs:21:6: 21:8
+          StorageLive(_11);                // scope 9 at $DIR/inst_combine_deref.rs:21:10: 21:12
+          _11 = _8;                        // scope 9 at $DIR/inst_combine_deref.rs:21:10: 21:12
+          StorageLive(_12);                // scope 9 at $DIR/inst_combine_deref.rs:21:14: 21:16
+          _12 = _9;                        // scope 9 at $DIR/inst_combine_deref.rs:21:14: 21:16
+          (_0.0: u64) = move _10;          // scope 9 at $DIR/inst_combine_deref.rs:21:5: 21:17
+          (_0.1: u64) = move _11;          // scope 9 at $DIR/inst_combine_deref.rs:21:5: 21:17
+          (_0.2: u64) = move _12;          // scope 9 at $DIR/inst_combine_deref.rs:21:5: 21:17
+          StorageDead(_12);                // scope 9 at $DIR/inst_combine_deref.rs:21:16: 21:17
+          StorageDead(_11);                // scope 9 at $DIR/inst_combine_deref.rs:21:16: 21:17
+          StorageDead(_10);                // scope 9 at $DIR/inst_combine_deref.rs:21:16: 21:17
+          StorageDead(_9);                 // scope 8 at $DIR/inst_combine_deref.rs:22:1: 22:2
+          StorageDead(_8);                 // scope 7 at $DIR/inst_combine_deref.rs:22:1: 22:2
+          StorageDead(_7);                 // scope 6 at $DIR/inst_combine_deref.rs:22:1: 22:2
+          StorageDead(_6);                 // scope 5 at $DIR/inst_combine_deref.rs:22:1: 22:2
+          StorageDead(_5);                 // scope 4 at $DIR/inst_combine_deref.rs:22:1: 22:2
+          StorageDead(_4);                 // scope 3 at $DIR/inst_combine_deref.rs:22:1: 22:2
+          StorageDead(_3);                 // scope 2 at $DIR/inst_combine_deref.rs:22:1: 22:2
+          StorageDead(_2);                 // scope 1 at $DIR/inst_combine_deref.rs:22:1: 22:2
+          StorageDead(_1);                 // scope 0 at $DIR/inst_combine_deref.rs:22:1: 22:2
+          return;                          // scope 0 at $DIR/inst_combine_deref.rs:22:2: 22:2
+      }
+  }
+