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.

5 files changed, 474 insertions, 68 deletions

diff --git a/compiler/rustc_mir_transform/src/copy_prop.rs b/compiler/rustc_mir_transform/src/copy_prop.rs
index 3922ed2fbf7..c565d6f13b1 100644
--- a/compiler/rustc_mir_transform/src/copy_prop.rs
+++ b/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 };
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 5c4b1ead4e9..277237a5515 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/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,
             //
diff --git a/compiler/rustc_mir_transform/src/normalize_array_len.rs b/compiler/rustc_mir_transform/src/normalize_array_len.rs
index b3b831bb4ab..3d61d33ce35 100644
--- a/compiler/rustc_mir_transform/src/normalize_array_len.rs
+++ b/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),
diff --git a/compiler/rustc_mir_transform/src/ref_prop.rs b/compiler/rustc_mir_transform/src/ref_prop.rs
new file mode 100644
index 00000000000..dafd2ae23a6
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/ref_prop.rs
@@ -0,0 +1,355 @@
+use rustc_data_structures::fx::FxHashSet;
+use rustc_index::bit_set::BitSet;
+use rustc_index::IndexVec;
+use rustc_middle::mir::visit::*;
+use rustc_middle::mir::*;
+use rustc_middle::ty::TyCtxt;
+use rustc_mir_dataflow::impls::MaybeStorageDead;
+use rustc_mir_dataflow::storage::always_storage_live_locals;
+use rustc_mir_dataflow::Analysis;
+
+use crate::ssa::{SsaLocals, StorageLiveLocals};
+use crate::MirPass;
+
+/// Propagate references using SSA analysis.
+///
+/// MIR building may produce a lot of borrow-dereference patterns.
+///
+/// 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.
+///
+/// There are 3 properties that need to be upheld for this transformation to be legal:
+/// - place stability: `PLACE` must refer to the same memory wherever it appears;
+/// - pointer liveness: we must not introduce dereferences of dangling pointers;
+/// - `&mut` borrow uniqueness.
+///
+/// # Stability
+///
+/// If `PLACE` is an indirect projection, if its of the form `(*LOCAL).PROJECTIONS` where:
+/// - `LOCAL` is SSA;
+/// - all projections in `PROJECTIONS` have a stable offset (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`;
+/// - all projections have a stable offset.
+///
+/// # Liveness
+///
+/// When performing a substitution, we must take care not to introduce uses of dangling locals.
+/// To ensure this, we walk the body with the `MaybeStorageDead` dataflow analysis:
+/// - if we want to replace `*x` by reborrow `*y` and `y` may be dead, we allow replacement and
+///   mark storage statements on `y` for removal;
+/// - if we want to replace `*x` by non-reborrow `y` and `y` must be live, we allow replacement;
+/// - if we want to replace `*x` by non-reborrow `y` and `y` may be dead, we do not replace.
+///
+/// # 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.
+pub struct ReferencePropagation;
+
+impl<'tcx> MirPass<'tcx> for ReferencePropagation {
+    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
+        sess.mir_opt_level() >= 4
+    }
+
+    #[instrument(level = "trace", skip(self, tcx, body))]
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        debug!(def_id = ?body.source.def_id());
+        propagate_ssa(tcx, body);
+    }
+}
+
+fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+    let ssa = SsaLocals::new(body);
+
+    let mut replacer = compute_replacement(tcx, body, &ssa);
+    debug!(?replacer.targets, ?replacer.allowed_replacements, ?replacer.storage_to_remove);
+
+    replacer.visit_body_preserves_cfg(body);
+
+    if replacer.any_replacement {
+        crate::simplify::remove_unused_definitions(body);
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+enum Value<'tcx> {
+    /// Not a pointer, or we can't know.
+    Unknown,
+    /// We know the value to be a pointer to this place.
+    /// The boolean indicates whether the reference is mutable, subject the uniqueness rule.
+    Pointer(Place<'tcx>, bool),
+}
+
+/// For each local, save the place corresponding to `*local`.
+#[instrument(level = "trace", skip(tcx, body))]
+fn compute_replacement<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    body: &Body<'tcx>,
+    ssa: &SsaLocals,
+) -> Replacer<'tcx> {
+    let always_live_locals = always_storage_live_locals(body);
+
+    // Compute which locals have a single `StorageLive` statement ever.
+    let storage_live = StorageLiveLocals::new(body, &always_live_locals);
+
+    // Compute `MaybeStorageDead` dataflow to check that we only replace when the pointee is
+    // definitely live.
+    let mut maybe_dead = MaybeStorageDead::new(always_live_locals)
+        .into_engine(tcx, body)
+        .iterate_to_fixpoint()
+        .into_results_cursor(body);
+
+    // Map for each local to the pointee.
+    let mut targets = IndexVec::from_elem(Value::Unknown, &body.local_decls);
+    // Set of locals for which we will remove their storage statement. This is useful for
+    // reborrowed references.
+    let mut storage_to_remove = BitSet::new_empty(body.local_decls.len());
+
+    let fully_replacable_locals = fully_replacable_locals(ssa);
+
+    // Returns true iff we can use `place` as a pointee.
+    //
+    // Note that we only need to verify that there is a single `StorageLive` statement, and we do
+    // not need to verify that it dominates all uses of that local.
+    //
+    // Consider the three statements:
+    //   SL : StorageLive(a)
+    //   DEF: b = &raw? mut? a
+    //   USE: stuff that uses *b
+    //
+    // First, we recall that DEF is checked to dominate USE. Now imagine for the sake of
+    // contradiction there is a DEF -> SL -> USE path. Consider two cases:
+    //
+    // - DEF dominates SL. We always have UB the first time control flow reaches DEF,
+    //   because the storage of `a` is dead. Since DEF dominates USE, that means we cannot
+    //   reach USE and so our optimization is ok.
+    //
+    // - DEF does not dominate SL. Then there is a `START_BLOCK -> SL` path not including DEF.
+    //   But we can extend this path to USE, meaning there is also a `START_BLOCK -> USE` path not
+    //   including DEF. This violates the DEF dominates USE condition, and so is impossible.
+    let is_constant_place = |place: Place<'_>| {
+        // We only allow `Deref` as the first projection, to avoid surprises.
+        if place.projection.first() == Some(&PlaceElem::Deref) {
+            // `place == (*some_local).xxx`, it is constant only if `some_local` is constant.
+            // We approximate constness using SSAness.
+            ssa.is_ssa(place.local) && place.projection[1..].iter().all(PlaceElem::is_stable_offset)
+        } else {
+            storage_live.has_single_storage(place.local)
+                && place.projection[..].iter().all(PlaceElem::is_stable_offset)
+        }
+    };
+
+    let mut can_perform_opt = |target: Place<'tcx>, loc: Location| {
+        if target.projection.first() == Some(&PlaceElem::Deref) {
+            // We are creating a reborrow. As `place.local` is a reference, removing the storage
+            // statements should not make it much harder for LLVM to optimize.
+            storage_to_remove.insert(target.local);
+            true
+        } else {
+            // This is a proper dereference. We can only allow it if `target` is live.
+            maybe_dead.seek_after_primary_effect(loc);
+            let maybe_dead = maybe_dead.contains(target.local);
+            !maybe_dead
+        }
+    };
+
+    for (local, rvalue, location) in ssa.assignments(body) {
+        debug!(?local);
+
+        // Only visit if we have something to do.
+        let Value::Unknown = targets[local] else { bug!() };
+
+        let ty = body.local_decls[local].ty;
+
+        // If this is not a reference or pointer, do nothing.
+        if !ty.is_any_ptr() {
+            debug!("not a reference or pointer");
+            continue;
+        }
+
+        // If this a mutable reference that we cannot fully replace, mark it as unknown.
+        if ty.is_mutable_ptr() && !fully_replacable_locals.contains(local) {
+            debug!("not fully replaceable");
+            continue;
+        }
+
+        debug!(?rvalue);
+        match rvalue {
+            // This is a copy, just use the value we have in store for the previous one.
+            // As we are visiting in `assignment_order`, ie. reverse postorder, `rhs` should
+            // have been visited before.
+            Rvalue::Use(Operand::Copy(place) | Operand::Move(place))
+            | Rvalue::CopyForDeref(place) => {
+                if let Some(rhs) = place.as_local() {
+                    let target = targets[rhs];
+                    if matches!(target, Value::Pointer(..)) {
+                        targets[local] = target;
+                    } else if ssa.is_ssa(rhs) {
+                        let refmut = body.local_decls[rhs].ty.is_mutable_ptr();
+                        targets[local] = Value::Pointer(tcx.mk_place_deref(rhs.into()), refmut);
+                    }
+                }
+            }
+            Rvalue::Ref(_, _, place) | Rvalue::AddressOf(_, place) => {
+                let mut place = *place;
+                // Try to see through `place` in order to collapse reborrow chains.
+                if place.projection.first() == Some(&PlaceElem::Deref)
+                    && let Value::Pointer(target, refmut) = targets[place.local]
+                    // Only see through immutable reference and pointers, as we do not know yet if
+                    // mutable references are fully replaced.
+                    && !refmut
+                    // Only collapse chain if the pointee is definitely live.
+                    && can_perform_opt(target, location)
+                {
+                    place = target.project_deeper(&place.projection[1..], tcx);
+                }
+                assert_ne!(place.local, local);
+                if is_constant_place(place) {
+                    targets[local] = Value::Pointer(place, ty.is_mutable_ptr());
+                }
+            }
+            // We do not know what to do, so keep as not-a-pointer.
+            _ => {}
+        }
+    }
+
+    debug!(?targets);
+
+    let mut finder = ReplacementFinder {
+        targets: &mut targets,
+        can_perform_opt,
+        allowed_replacements: FxHashSet::default(),
+    };
+    let reachable_blocks = traversal::reachable_as_bitset(body);
+    for (bb, bbdata) in body.basic_blocks.iter_enumerated() {
+        // Only visit reachable blocks as we rely on dataflow.
+        if reachable_blocks.contains(bb) {
+            finder.visit_basic_block_data(bb, bbdata);
+        }
+    }
+
+    let allowed_replacements = finder.allowed_replacements;
+    return Replacer {
+        tcx,
+        targets,
+        storage_to_remove,
+        allowed_replacements,
+        any_replacement: false,
+    };
+
+    struct ReplacementFinder<'a, 'tcx, F> {
+        targets: &'a mut IndexVec<Local, Value<'tcx>>,
+        can_perform_opt: F,
+        allowed_replacements: FxHashSet<(Local, Location)>,
+    }
+
+    impl<'tcx, F> Visitor<'tcx> for ReplacementFinder<'_, 'tcx, F>
+    where
+        F: FnMut(Place<'tcx>, Location) -> bool,
+    {
+        fn visit_place(&mut self, place: &Place<'tcx>, ctxt: PlaceContext, loc: Location) {
+            if matches!(ctxt, PlaceContext::NonUse(_)) {
+                // There is no need to check liveness for non-uses.
+                return;
+            }
+
+            if let Value::Pointer(target, refmut) = self.targets[place.local]
+                && place.projection.first() == Some(&PlaceElem::Deref)
+            {
+                let perform_opt = (self.can_perform_opt)(target, loc);
+                if perform_opt {
+                    self.allowed_replacements.insert((target.local, loc));
+                } else if refmut {
+                    // This mutable reference is not fully replacable, so drop it.
+                    self.targets[place.local] = Value::Unknown;
+                }
+            }
+        }
+    }
+}
+
+/// Compute the set of locals that can be fully replaced.
+///
+/// We consider a local to be replacable iff it's only used in a `Deref` projection `*_local` or
+/// non-use position (like storage statements and debuginfo).
+fn fully_replacable_locals(ssa: &SsaLocals) -> BitSet<Local> {
+    let mut replacable = BitSet::new_empty(ssa.num_locals());
+
+    // First pass: for each local, whether its uses can be fully replaced.
+    for local in ssa.locals() {
+        if ssa.num_direct_uses(local) == 0 {
+            replacable.insert(local);
+        }
+    }
+
+    // Second pass: a local can only be fully replaced if all its copies can.
+    ssa.meet_copy_equivalence(&mut replacable);
+
+    replacable
+}
+
+/// Utility to help performing subtitution of `*pattern` by `target`.
+struct Replacer<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    targets: IndexVec<Local, Value<'tcx>>,
+    storage_to_remove: BitSet<Local>,
+    allowed_replacements: FxHashSet<(Local, Location)>,
+    any_replacement: bool,
+}
+
+impl<'tcx> MutVisitor<'tcx> for Replacer<'tcx> {
+    fn tcx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn visit_place(&mut self, place: &mut Place<'tcx>, ctxt: PlaceContext, loc: Location) {
+        if let Value::Pointer(target, _) = self.targets[place.local]
+            && place.projection.first() == Some(&PlaceElem::Deref)
+        {
+            let perform_opt = matches!(ctxt, PlaceContext::NonUse(_))
+                || self.allowed_replacements.contains(&(target.local, loc));
+
+            if perform_opt {
+                *place = target.project_deeper(&place.projection[1..], self.tcx);
+                self.any_replacement = true;
+            }
+        } else {
+            self.super_place(place, ctxt, loc);
+        }
+    }
+
+    fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
+        match stmt.kind {
+            StatementKind::StorageLive(l) | StatementKind::StorageDead(l)
+                if self.storage_to_remove.contains(l) =>
+            {
+                stmt.make_nop();
+            }
+            // Do not remove assignments as they may still be useful for debuginfo.
+            _ => self.super_statement(stmt, loc),
+        }
+    }
+}
diff --git a/compiler/rustc_mir_transform/src/ssa.rs b/compiler/rustc_mir_transform/src/ssa.rs
index ec8d42c1652..05a7b226f0c 100644
--- a/compiler/rustc_mir_transform/src/ssa.rs
+++ b/compiler/rustc_mir_transform/src/ssa.rs
@@ -1,3 +1,10 @@
+//! We denote as "SSA" the set of locals that verify the following properties:
+//! 1/ They are only assigned-to once, either as a function parameter, or in an assign statement;
+//! 2/ This single assignment dominates all uses;
+//!
+//! As a consequence of rule 2, we consider that borrowed locals are not SSA, even if they are
+//! `Freeze`, as we do not track that the assignment dominates all uses of the borrow.
+
 use either::Either;
 use rustc_data_structures::graph::dominators::Dominators;
 use rustc_index::bit_set::BitSet;
