Extend GVN to perform local value numbering.

Author: cjgillot

compiler/rustc_mir_transform/src/gvn.rs

@@ -1,4 +1,4 @@
-//! Global value numbering.
+//! Value numbering.
 //!
 //! MIR may contain repeated and/or redundant computations. The objective of this pass is to detect
 //! such redundancies and re-use the already-computed result when possible.
@@ -8,15 +8,23 @@
 //!
 //! We traverse all assignments `x = rvalue` and operands.
 //!
-//! For each SSA one, we compute a symbolic representation of values that are assigned to SSA
+//! For each assignment, we compute a symbolic representation of values that are assigned to SSA
 //! locals. This symbolic representation is defined by the `Value` enum. Each produced instance of
 //! `Value` is interned as a `VnIndex`, which allows us to cheaply compute identical values.
 //!
-//! For each non-SSA
-//! one, we compute the `VnIndex` of the rvalue. If this `VnIndex` is associated to a constant, we
-//! replace the rvalue/operand by that constant. Otherwise, if there is an SSA local `y`
-//! associated to this `VnIndex`, and if its definition location strictly dominates the assignment
-//! to `x`, we replace the assignment by `x = y`.
+//! If the local is SSA, we append it into the mapping `rev_locals_ssa[value]` for later reuse.
+//! That mapping accumulates the knowledge across basic blocks.
+//!
+//! If the local is non-SSA, we remove it from `rev_locals_non_ssa[old_value]` and append it to
+//! `rev_locals_non_ssa[new_value]`. That mapping is cleared at the beginning of each basic block,
+//! to ensure we do not carry information across blocks.
+//!
+//! If the computed `VnIndex` is associated to a constant, we replace the rvalue/operand by that
+//! constant. Otherwise, if there is a local `y` associated to this `VnIndex`:
+//! - if `y` is SSA and its definition location strictly dominates the assignment to `x`, we
+//!   replace the assignment by `x = y`;
+//! - if `y` is not SSA, then the assignment happens earlier in the same block, so we replace the
+//!   assignment by `x = y`.
 //!
 //! By opportunity, this pass simplifies some `Rvalue`s based on the accumulated knowledge.
 //!
@@ -97,7 +105,7 @@ use rustc_const_eval::interpret::{
     ImmTy, Immediate, InterpCx, MemPlaceMeta, MemoryKind, OpTy, Projectable, Scalar,
     intern_const_alloc_for_constprop,
 };
-use rustc_data_structures::fx::FxHasher;
+use rustc_data_structures::fx::{FxHashMap, FxHasher};
 use rustc_data_structures::graph::dominators::Dominators;
 use rustc_hir::def::DefKind;
 use rustc_index::bit_set::DenseBitSet;
@@ -358,8 +366,11 @@ struct VnState<'body, 'a, 'tcx> {
     /// Value stored in each local.
     locals: IndexVec<Local, Option<VnIndex>>,
     /// Locals that are assigned that value.
-    // This vector does not hold all the values of `VnIndex` that we create.
-    rev_locals: IndexVec<VnIndex, SmallVec<[Local; 1]>>,
+    // This vector holds the locals that are SSA.
+    rev_locals_ssa: IndexVec<VnIndex, SmallVec<[Local; 1]>>,
+    // This map holds the locals that are not SSA. This map is cleared at the end of each block.
+    // Therefore, we do not need a location, the local always appears before the current location.
+    rev_locals_non_ssa: FxHashMap<VnIndex, SmallVec<[Local; 1]>>,
     values: ValueSet<'a, 'tcx>,
     /// Values evaluated as constants if possible.
     evaluated: IndexVec<VnIndex, Option<OpTy<'tcx>>>,
@@ -394,7 +405,8 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             local_decls,
             is_coroutine: body.coroutine.is_some(),
             locals: IndexVec::from_elem(None, local_decls),
-            rev_locals: IndexVec::with_capacity(num_values),
+            rev_locals_ssa: IndexVec::with_capacity(num_values),
+            rev_locals_non_ssa: FxHashMap::default(),
             values: ValueSet::new(num_values),
             evaluated: IndexVec::with_capacity(num_values),
             derefs: Vec::new(),
@@ -413,11 +425,11 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
     fn insert(&mut self, ty: Ty<'tcx>, value: Value<'a, 'tcx>) -> VnIndex {
         let (index, new) = self.values.insert(ty, value);
         if new {
-            // Grow `evaluated` and `rev_locals` here to amortize the allocations.
+            // Grow `evaluated` and `rev_locals_ssa` here to amortize the allocations.
             let evaluated = self.eval_to_const(index);
             let _index = self.evaluated.push(evaluated);
             debug_assert_eq!(index, _index);
-            let _index = self.rev_locals.push(SmallVec::new());
+            let _index = self.rev_locals_ssa.push(Default::default());
             debug_assert_eq!(index, _index);
         }
         index
@@ -430,7 +442,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         let index = self.values.insert_unique(ty, Value::Opaque);
         let _index = self.evaluated.push(None);
         debug_assert_eq!(index, _index);
-        let _index = self.rev_locals.push(SmallVec::new());
+        let _index = self.rev_locals_ssa.push(SmallVec::new());
         debug_assert_eq!(index, _index);
         index
     }
