Auto merge of #106285 - cjgillot:refprop-ssa, r=JakobDegen

Implement SSA-based reference propagation

Rust has a tendency to create a lot of short-lived borrows, in particular for method calls. This PR aims to remove those short-lived borrows with a const-propagation dedicated to pointers to local places.

This pass aims to transform the following pattern:
```
  _1 = &raw? mut? PLACE;
  _3 = *_1;
  _4 = &raw? mut? *_1;
```

Into
```
  _1 = &raw? mut? PLACE;
  _3 = PLACE;
  _4 = &raw? mut? PLACE;
```

where `PLACE` is a direct or an indirect place expression.

By removing indirection, this pass should help both dest-prop and const-prop to handle more cases.
This optimization is distinct from const-prop and dataflow const-prop since the borrow-reborrow patterns needs to preserve borrowck invariants, especially the uniqueness property of mutable references.

The pointed-to places are computed using a SSA analysis. We suppose that removable borrows are typically temporaries from autoref, so they are by construction assigned only once, and a SSA analysis is enough to catch them. For each local, we store both where and how it is used, in order to efficiently compute the all-or-nothing property. Thanks to `Derefer`, we only have to track locals, not places in general.

---

There are 3 properties that need to be upheld for this transformation to be legal:
- place constness: `PLACE` must refer to the same memory wherever it appears;
- pointer liveness: we must not introduce dereferences of dangling pointers;
- `&mut` borrow uniqueness.

## Constness

If `PLACE` is an indirect projection, if its of the form `(*LOCAL).PROJECTIONS` where:
- `LOCAL` is SSA;
- all projections in `PROJECTIONS` are constant (no dereference and no indexing).

If `PLACE` is a direct projection of a local, we consider it as constant if:
- the local is always live, or it has a single `StorageLive` that dominates all uses;
- all projections are constant.

# Liveness

When performing a substitution, we must take care not to introduce uses of dangling locals.

Using a dangling borrow is UB. Therefore, we assume that for any use of `*x`, where `x` is a borrow, the pointed-to memory is live.

Limitations:
- occurrences of `*x` in an `&raw mut? *x` are accepted;
- raw pointers are allowed to be dangling.

In those 2 case, we do not substitute anything, to be on the safe side.

**Open question:** we do not differentiate borrows of ZST and non-ZST. The UB rules may be
different depending on the layout. Having a different treatment would effectively prevent this
pass from running on polymorphic MIR, which defeats the purpose of MIR opts.

## Uniqueness

For `&mut` borrows, we also need to preserve the uniqueness property:
we must avoid creating a state where we interleave uses of `*_1` and `_2`.
To do it, we only perform full substitution of mutable borrows:
we replace either all or none of the occurrences of `*_1`.

Some care has to be taken when `_1` is copied in other locals.
```
   _1 = &raw? mut? _2;
   _3 = *_1;
   _4 = _1
   _5 = *_4
```
In such cases, fully substituting `_1` means fully substituting all of the copies.

For immutable borrows, we do not need to preserve such uniqueness property,
so we perform all the possible substitutions without removing the `_1 = &_2` statement.

compiler/rustc_middle/src/mir/mod.rs

@@ -1524,6 +1524,19 @@ impl<V, T> ProjectionElem<V, T> {
         }
     }
 
+    /// Returns `true` if the target of this projection always refers to the same memory region
+    /// whatever the state of the program.
+    pub fn is_stable_offset(&self) -> bool {
+        match self {
+            Self::Deref | Self::Index(_) => false,
+            Self::Field(_, _)
+            | Self::OpaqueCast(_)
+            | Self::ConstantIndex { .. }
+            | Self::Subslice { .. }
+            | Self::Downcast(_, _) => true,
+        }
+    }
+
     /// Returns `true` if this is a `Downcast` projection with the given `VariantIdx`.
     pub fn is_downcast_to(&self, v: VariantIdx) -> bool {
         matches!(*self, Self::Downcast(_, x) if x == v)

compiler/rustc_mir_dataflow/src/impls/mod.rs

@@ -26,7 +26,7 @@ pub use self::borrowed_locals::borrowed_locals;
 pub use self::borrowed_locals::MaybeBorrowedLocals;
 pub use self::liveness::MaybeLiveLocals;
 pub use self::liveness::MaybeTransitiveLiveLocals;
-pub use self::storage_liveness::{MaybeRequiresStorage, MaybeStorageLive};
+pub use self::storage_liveness::{MaybeRequiresStorage, MaybeStorageDead, MaybeStorageLive};
 
 /// `MaybeInitializedPlaces` tracks all places that might be
 /// initialized upon reaching a particular point in the control flow

compiler/rustc_mir_dataflow/src/impls/storage_liveness.rs

@@ -74,6 +74,73 @@ impl<'tcx, 'a> crate::GenKillAnalysis<'tcx> for MaybeStorageLive<'a> {
     }
 }
 
