Extend GVN to perform local value numbering.

Author: cjgillot

compiler/rustc_mir_transform/src/gvn.rs

@@ -131,11 +131,6 @@ impl<'tcx> crate::MirPass<'tcx> for GVN {
         let mut state =
             VnState::new(tcx, body, typing_env, &ssa, dominators, &body.local_decls, &arena);
 
-        for local in body.args_iter().filter(|&local| ssa.is_ssa(local)) {
-            let opaque = state.new_opaque(body.local_decls[local].ty);
-            state.assign(local, opaque);
-        }
-
         let reverse_postorder = body.basic_blocks.reverse_postorder().to_vec();
         for bb in reverse_postorder {
             let data = &mut body.basic_blocks.as_mut_preserves_cfg()[bb];
@@ -232,10 +227,10 @@ struct VnState<'body, 'a, 'tcx> {
     local_decls: &'body LocalDecls<'tcx>,
     is_coroutine: bool,
     /// Value stored in each local.
-    locals: IndexVec<Local, Option<VnIndex>>,
+    locals: IndexVec<Local, 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]>>,
+    rev_locals: IndexVec<VnIndex, SmallVec<[(Local, Location); 1]>>,
     values: FxIndexSet<(Value<'a, 'tcx>, Ty<'tcx>)>,
     /// Values evaluated as constants if possible.
     evaluated: IndexVec<VnIndex, Option<OpTy<'tcx>>>,
@@ -266,12 +261,12 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         let num_values =
             2 * body.basic_blocks.iter().map(|bbdata| bbdata.statements.len()).sum::<usize>()
                 + 4 * body.basic_blocks.len();
-        VnState {
+        let mut this = VnState {
             tcx,
             ecx: InterpCx::new(tcx, DUMMY_SP, typing_env, DummyMachine),
             local_decls,
             is_coroutine: body.coroutine.is_some(),
-            locals: IndexVec::from_elem(None, local_decls),
+            locals: IndexVec::with_capacity(body.local_decls.len()),
             rev_locals: IndexVec::with_capacity(num_values),
             values: FxIndexSet::with_capacity_and_hasher(num_values, Default::default()),
             evaluated: IndexVec::with_capacity(num_values),
@@ -281,7 +276,14 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             dominators,
             reused_locals: DenseBitSet::new_empty(local_decls.len()),
             arena,
+        };
+        let init_loc = Location { block: START_BLOCK, statement_index: 0 };
+        for decl in body.local_decls.iter() {
+            let value = this.new_opaque(decl.ty);
+            let local = this.locals.push(value);
+            this.rev_locals[value].push((local, init_loc));
         }
+        this
     }
 
     fn typing_env(&self) -> ty::TypingEnv<'tcx> {
@@ -335,7 +337,7 @@ 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.locals[place.local];
             // Skip the initial `Deref`.
             projection.next();
             AddressBase::Deref(base)
@@ -346,7 +348,8 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         let projection = self
             .arena
             .try_alloc_from_iter(
-                projection.map(|proj| proj.try_map(|value| self.locals[value], |ty| ty).ok_or(())),
+                projection
+                    .map(|proj| proj.try_map(|value| Some(self.locals[value]), |ty| ty).ok_or(())),
             )
             .ok()?;
         let value = Value::Address { base, projection, kind, provenance: self.next_opaque() };
@@ -363,12 +366,39 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         self.values.get_index(index.as_usize()).unwrap().1
     }
 
-    /// 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);
+    fn assign(&mut self, local: Local, value: VnIndex, loc: Location) {
+        self.locals[local] = value;
+        self.rev_locals[value].push((local, loc));
+    }
+
+    #[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;
+            }
+            let value = this.locals[local];
+            this.rev_locals[value].retain(|(l, _)| *l != local);
+            this.locals[local] = this.new_opaque(this.ty(value));
+            #[cfg(debug_assertions)]
+            for local_vec in this.rev_locals.iter() {
+                for (other, _) in local_vec {
+                    debug_assert_ne!(*other, 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_constant(&mut self, value: Const<'tcx>) -> VnIndex {
@@ -633,7 +663,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.locals[local];
                 let place_ty = PlaceTy::from_ty(self.ty(local));
                 (place_ty, local)
             }
@@ -744,7 +774,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.locals[place.local]
             && let Some(new_local) = self.try_as_local(base, location)
             && place.local != new_local
         {
@@ -756,9 +786,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.locals[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)
@@ -794,7 +823,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.locals[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() {
@@ -805,7 +834,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.locals[value]), |ty| ty) else {
                 return Err(place_ref);
             };
             let Some(ty_and_value) = self.project(place_ty, value, proj) else {
@@ -1017,6 +1046,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         None
     }
 
+    #[instrument(level = "trace", skip(self), ret)]
     fn simplify_aggregate(
         &mut self,
         lhs: &Place<'tcx>,
@@ -1661,6 +1691,7 @@ fn op_to_prop_const<'tcx>(
 impl<'tcx> VnState<'_, '_, 'tcx> {
     /// If either [`Self::try_as_constant`] as [`Self::try_as_place`] succeeds,
     /// returns that result as an [`Operand`].
+    #[instrument(level = "trace", skip(self), ret)]
     fn try_as_operand(&mut self, index: VnIndex, location: Location) -> Option<Operand<'tcx>> {
         if let Some(const_) = self.try_as_constant(index) {
             Some(Operand::Constant(Box::new(const_)))
@@ -1673,6 +1704,7 @@ impl<'tcx> VnState<'_, '_, 'tcx> {
     }
 
     /// If `index` is a `Value::Constant`, return the `Constant` to be put in the MIR.
+    #[instrument(level = "trace", skip(self), ret)]
     fn try_as_constant(&mut self, index: VnIndex) -> Option<ConstOperand<'tcx>> {
         // This was already constant in MIR, do not change it. If the constant is not
         // deterministic, adding an additional mention of it in MIR will not give the same value as
@@ -1730,12 +1762,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.
+    #[instrument(level = "trace", skip(self), ret)]
     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()
+            .find(|&&(other, assign_loc)| {
+                self.ssa.assignment_dominates(&self.dominators, other, loc)
+                    || (assign_loc.block == loc.block
+                        && assign_loc.statement_index < loc.statement_index)
+            })
+            .map(|&(other, _)| other)
     }
 }
 
@@ -1744,11 +1781,25 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
         self.tcx
     }
 
+    fn visit_basic_block_data(&mut self, block: BasicBlock, bbdata: &mut BasicBlockData<'tcx>) {
+        for local_set in self.rev_locals.iter_mut() {
+            local_set.retain(|(local, _)| self.ssa.is_ssa(*local));
+        }
+        for local in self.locals.indices() {
+            if !self.ssa.is_ssa(local) {
+                let current = self.locals[local];
+                let new = self.new_opaque(self.ty(current));
+                self.locals[local] = new;
+                self.rev_locals[new].push((local, Location { block, statement_index: 0 }));
+            }
+        }
+        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);
     }
@@ -1779,32 +1830,27 @@ 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.
             && self.local_decls[local].ty == rvalue_ty
         {
             let value = value.unwrap_or_else(|| self.new_opaque(rvalue_ty));
-            self.assign(local, value);
+            self.assign(local, value, location);
         }
     }
 
     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, location);
         }
         // Function calls and ASM may invalidate (nested) derefs. We must handle them carefully.
         // Currently, only preserving derefs for trivial terminators like SwitchInt and Goto.