@@ -448,17 +460,19 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
 
         let mut projection = place.projection.iter();
         let base = if place.is_indirect_first_projection() {
-            let base = self.locals[place.local]?;
+            let base = self.local(place.local);
             // Skip the initial `Deref`.
             projection.next();
             AddressBase::Deref(base)
         } else {
             AddressBase::Local(place.local)
         };
+        let arena = self.arena;
+
         // Do not try evaluating inside `Index`, this has been done by `simplify_place_projection`.
         let projection =
-            projection.map(|proj| proj.try_map(|index| self.locals[index], |ty| ty).ok_or(()));
-        let projection = self.arena.try_alloc_from_iter(projection).ok()?;
+            projection.map(|proj| proj.try_map(|index| Some(self.local(index)), |ty| ty).ok_or(()));
+        let projection = arena.try_alloc_from_iter(projection).ok()?;
 
         let index = self.values.insert_unique(ty, |provenance| Value::Address {
             base,
@@ -469,7 +483,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         let evaluated = self.eval_to_const(index);
         let _index = self.evaluated.push(evaluated);
         debug_assert_eq!(index, _index);
-        let _index = self.rev_locals.push(SmallVec::new());
+        let _index = self.rev_locals_ssa.push(SmallVec::new());
         debug_assert_eq!(index, _index);
 
         Some(index)
@@ -494,7 +508,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             let evaluated = self.eval_to_const(index);
             let _index = self.evaluated.push(evaluated);
             debug_assert_eq!(index, _index);
-            let _index = self.rev_locals.push(SmallVec::new());
+            let _index = self.rev_locals_ssa.push(SmallVec::new());
             debug_assert_eq!(index, _index);
         }
         index
@@ -510,12 +524,49 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         self.values.ty(index)
     }
 
