Do not use for_each_assignment_mut to iterate over assignment statements

`for_each_assignment_mut` can skip assignment statements with side effects,
which can result in some assignment statements retrieving outdated value.
For example, it may skip a dereference assignment statement.

Author: dianqk

compiler/rustc_mir_transform/src/gvn.rs

@@ -3,14 +3,16 @@
 //! 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.
 //!
-//! In a first pass, 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.
-//!
 //! From those assignments, we construct a mapping `VnIndex -> Vec<(Local, Location)>` of available
 //! values, the locals in which they are stored, and the assignment location.
 //!
-//! In a second pass, we traverse all (non SSA) assignments `x = rvalue` and operands. For each
+//! We traverse all assignments `x = rvalue` and operands.
+//!
+//! For each SSA one, 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
@@ -107,7 +109,7 @@ use rustc_span::def_id::DefId;
 use smallvec::SmallVec;
 use tracing::{debug, instrument, trace};
 
-use crate::ssa::{AssignedValue, SsaLocals};
+use crate::ssa::SsaLocals;
 
 pub(super) struct GVN;
 
@@ -126,31 +128,11 @@ impl<'tcx> crate::MirPass<'tcx> for GVN {
         let dominators = body.basic_blocks.dominators().clone();
 
         let mut state = VnState::new(tcx, body, typing_env, &ssa, dominators, &body.local_decls);
-        ssa.for_each_assignment_mut(
-            body.basic_blocks.as_mut_preserves_cfg(),
-            |local, value, location| {
-                let value = match value {
-                    // We do not know anything of this assigned value.
-                    AssignedValue::Arg | AssignedValue::Terminator => None,
-                    // Try to get some insight.
-                    AssignedValue::Rvalue(rvalue) => {
-                        let value = state.simplify_rvalue(rvalue, location);
-                        // FIXME(#112651) `rvalue` may have a subtype to `local`. We can only mark
-                        // `local` as reusable if we have an exact type match.
-                        if state.local_decls[local].ty != rvalue.ty(state.local_decls, tcx) {
-                            return;
-                        }
-                        value
-                    }
-                };
-                // `next_opaque` is `Some`, so `new_opaque` must return `Some`.
-                let value = value.or_else(|| state.new_opaque()).unwrap();
-                state.assign(local, value);
-            },
-        );
 
-        // Stop creating opaques during replacement as it is useless.
-        state.next_opaque = None;
+        for local in body.args_iter().filter(|&local| ssa.is_ssa(local)) {
+            let opaque = state.new_opaque().unwrap();
+            state.assign(local, opaque);
+        }
 
         let reverse_postorder = body.basic_blocks.reverse_postorder().to_vec();
         for bb in reverse_postorder {
@@ -246,7 +228,6 @@ struct VnState<'body, 'tcx> {
     locals: IndexVec<Local, Option<VnIndex>>,
     /// Locals that are assigned that value.
     // This vector does not hold all the values of `VnIndex` that we create.
-    // It stops at the largest value created in the first phase of collecting assignments.
     rev_locals: IndexVec<VnIndex, SmallVec<[Local; 1]>>,
     values: FxIndexSet<Value<'tcx>>,
     /// Values evaluated as constants if possible.
@@ -341,6 +322,7 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
     /// Record that `local` is assigned `value`. `local` must be SSA.
     #[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);
 
         // Only register the value if its type is `Sized`, as we will emit copies of it.
@@ -1638,15 +1620,19 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, 'tcx> {
         if let StatementKind::Assign(box (ref mut lhs, ref mut rvalue)) = stmt.kind {
             self.simplify_place_projection(lhs, location);
 
-            // Do not try to simplify a constant, it's already in canonical shape.
-            if matches!(rvalue, Rvalue::Use(Operand::Constant(_))) {
-                return;
-            }
-
-            let value = lhs
-                .as_local()
-                .and_then(|local| self.locals[local])
-                .or_else(|| self.simplify_rvalue(rvalue, location));
+            let value = self.simplify_rvalue(rvalue, location);
+            let value = if let Some(local) = lhs.as_local()
+                && self.ssa.is_ssa(local)
+                // 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(self.local_decls, self.tcx)
+            {
+                let value = value.or_else(|| self.new_opaque()).unwrap();
+                self.assign(local, value);
+                Some(value)
+            } else {
+                value
+            };
             let Some(value) = value else { return };
 
             if let Some(const_) = self.try_as_constant(value) {
@@ -1662,6 +1648,17 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, 'tcx> {
         }
         self.super_statement(stmt, location);
     }
+
+    fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
+        if let Terminator { kind: TerminatorKind::Call { destination, .. }, .. } = terminator
+            && let Some(local) = destination.as_local()
+            && self.ssa.is_ssa(local)
+        {
+            let opaque = self.new_opaque().unwrap();
+            self.assign(local, opaque);
+        }
+        self.super_terminator(terminator, location);
+    }
 }
 
 struct StorageRemover<'tcx> {

compiler/rustc_mir_transform/src/ssa.rs

@@ -32,12 +32,6 @@ pub(super) struct SsaLocals {
     borrowed_locals: BitSet<Local>,
 }
 
-pub(super) enum AssignedValue<'a, 'tcx> {
-    Arg,
-    Rvalue(&'a mut Rvalue<'tcx>),
-    Terminator,
-}
-
 impl SsaLocals {
     pub(super) fn new<'tcx>(
         tcx: TyCtxt<'tcx>,
@@ -152,37 +146,6 @@ impl SsaLocals {
         })
     }
 