@@ -5,7 +12,6 @@ use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::middle::resolve_bound_vars::Set1;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
-use rustc_middle::ty::{ParamEnv, TyCtxt};
 
 #[derive(Debug)]
 pub struct SsaLocals {
@@ -17,6 +23,9 @@ pub struct SsaLocals {
     assignment_order: Vec<Local>,
     /// Copy equivalence classes between locals. See `copy_classes` for documentation.
     copy_classes: IndexVec<Local, Local>,
+    /// Number of "direct" uses of each local, ie. uses that are not dereferences.
+    /// We ignore non-uses (Storage statements, debuginfo).
+    direct_uses: IndexVec<Local, u32>,
 }
 
 /// We often encounter MIR bodies with 1 or 2 basic blocks. In those cases, it's unnecessary to
@@ -26,48 +35,48 @@ struct SmallDominators {
     inner: Option<Dominators<BasicBlock>>,
 }
 
-trait DomExt {
-    fn dominates(self, _other: Self, dominators: &SmallDominators) -> bool;
-}
-
-impl DomExt for Location {
-    fn dominates(self, other: Location, dominators: &SmallDominators) -> bool {
-        if self.block == other.block {
-            self.statement_index <= other.statement_index
+impl SmallDominators {
+    fn dominates(&self, first: Location, second: Location) -> bool {
+        if first.block == second.block {
+            first.statement_index <= second.statement_index
+        } else if let Some(inner) = &self.inner {
+            inner.dominates(first.block, second.block)
         } else {
-            dominators.dominates(self.block, other.block)
+            first.block < second.block
         }
     }
-}
 
-impl SmallDominators {
-    fn dominates(&self, dom: BasicBlock, node: BasicBlock) -> bool {
-        if let Some(inner) = &self.inner { inner.dominates(dom, node) } else { dom < node }
+    fn check_dominates(&mut self, set: &mut Set1<LocationExtended>, loc: Location) {
+        let assign_dominates = match *set {
+            Set1::Empty | Set1::Many => false,
+            Set1::One(LocationExtended::Arg) => true,
+            Set1::One(LocationExtended::Plain(assign)) => {
+                self.dominates(assign.successor_within_block(), loc)
+            }
+        };
+        // We are visiting a use that is not dominated by an assignment.
+        // Either there is a cycle involved, or we are reading for uninitialized local.
+        // Bail out.
+        if !assign_dominates {
+            *set = Set1::Many;
+        }
     }
 }
 
 impl SsaLocals {
-    pub fn new<'tcx>(
-        tcx: TyCtxt<'tcx>,
-        param_env: ParamEnv<'tcx>,
-        body: &Body<'tcx>,
-        borrowed_locals: &BitSet<Local>,
-    ) -> SsaLocals {
+    pub fn new<'tcx>(body: &Body<'tcx>) -> SsaLocals {
         let assignment_order = Vec::with_capacity(body.local_decls.len());
 
         let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls);
         let dominators =
             if body.basic_blocks.len() > 2 { Some(body.basic_blocks.dominators()) } else { None };
         let dominators = SmallDominators { inner: dominators };
-        let mut visitor = SsaVisitor { assignments, assignment_order, dominators };
 
-        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_locals.contains(local) && !decl.ty.is_freeze(tcx, param_env) {
-                visitor.assignments[local] = Set1::Many;
-            }
+        let direct_uses = IndexVec::from_elem(0, &body.local_decls);
+        let mut visitor = SsaVisitor { assignments, assignment_order, dominators, direct_uses };
+
+        for local in body.args_iter() {
+            visitor.assignments[local] = Set1::One(LocationExtended::Arg);
         }
 
         if body.basic_blocks.len() > 2 {
@@ -85,36 +94,51 @@ impl SsaLocals {
         }
 
         debug!(?visitor.assignments);
+        debug!(?visitor.direct_uses);
 
         visitor
             .assignment_order
             .retain(|&local| matches!(visitor.assignments[local], Set1::One(_)));
         debug!(?visitor.assignment_order);
 
-        let copy_classes = compute_copy_classes(&visitor, body);
+        let copy_classes = compute_copy_classes(&mut visitor, body);
 
         SsaLocals {
             assignments: visitor.assignments,
             assignment_order: visitor.assignment_order,
+            direct_uses: visitor.direct_uses,
             copy_classes,
         }
     }
 
+    pub fn num_locals(&self) -> usize {
+        self.assignments.len()
+    }
+
+    pub fn locals(&self) -> impl Iterator<Item = Local> {
+        self.assignments.indices()
+    }
+
     pub fn is_ssa(&self, local: Local) -> bool {
         matches!(self.assignments[local], Set1::One(_))
     }
 
+    /// Return the number of uses if a local that are not "Deref".
+    pub fn num_direct_uses(&self, local: Local) -> u32 {
+        self.direct_uses[local]
+    }
+
     pub fn assignments<'a, 'tcx>(
         &'a self,
         body: &'a Body<'tcx>,
-    ) -> impl Iterator<Item = (Local, &'a Rvalue<'tcx>)> + 'a {
+    ) -> impl Iterator<Item = (Local, &'a Rvalue<'tcx>, Location)> + 'a {
         self.assignment_order.iter().filter_map(|&local| {
             if let Set1::One(LocationExtended::Plain(loc)) = self.assignments[local] {
                 // `loc` must point to a direct assignment to `local`.
                 let Either::Left(stmt) = body.stmt_at(loc) else { bug!() };
                 let Some((target, rvalue)) = stmt.kind.as_assign() else { bug!() };
                 assert_eq!(target.as_local(), Some(local));
-                Some((local, rvalue))
+                Some((local, rvalue, loc))
             } else {
                 None
             }
@@ -177,30 +201,14 @@ struct SsaVisitor {
     dominators: SmallDominators,
     assignments: IndexVec<Local, Set1<LocationExtended>>,
     assignment_order: Vec<Local>,
-}
-
-impl SsaVisitor {
-    fn check_assignment_dominates(&mut self, local: Local, loc: Location) {
-        let set = &mut self.assignments[local];
-        let assign_dominates = match *set {
-            Set1::Empty | Set1::Many => false,
-            Set1::One(LocationExtended::Arg) => true,
-            Set1::One(LocationExtended::Plain(assign)) => {
-                assign.successor_within_block().dominates(loc, &self.dominators)
-            }
-        };
-        // We are visiting a use that is not dominated by an assignment.
-        // Either there is a cycle involved, or we are reading for uninitialized local.
-        // Bail out.
-        if !assign_dominates {
-            *set = Set1::Many;
-        }
-    }
+    direct_uses: IndexVec<Local, u32>,
 }
 
 impl<'tcx> Visitor<'tcx> for SsaVisitor {
     fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) {
         match ctxt {
+            PlaceContext::MutatingUse(MutatingUseContext::Projection)
+            | PlaceContext::NonMutatingUse(NonMutatingUseContext::Projection) => bug!(),
             PlaceContext::MutatingUse(MutatingUseContext::Store) => {
                 self.assignments[local].insert(LocationExtended::Plain(loc));
                 if let Set1::One(_) = self.assignments[local] {
@@ -209,12 +217,20 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor {
                 }
             }
             // Anything can happen with raw pointers, so remove them.
-            PlaceContext::NonMutatingUse(NonMutatingUseContext::AddressOf)
-            | PlaceContext::MutatingUse(_) => self.assignments[local] = Set1::Many,
-            // Immutable borrows are taken into account in `SsaLocals::new` by
-            // removing non-freeze locals.
+            // We do not verify that all uses of the borrow dominate the assignment to `local`,
+            // so we have to remove them too.
+            PlaceContext::NonMutatingUse(
+                NonMutatingUseContext::SharedBorrow
+                | NonMutatingUseContext::ShallowBorrow
+                | NonMutatingUseContext::UniqueBorrow
+                | NonMutatingUseContext::AddressOf,
+            )
+            | PlaceContext::MutatingUse(_) => {
+                self.assignments[local] = Set1::Many;
+            }
             PlaceContext::NonMutatingUse(_) => {
-                self.check_assignment_dominates(local, loc);
+                self.dominators.check_dominates(&mut self.assignments[local], loc);
+                self.direct_uses[local] += 1;
             }
             PlaceContext::NonUse(_) => {}
         }
@@ -224,20 +240,22 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor {
         if place.projection.first() == Some(&PlaceElem::Deref) {
             // Do not do anything for storage statements and debuginfo.
             if ctxt.is_use() {
-                // A use through a `deref` only reads from the local, and cannot write to it.
-                let new_ctxt = PlaceContext::NonMutatingUse(NonMutatingUseContext::Projection);
+                // Only change the context if it is a real use, not a "use" in debuginfo.
+                let new_ctxt = PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy);
 
                 self.visit_projection(place.as_ref(), new_ctxt, loc);
-                self.check_assignment_dominates(place.local, loc);
+                self.dominators.check_dominates(&mut self.assignments[place.local], loc);
             }
             return;
+        } else {
+            self.visit_projection(place.as_ref(), ctxt, loc);
+            self.visit_local(place.local, ctxt, loc);
         }
-        self.super_place(place, ctxt, loc);
     }
 }
 
 #[instrument(level = "trace", skip(ssa, body))]
-fn compute_copy_classes(ssa: &SsaVisitor, body: &Body<'_>) -> IndexVec<Local, Local> {
+fn compute_copy_classes(ssa: &mut SsaVisitor, body: &Body<'_>) -> IndexVec<Local, Local> {
     let mut copies = IndexVec::from_fn_n(|l| l, body.local_decls.len());
 
     for &local in &ssa.assignment_order {
@@ -267,9 +285,11 @@ fn compute_copy_classes(ssa: &SsaVisitor, body: &Body<'_>) -> IndexVec<Local, Lo
         // We visit in `assignment_order`, ie. reverse post-order, so `rhs` has been
         // visited before `local`, and we just have to copy the representing local.
         copies[local] = copies[rhs];
+        ssa.direct_uses[rhs] -= 1;
     }
 
     debug!(?copies);
+    debug!(?ssa.direct_uses);
 
     // Invariant: `copies` must point to the head of an equivalence class.
     #[cfg(debug_assertions)]
@@ -279,3 +299,36 @@ fn compute_copy_classes(ssa: &SsaVisitor, body: &Body<'_>) -> IndexVec<Local, Lo
 
     copies
 }
+
+#[derive(Debug)]
+pub(crate) struct StorageLiveLocals {
+    /// Set of "StorageLive" statements for each local.
+    storage_live: IndexVec<Local, Set1<LocationExtended>>,
+}
+
+impl StorageLiveLocals {
+    pub(crate) fn new(
+        body: &Body<'_>,
+        always_storage_live_locals: &BitSet<Local>,
+    ) -> StorageLiveLocals {
+        let mut storage_live = IndexVec::from_elem(Set1::Empty, &body.local_decls);
+        for local in always_storage_live_locals.iter() {
+            storage_live[local] = Set1::One(LocationExtended::Arg);
+        }
+        for (block, bbdata) in body.basic_blocks.iter_enumerated() {
+            for (statement_index, statement) in bbdata.statements.iter().enumerate() {
+                if let StatementKind::StorageLive(local) = statement.kind {
+                    storage_live[local]
+                        .insert(LocationExtended::Plain(Location { block, statement_index }));
+                }
+            }
+        }
+        debug!(?storage_live);
+        StorageLiveLocals { storage_live }
+    }
+
+    #[inline]
+    pub(crate) fn has_single_storage(&self, local: Local) -> bool {
+        matches!(self.storage_live[local], Set1::One(_))
+    }
+}