-    /// Record that `local` is assigned `value`. `local` must be SSA.
+    /// Record that `local` is assigned `value`.
     #[instrument(level = "trace", skip(self))]
     fn assign(&mut self, local: Local, value: VnIndex) {
-        debug_assert!(self.ssa.is_ssa(local));
         self.locals[local] = Some(value);
-        self.rev_locals[value].push(local);
+        if self.ssa.is_ssa(local) {
+            self.rev_locals_ssa[value].push(local);
+        } else {
+            self.rev_locals_non_ssa.entry(value).or_default().push(local);
+        }
+    }
+
+    /// Return the value assigned to a local, or assign an opaque value and return it.
+    #[instrument(level = "trace", skip(self), ret)]
+    fn local(&mut self, local: Local) -> VnIndex {
+        if let Some(value) = self.locals[local] {
+            return value;
+        }
+        let value = self.new_opaque(self.local_decls[local].ty);
+        self.locals[local] = Some(value);
+        self.rev_locals_non_ssa.entry(value).or_default().push(local);
+        value
+    }
+
+    #[instrument(level = "trace", skip(self))]
+    fn discard_place(&mut self, place: Place<'tcx>) {
+        let discard_local = |this: &mut Self, local| {
+            if this.ssa.is_ssa(local) {
+                return;
+            }
+            if let Some(value) = this.locals[local].take() {
+                this.rev_locals_non_ssa.entry(value).or_default().retain(|l| *l != local);
+            }
+        };
+        if place.is_indirect_first_projection() {
+            // Non-local mutation maybe invalidate deref.
+            self.invalidate_derefs();
+            // Remove stored value from borrowed locals.
+            for local in self.ssa.borrowed_locals().iter() {
+                discard_local(self, local);
+            }
+        } else {
+            discard_local(self, place.local);
+        }
     }
 
     fn insert_bool(&mut self, flag: bool) -> VnIndex {
@@ -759,7 +810,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         let (mut place_ty, mut value) = match base {
             // The base is a local, so we take the local's value and project from it.
             AddressBase::Local(local) => {
-                let local = self.locals[local]?;
+                let local = self.local(local);
                 let place_ty = PlaceTy::from_ty(self.ty(local));
                 (place_ty, local)
             }
@@ -865,7 +916,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         // If the projection is indirect, we treat the local as a value, so can replace it with
         // another local.
         if place.is_indirect_first_projection()
-            && let Some(base) = self.locals[place.local]
+            && let base = self.local(place.local)
             && let Some(new_local) = self.try_as_local(base, location)
             && place.local != new_local
         {
@@ -877,9 +928,8 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
 
         for i in 0..projection.len() {
             let elem = projection[i];
-            if let ProjectionElem::Index(idx_local) = elem
-                && let Some(idx) = self.locals[idx_local]
-            {
+            if let ProjectionElem::Index(idx_local) = elem {
+                let idx = self.local(idx_local);
                 if let Some(offset) = self.evaluated[idx].as_ref()
                     && let Some(offset) = self.ecx.read_target_usize(offset).discard_err()
                     && let Some(min_length) = offset.checked_add(1)
@@ -915,7 +965,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         let mut place_ref = place.as_ref();
 
         // Invariant: `value` holds the value up-to the `index`th projection excluded.
-        let Some(mut value) = self.locals[place.local] else { return Err(place_ref) };
+        let mut value = self.local(place.local);
         // Invariant: `value` has type `place_ty`, with optional downcast variant if needed.
         let mut place_ty = PlaceTy::from_ty(self.local_decls[place.local].ty);
         for (index, proj) in place.projection.iter().enumerate() {
@@ -926,7 +976,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
                 place_ref = PlaceRef { local, projection: &place.projection[index..] };
             }
 
-            let Some(proj) = proj.try_map(|value| self.locals[value], |ty| ty) else {
+            let Some(proj) = proj.try_map(|value| Some(self.local(value)), |ty| ty) else {
                 return Err(place_ref);
             };
             let Some(ty_and_value) = self.project(place_ty, value, proj) else {
@@ -1809,11 +1859,17 @@ impl<'tcx> VnState<'_, '_, 'tcx> {
     /// If there is a local which is assigned `index`, and its assignment strictly dominates `loc`,
     /// return it. If you used this local, add it to `reused_locals` to remove storage statements.
     fn try_as_local(&mut self, index: VnIndex, loc: Location) -> Option<Local> {
-        let other = self.rev_locals.get(index)?;
-        other
-            .iter()
-            .find(|&&other| self.ssa.assignment_dominates(&self.dominators, other, loc))
-            .copied()
+        if let Some(ssa) = self.rev_locals_ssa.get(index)
+            && let Some(other) = ssa
+                .iter()
+                .find(|&&other| self.ssa.assignment_dominates(&self.dominators, other, loc))
+        {
+            Some(*other)
+        } else if let Some(non_ssa) = self.rev_locals_non_ssa.get(&index) {
+            non_ssa.first().copied()
+        } else {
+            None
+        }
     }
 }
 
@@ -1822,11 +1878,20 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
         self.tcx
     }
 
+    fn visit_basic_block_data(&mut self, block: BasicBlock, bbdata: &mut BasicBlockData<'tcx>) {
+        self.rev_locals_non_ssa.clear();
+        for local in self.locals.indices() {
+            if !self.ssa.is_ssa(local) {
+                self.locals[local] = None;
+            }
+        }
+        self.super_basic_block_data(block, bbdata);
+    }
+
     fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
         self.simplify_place_projection(place, location);
-        if context.is_mutating_use() && place.is_indirect() {
-            // Non-local mutation maybe invalidate deref.
-            self.invalidate_derefs();
+        if context.is_mutating_use() {
+            self.discard_place(*place);
         }
         self.super_place(place, context, location);
     }
@@ -1865,13 +1930,9 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
             }
         }
 
-        if lhs.is_indirect() {
-            // Non-local mutation maybe invalidate deref.
-            self.invalidate_derefs();
-        }
+        self.discard_place(*lhs);
 
         if let Some(local) = lhs.as_local()
-            && self.ssa.is_ssa(local)
             && let rvalue_ty = rvalue.ty(self.local_decls, self.tcx)
             // FIXME(#112651) `rvalue` may have a subtype to `local`. We can only mark
             // `local` as reusable if we have an exact type match.
@@ -1883,14 +1944,13 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
     }
 
     fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
-        if let Terminator { kind: TerminatorKind::Call { destination, .. }, .. } = terminator {
-            if let Some(local) = destination.as_local()
-                && self.ssa.is_ssa(local)
-            {
-                let ty = self.local_decls[local].ty;
-                let opaque = self.new_opaque(ty);
-                self.assign(local, opaque);
-            }
+        self.super_terminator(terminator, location);
+        if let Terminator { kind: TerminatorKind::Call { destination, .. }, .. } = terminator
+            && let Some(local) = destination.as_local()
+        {
+            let ty = self.local_decls[local].ty;
+            let opaque = self.new_opaque(ty);
+            self.assign(local, opaque);
         }
         // Function calls and ASM may invalidate (nested) derefs. We must handle them carefully.
         // Currently, only preserving derefs for trivial terminators like SwitchInt and Goto.
@@ -1901,7 +1961,6 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
         if !safe_to_preserve_derefs {
             self.invalidate_derefs();
         }
-        self.super_terminator(terminator, location);
     }
 }