Implement SSA-based reference propagation.

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -76,7 +76,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

@@ -41,7 +41,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),

compiler/rustc_mir_transform/src/ref_prop.rs

@@ -0,0 +1,324 @@
+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::{borrowed_locals, MaybeStorageDead};
+use rustc_mir_dataflow::storage::always_storage_live_locals;
+use rustc_mir_dataflow::Analysis;
+
+use crate::ssa::SsaLocals;
+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` that dominates all uses;
+/// - 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 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 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> {
+    // Compute `MaybeStorageDead` dataflow to check that we only replace when the pointee is
+    // definitely live.
+    let always_live_locals = always_storage_live_locals(body);
+    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);
+
+    let mut can_perform_opt = |target: Place<'tcx>, loc: Location| {
+        maybe_dead.seek_after_primary_effect(loc);
+        let maybe_dead = maybe_dead.contains(target.local);
+
+        if target.projection.first() == Some(&PlaceElem::Deref) {
+            // We are creating a reborrow. As `place.local` is a reference, removing the
+            // `StorageDead` is fine.
+            if maybe_dead {
+                storage_to_remove.insert(target.local);
+            }
+            true
+        } else {
+            // This is a proper dereference. We can only allow it if `target` is live.
+            !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 ssa.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),
+        }
+    }
+}

compiler/rustc_mir_transform/src/ssa.rs

@@ -6,6 +6,7 @@ use rustc_middle::middle::resolve_bound_vars::Set1;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{ParamEnv, TyCtxt};
+use rustc_mir_dataflow::storage::always_storage_live_locals;
 
 #[derive(Debug)]
 pub struct SsaLocals {
@@ -17,6 +18,12 @@ 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>,
+    /// Set of "StorageLive" statements for each local. When the "StorageLive" statement does not
+    /// dominate all uses of the local, we mark it as `Set1::Many`.
+    storage_live: IndexVec<Local, Set1<LocationExtended>>,
 }
 
 /// We often encounter MIR bodies with 1 or 2 basic blocks. In those cases, it's unnecessary to
@@ -26,23 +33,31 @@ 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;
+        }
     }
 }
 
@@ -59,7 +74,11 @@ impl SsaLocals {
         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 };
+
+        let direct_uses = IndexVec::from_elem(0, &body.local_decls);
+        let storage_live = IndexVec::from_elem(Set1::Empty, &body.local_decls);
+        let mut visitor =
+            SsaVisitor { assignments, assignment_order, dominators, direct_uses, storage_live };
 
         for (local, decl) in body.local_decls.iter_enumerated() {
             if matches!(body.local_kind(local), LocalKind::Arg) {
@@ -70,6 +89,10 @@ impl SsaLocals {
             }
         }
 
+        for local in always_storage_live_locals(body).iter() {
+            visitor.storage_live[local] = Set1::One(LocationExtended::Arg);
+        }
+
         if body.basic_blocks.len() > 2 {
             for (bb, data) in traversal::reverse_postorder(body) {
                 visitor.visit_basic_block_data(bb, data);
@@ -85,36 +108,66 @@ impl SsaLocals {
         }
 
         debug!(?visitor.assignments);
+        debug!(?visitor.direct_uses);
+        debug!(?visitor.storage_live);
 
         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,
+            storage_live: visitor.storage_live,
             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(_))
     }
 
+    /// Returns true iff we can use `p` as a pointee.
+    pub fn is_constant_place(&self, p: Place<'_>) -> bool {
+        // We only allow `Deref` as the first projection, to avoid surprises.
+        if p.projection.first() == Some(&PlaceElem::Deref) {
+            // `p == (*some_local).xxx`, it is constant only if `some_local` is constant.
+            // We approximate constness using SSAness.
+            self.is_ssa(p.local) && p.projection[1..].iter().all(PlaceElem::is_stable_offset)
+        } else {
+            matches!(self.storage_live[p.local], Set1::One(_))
+                && p.projection[..].iter().all(PlaceElem::is_stable_offset)
+        }
+    }
+
+    /// 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,25 +230,8 @@ 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>,
+    storage_live: IndexVec<Local, Set1<LocationExtended>>,
 }
 
 impl<'tcx> Visitor<'tcx> for SsaVisitor {
@@ -207,14 +243,23 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor {
                     // Only record if SSA-like, to avoid growing the vector needlessly.
                     self.assignment_order.push(local);
                 }
+                self.dominators.check_dominates(&mut self.storage_live[local], loc);
             }
             // Anything can happen with raw pointers, so remove them.
             PlaceContext::NonMutatingUse(NonMutatingUseContext::AddressOf)
-            | PlaceContext::MutatingUse(_) => self.assignments[local] = Set1::Many,
+            | PlaceContext::MutatingUse(_) => {
+                self.assignments[local] = Set1::Many;
+                self.dominators.check_dominates(&mut self.storage_live[local], loc);
+            }
             // Immutable borrows are taken into account in `SsaLocals::new` by
             // removing non-freeze locals.
             PlaceContext::NonMutatingUse(_) => {
-                self.check_assignment_dominates(local, loc);
+                self.dominators.check_dominates(&mut self.assignments[local], loc);
+                self.dominators.check_dominates(&mut self.storage_live[local], loc);
+                self.direct_uses[local] += 1;
+            }
+            PlaceContext::NonUse(NonUseContext::StorageLive) => {
+                self.storage_live[local].insert(LocationExtended::Plain(loc));
             }
             PlaceContext::NonUse(_) => {}
         }
@@ -224,11 +269,12 @@ 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.
+                // Only change the context if it is a real use, not a "use" in debuginfo.
                 let new_ctxt = PlaceContext::NonMutatingUse(NonMutatingUseContext::Projection);
 
                 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);
+                self.dominators.check_dominates(&mut self.storage_live[place.local], loc);
             }
             return;
         }
@@ -237,7 +283,7 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor {
 }
 
 #[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 +313,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)]