-    pub(super) fn for_each_assignment_mut<'tcx>(
-        &self,
-        basic_blocks: &mut IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
-        mut f: impl FnMut(Local, AssignedValue<'_, 'tcx>, Location),
-    ) {
-        for &local in &self.assignment_order {
-            match self.assignments[local] {
-                Set1::One(DefLocation::Argument) => f(local, AssignedValue::Arg, Location {
-                    block: START_BLOCK,
-                    statement_index: 0,
-                }),
-                Set1::One(DefLocation::Assignment(loc)) => {
-                    let bb = &mut basic_blocks[loc.block];
-                    // `loc` must point to a direct assignment to `local`.
-                    let stmt = &mut bb.statements[loc.statement_index];
-                    let StatementKind::Assign(box (target, ref mut rvalue)) = stmt.kind else {
-                        bug!()
-                    };
-                    assert_eq!(target.as_local(), Some(local));
-                    f(local, AssignedValue::Rvalue(rvalue), loc)
-                }
-                Set1::One(DefLocation::CallReturn { call, .. }) => {
-                    let bb = &mut basic_blocks[call];
-                    let loc = Location { block: call, statement_index: bb.statements.len() };
-                    f(local, AssignedValue::Terminator, loc)
-                }
-                _ => {}
-            }
-        }
-    }
-
     /// Compute the equivalence classes for locals, based on copy statements.
     ///
     /// The returned vector maps each local to the one it copies. In the following case:

tests/mir-opt/const_prop/control_flow_simplification.hello.GVN.panic-abort.diff

@@ -8,8 +8,9 @@
 
       bb0: {
           StorageLive(_1);
-          _1 = const <bool as NeedsDrop>::NEEDS;
+-         _1 = const <bool as NeedsDrop>::NEEDS;
 -         switchInt(move _1) -> [0: bb2, otherwise: bb1];
++         _1 = const false;
 +         switchInt(const false) -> [0: bb2, otherwise: bb1];
       }
 

tests/mir-opt/const_prop/control_flow_simplification.hello.GVN.panic-unwind.diff

@@ -8,8 +8,9 @@
 
       bb0: {
           StorageLive(_1);
-          _1 = const <bool as NeedsDrop>::NEEDS;
+-         _1 = const <bool as NeedsDrop>::NEEDS;
 -         switchInt(move _1) -> [0: bb2, otherwise: bb1];
++         _1 = const false;
 +         switchInt(const false) -> [0: bb2, otherwise: bb1];
       }