@@ -1815,7 +1861,6 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
         if !safe_to_preserve_derefs {
             self.invalidate_derefs();
         }
-        self.super_terminator(terminator, location);
     }
 }
 

tests/mir-opt/const_prop/boxes.main.GVN.panic-abort.diff

@@ -30,16 +30,21 @@
       }
 
       bb1: {
-          StorageLive(_7);
+-         StorageLive(_7);
++         nop;
           _7 = ShallowInitBox(move _6, i32);
           _8 = copy ((_7.0: std::ptr::Unique<i32>).0: std::ptr::NonNull<i32>) as *const i32 (Transmute);
           (*_8) = const 42_i32;
-          _3 = move _7;
-          StorageDead(_7);
-          _9 = copy ((_3.0: std::ptr::Unique<i32>).0: std::ptr::NonNull<i32>) as *const i32 (Transmute);
-          _2 = copy (*_9);
+-         _3 = move _7;
+-         StorageDead(_7);
+-         _9 = copy ((_3.0: std::ptr::Unique<i32>).0: std::ptr::NonNull<i32>) as *const i32 (Transmute);
+-         _2 = copy (*_9);
 -         _1 = Add(move _2, const 0_i32);
 -         StorageDead(_2);
++         _3 = copy _7;
++         nop;
++         _9 = copy _8;
++         _2 = copy (*_8);
 +         _1 = copy _2;
 +         nop;
           drop(_3) -> [return: bb2, unwind unreachable];

tests/mir-opt/const_prop/boxes.main.GVN.panic-unwind.diff

@@ -30,16 +30,21 @@
       }
 
       bb1: {
-          StorageLive(_7);
+-         StorageLive(_7);
++         nop;
           _7 = ShallowInitBox(move _6, i32);
           _8 = copy ((_7.0: std::ptr::Unique<i32>).0: std::ptr::NonNull<i32>) as *const i32 (Transmute);
           (*_8) = const 42_i32;
-          _3 = move _7;
-          StorageDead(_7);
-          _9 = copy ((_3.0: std::ptr::Unique<i32>).0: std::ptr::NonNull<i32>) as *const i32 (Transmute);
-          _2 = copy (*_9);
+-         _3 = move _7;
+-         StorageDead(_7);
+-         _9 = copy ((_3.0: std::ptr::Unique<i32>).0: std::ptr::NonNull<i32>) as *const i32 (Transmute);
+-         _2 = copy (*_9);
 -         _1 = Add(move _2, const 0_i32);
 -         StorageDead(_2);
++         _3 = copy _7;
++         nop;
++         _9 = copy _8;
++         _2 = copy (*_8);
 +         _1 = copy _2;
 +         nop;
           drop(_3) -> [return: bb2, unwind: bb3];

tests/mir-opt/const_prop/mutable_variable.main.GVN.diff

@@ -17,7 +17,8 @@
           _1 = const 42_i32;
           _1 = const 99_i32;
           StorageLive(_2);
-          _2 = copy _1;
+-         _2 = copy _1;
++         _2 = const 99_i32;
           _0 = const ();
           StorageDead(_2);
           StorageDead(_1);

tests/mir-opt/const_prop/mutable_variable.rs

@@ -7,7 +7,7 @@ fn main() {
     // CHECK: debug y => [[y:_.*]];
     // CHECK: [[x]] = const 42_i32;
     // CHECK: [[x]] = const 99_i32;
-    // CHECK: [[y]] = copy [[x]];
+    // CHECK: [[y]] = const 99_i32;
     let mut x = 42;
     x = 99;
     let y = x;

tests/mir-opt/const_prop/mutable_variable_no_prop.main.GVN.diff

@@ -19,17 +19,21 @@
           StorageLive(_1);
           _1 = const 42_u32;
           StorageLive(_2);
-          StorageLive(_3);
+-         StorageLive(_3);
++         nop;
           StorageLive(_4);
           _4 = const {ALLOC0: *mut u32};
           _3 = copy (*_4);
-          _1 = move _3;
-          StorageDead(_3);
+-         _1 = move _3;
+-         StorageDead(_3);
++         _1 = copy _3;
++         nop;
           StorageDead(_4);
           _2 = const ();
           StorageDead(_2);
           StorageLive(_5);
-          _5 = copy _1;
+-         _5 = copy _1;
++         _5 = copy _3;
           _0 = const ();
           StorageDead(_5);
           StorageDead(_1);

tests/mir-opt/const_prop/mutable_variable_no_prop.rs

@@ -10,8 +10,8 @@ fn main() {
     // CHECK: debug y => [[y:_.*]];
     // CHECK: [[x]] = const 42_u32;
     // CHECK: [[tmp:_.*]] = copy (*{{_.*}});
-    // CHECK: [[x]] = move [[tmp]];
-    // CHECK: [[y]] = copy [[x]];
+    // CHECK: [[x]] = copy [[tmp]];
+    // CHECK: [[y]] = copy [[tmp]];
     let mut x = 42;
     unsafe {
         x = STATIC;