+#[derive(Clone)]
+pub struct MaybeStorageDead {
+    always_live_locals: BitSet<Local>,
+}
+
+impl MaybeStorageDead {
+    pub fn new(always_live_locals: BitSet<Local>) -> Self {
+        MaybeStorageDead { always_live_locals }
+    }
+}
+
+impl<'tcx> crate::AnalysisDomain<'tcx> for MaybeStorageDead {
+    type Domain = BitSet<Local>;
+
+    const NAME: &'static str = "maybe_storage_dead";
+
+    fn bottom_value(&self, body: &mir::Body<'tcx>) -> Self::Domain {
+        // bottom = live
+        BitSet::new_empty(body.local_decls.len())
+    }
+
+    fn initialize_start_block(&self, body: &mir::Body<'tcx>, on_entry: &mut Self::Domain) {
+        assert_eq!(body.local_decls.len(), self.always_live_locals.domain_size());
+        // Do not iterate on return place and args, as they are trivially always live.
+        for local in body.vars_and_temps_iter() {
+            if !self.always_live_locals.contains(local) {
+                on_entry.insert(local);
+            }
+        }
+    }
+}
+
+impl<'tcx> crate::GenKillAnalysis<'tcx> for MaybeStorageDead {
+    type Idx = Local;
+
+    fn statement_effect(
+        &self,
+        trans: &mut impl GenKill<Self::Idx>,
+        stmt: &mir::Statement<'tcx>,
+        _: Location,
+    ) {
+        match stmt.kind {
+            StatementKind::StorageLive(l) => trans.kill(l),
+            StatementKind::StorageDead(l) => trans.gen(l),
+            _ => (),
+        }
+    }
+
+    fn terminator_effect(
+        &self,
+        _trans: &mut impl GenKill<Self::Idx>,
+        _: &mir::Terminator<'tcx>,
+        _: Location,
+    ) {
+        // Terminators have no effect
+    }
+
+    fn call_return_effect(
+        &self,
+        _trans: &mut impl GenKill<Self::Idx>,
+        _block: BasicBlock,
+        _return_places: CallReturnPlaces<'_, 'tcx>,
+    ) {
+        // Nothing to do when a call returns successfully
+    }
+}
+
 type BorrowedLocalsResults<'a, 'tcx> = ResultsRefCursor<'a, 'a, 'tcx, MaybeBorrowedLocals>;
 
 /// Dataflow analysis that determines whether each local requires storage at a

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -33,9 +33,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 borrowed_locals = borrowed_locals(body);
-    let ssa = SsaLocals::new(tcx, param_env, body, &borrowed_locals);
+    let ssa = SsaLocals::new(body);
 
     let fully_moved = fully_moved_locals(&ssa, body);
     debug!(?fully_moved);
@@ -76,7 +75,7 @@ fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
 fn fully_moved_locals(ssa: &SsaLocals, body: &Body<'_>) -> BitSet<Local> {
     let mut fully_moved = BitSet::new_filled(body.local_decls.len());
 
-    for (_, rvalue) in ssa.assignments(body) {
+    for (_, rvalue, _) in ssa.assignments(body) {
         let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) | Rvalue::CopyForDeref(place))
             = rvalue
         else { continue };

compiler/rustc_mir_transform/src/lib.rs

@@ -84,6 +84,7 @@ mod match_branches;
 mod multiple_return_terminators;
 mod normalize_array_len;
 mod nrvo;
+mod ref_prop;
 mod remove_noop_landing_pads;
 mod remove_storage_markers;
 mod remove_uninit_drops;
@@ -559,6 +560,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &separate_const_switch::SeparateConstSwitch,
             &simplify::SimplifyLocals::BeforeConstProp,
             &copy_prop::CopyProp,
+            &ref_prop::ReferencePropagation,
             &const_prop::ConstProp,
             &dataflow_const_prop::DataflowConstProp,
             //

compiler/rustc_mir_transform/src/normalize_array_len.rs

@@ -7,7 +7,6 @@ use rustc_index::IndexVec;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, TyCtxt};
-use rustc_mir_dataflow::impls::borrowed_locals;
 
 pub struct NormalizeArrayLen;
 
@@ -24,9 +23,7 @@ impl<'tcx> MirPass<'tcx> for NormalizeArrayLen {
 }
 
 fn normalize_array_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-    let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
-    let borrowed_locals = borrowed_locals(body);
-    let ssa = SsaLocals::new(tcx, param_env, body, &borrowed_locals);
+    let ssa = SsaLocals::new(body);
 
     let slice_lengths = compute_slice_length(tcx, &ssa, body);
     debug!(?slice_lengths);
@@ -41,7 +38,7 @@ fn compute_slice_length<'tcx>(
 ) -> IndexVec<Local, Option<ty::Const<'tcx>>> {
     let mut slice_lengths = IndexVec::from_elem(None, &body.local_decls);
 
-    for (local, rvalue) in ssa.assignments(body) {
+    for (local, rvalue, _) in ssa.assignments(body) {
         match rvalue {
             Rvalue::Cast(
                 CastKind::Pointer(ty::adjustment::PointerCast::Unsize),