Separate checking rvalue from evaluation.

2 files changed, 104 insertions, 113 deletions

diff --git a/compiler/rustc_mir_transform/src/const_prop.rs b/compiler/rustc_mir_transform/src/const_prop.rs
index fa8c22ceb40..6e590981357 100644
--- a/compiler/rustc_mir_transform/src/const_prop.rs
+++ b/compiler/rustc_mir_transform/src/const_prop.rs
@@ -470,7 +470,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         }
     }
 
-    fn const_prop(&mut self, rvalue: &Rvalue<'tcx>, place: Place<'tcx>) -> Option<()> {
+    fn check_rvalue(&mut self, rvalue: &Rvalue<'tcx>) -> Option<()> {
         // Perform any special handling for specific Rvalue types.
         // Generally, checks here fall into one of two categories:
         //   1. Additional checking to provide useful lints to the user
@@ -527,7 +527,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
             return None;
         }
 
-        self.eval_rvalue_with_identities(rvalue, place)
+        Some(())
     }
 
     // Attempt to use algebraic identities to eliminate constant expressions
@@ -595,7 +595,16 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         }))
     }
 
-    fn replace_with_const(&mut self, rval: &mut Rvalue<'tcx>, value: &OpTy<'tcx>) {
+    fn replace_with_const(&mut self, place: Place<'tcx>, rval: &mut Rvalue<'tcx>) {
+        // This will return None if the above `const_prop` invocation only "wrote" a
+        // type whose creation requires no write. E.g. a generator whose initial state
+        // consists solely of uninitialized memory (so it doesn't capture any locals).
+        let Some(ref value) = self.get_const(place) else { return };
+        if !self.should_const_prop(value) {
+            return;
+        }
+        trace!("replacing {:?}={:?} with {:?}", place, rval, value);
+
         if let Rvalue::Use(Operand::Constant(c)) = rval {
             match c.literal {
                 ConstantKind::Ty(c) if matches!(c.kind(), ConstKind::Unevaluated(..)) => {}
@@ -688,6 +697,19 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
             _ => false,
         }
     }
+
+    fn ensure_not_propagated(&mut self, local: Local) {
+        if cfg!(debug_assertions) {
+            assert!(
+                self.get_const(local.into()).is_none()
+                    || self
+                        .layout_of(self.local_decls[local].ty)
+                        .map_or(true, |layout| layout.is_zst()),
+                "failed to remove values for `{local:?}`, value={:?}",
+                self.get_const(local.into()),
+            )
+        }
+    }
 }
 
 /// The mode that `ConstProp` is allowed to run in for a given `Local`.
@@ -827,44 +849,23 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
         self.eval_constant(constant);
     }
 
-    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
-        trace!("visit_statement: {:?}", statement);
-
-        // Recurse into statement before applying the assignment.
-        self.super_statement(statement, location);
-
-        match statement.kind {
-            StatementKind::Assign(box (place, ref mut rval)) => {
-                let can_const_prop = self.ecx.machine.can_const_prop[place.local];
-                if let Some(()) = self.const_prop(rval, place) {
-                    // This will return None if the above `const_prop` invocation only "wrote" a
-                    // type whose creation requires no write. E.g. a generator whose initial state
-                    // consists solely of uninitialized memory (so it doesn't capture any locals).
-                    if let Some(ref value) = self.get_const(place) && self.should_const_prop(value) {
-                        trace!("replacing {:?} with {:?}", rval, value);
-                        self.replace_with_const(rval, value);
-                        if can_const_prop == ConstPropMode::FullConstProp
-                            || can_const_prop == ConstPropMode::OnlyInsideOwnBlock
-                        {
-                            trace!("propagated into {:?}", place);
-                        }
-                    }
-                    match can_const_prop {
-                        ConstPropMode::OnlyInsideOwnBlock => {
-                            trace!(
-                                "found local restricted to its block. \
-                                Will remove it from const-prop after block is finished. Local: {:?}",
-                                place.local
-                            );
-                        }
-                        ConstPropMode::NoPropagation => {
-                            trace!("can't propagate into {:?}", place);
-                            if place.local != RETURN_PLACE {
-                                Self::remove_const(&mut self.ecx, place.local);
-                            }
-                        }
-                        ConstPropMode::FullConstProp => {}
-                    }
+    fn visit_assign(
+        &mut self,
+        place: &mut Place<'tcx>,
+        rvalue: &mut Rvalue<'tcx>,
+        location: Location,
+    ) {
+        self.super_assign(place, rvalue, location);
+
+        let Some(()) = self.check_rvalue(rvalue) else { return };
+
+        match self.ecx.machine.can_const_prop[place.local] {
+            // Do nothing if the place is indirect.
+            _ if place.is_indirect() => {}
+            ConstPropMode::NoPropagation => self.ensure_not_propagated(place.local),
+            ConstPropMode::OnlyInsideOwnBlock | ConstPropMode::FullConstProp => {
+                if let Some(()) = self.eval_rvalue_with_identities(rvalue, *place) {
+                    self.replace_with_const(*place, rvalue);
                 } else {
                     // Const prop failed, so erase the destination, ensuring that whatever happens
                     // from here on, does not know about the previous value.
@@ -884,8 +885,21 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
                     Self::remove_const(&mut self.ecx, place.local);
                 }
             }
+        }
+    }
+
+    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
+        trace!("visit_statement: {:?}", statement);
+
+        // Recurse into statement before applying the assignment.
+        self.super_statement(statement, location);
+
+        match statement.kind {
             StatementKind::SetDiscriminant { ref place, .. } => {
                 match self.ecx.machine.can_const_prop[place.local] {
+                    // Do nothing if the place is indirect.
+                    _ if place.is_indirect() => {}
+                    ConstPropMode::NoPropagation => self.ensure_not_propagated(place.local),
                     ConstPropMode::FullConstProp | ConstPropMode::OnlyInsideOwnBlock => {
                         if self.ecx.statement(statement).is_ok() {
                             trace!("propped discriminant into {:?}", place);
@@ -893,9 +907,6 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
                             Self::remove_const(&mut self.ecx, place.local);
                         }
                     }
-                    ConstPropMode::NoPropagation => {
-                        Self::remove_const(&mut self.ecx, place.local);
-                    }
                 }
             }
             StatementKind::StorageLive(local) => {
@@ -958,19 +969,6 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
     fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
         self.super_basic_block_data(block, data);
 
-        let ensure_not_propagated = |this: &mut Self, local: Local| {
-            if cfg!(debug_assertions) {
-                assert!(
-                    this.get_const(local.into()).is_none()
-                        || this
-                            .layout_of(this.local_decls[local].ty)
-                            .map_or(true, |layout| layout.is_zst()),
-                    "failed to remove values for `{local:?}`, value={:?}",
-                    this.get_const(local.into()),
-                )
-            }
-        };
-
         // We remove all Locals which are restricted in propagation to their containing blocks and
         // which were modified in the current block.
         // Take it out of the ecx so we can get a mutable reference to the ecx for `remove_const`.
@@ -978,10 +976,10 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
         for (local, &mode) in can_const_prop.iter_enumerated() {
             match mode {
                 ConstPropMode::FullConstProp => {}
-                ConstPropMode::NoPropagation => ensure_not_propagated(self, local),
+                ConstPropMode::NoPropagation => self.ensure_not_propagated(local),
                 ConstPropMode::OnlyInsideOwnBlock => {
                     Self::remove_const(&mut self.ecx, local);
-                    ensure_not_propagated(self, local);
+                    self.ensure_not_propagated(local);
                 }
             }
         }
diff --git a/compiler/rustc_mir_transform/src/const_prop_lint.rs b/compiler/rustc_mir_transform/src/const_prop_lint.rs
index b7835f842ba..28198781fd9 100644
--- a/compiler/rustc_mir_transform/src/const_prop_lint.rs
+++ b/compiler/rustc_mir_transform/src/const_prop_lint.rs
@@ -405,12 +405,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         Some(())
     }
 
-    fn const_prop(
-        &mut self,
-        rvalue: &Rvalue<'tcx>,
-        source_info: SourceInfo,
-        place: Place<'tcx>,
-    ) -> Option<()> {
+    fn check_rvalue(&mut self, rvalue: &Rvalue<'tcx>, source_info: SourceInfo) -> Option<()> {
         // Perform any special handling for specific Rvalue types.
         // Generally, checks here fall into one of two categories:
         //   1. Additional checking to provide useful lints to the user
@@ -486,7 +481,20 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
             return None;
         }
 
-        self.use_ecx(source_info, |this| this.ecx.eval_rvalue_into_place(rvalue, place))
+        Some(())
+    }
+
+    fn ensure_not_propagated(&mut self, local: Local) {
+        if cfg!(debug_assertions) {
+            assert!(
+                self.get_const(local.into()).is_none()
+                    || self
+                        .layout_of(self.local_decls[local].ty)
+                        .map_or(true, |layout| layout.is_zst()),
+                "failed to remove values for `{local:?}`, value={:?}",
+                self.get_const(local.into()),
+            )
+        }
     }
 }
 
@@ -507,35 +515,21 @@ impl<'tcx> Visitor<'tcx> for ConstPropagator<'_, 'tcx> {
         self.eval_constant(constant, self.source_info.unwrap());
     }
 
-    fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
-        trace!("visit_statement: {:?}", statement);
-        let source_info = statement.source_info;
-        self.source_info = Some(source_info);
-
-        // Recurse into statement before applying the assignment.
-        self.super_statement(statement, location);
-
-        match statement.kind {
-            StatementKind::Assign(box (place, ref rval)) => {
-                let can_const_prop = self.ecx.machine.can_const_prop[place.local];
-                if let Some(()) = self.const_prop(rval, source_info, place) {
-                    match can_const_prop {
-                        ConstPropMode::OnlyInsideOwnBlock => {
-                            trace!(
-                                "found local restricted to its block. \
-                                Will remove it from const-prop after block is finished. Local: {:?}",
-                                place.local
-                            );
-                        }
-                        ConstPropMode::NoPropagation => {
-                            trace!("can't propagate into {:?}", place);
-                            if place.local != RETURN_PLACE {
-                                Self::remove_const(&mut self.ecx, place.local);
-                            }
-                        }
-                        ConstPropMode::FullConstProp => {}
-                    }
-                } else {
+    fn visit_assign(&mut self, place: &Place<'tcx>, rvalue: &Rvalue<'tcx>, location: Location) {
+        self.super_assign(place, rvalue, location);
+
+        let source_info = self.source_info.unwrap();
+        let Some(()) = self.check_rvalue(rvalue, source_info) else { return };
+
+        match self.ecx.machine.can_const_prop[place.local] {
+            // Do nothing if the place is indirect.
+            _ if place.is_indirect() => {}
+            ConstPropMode::NoPropagation => self.ensure_not_propagated(place.local),
+            ConstPropMode::OnlyInsideOwnBlock | ConstPropMode::FullConstProp => {
+                if self
+                    .use_ecx(source_info, |this| this.ecx.eval_rvalue_into_place(rvalue, *place))
+                    .is_none()
+                {
                     // Const prop failed, so erase the destination, ensuring that whatever happens
                     // from here on, does not know about the previous value.
                     // This is important in case we have
@@ -554,8 +548,23 @@ impl<'tcx> Visitor<'tcx> for ConstPropagator<'_, 'tcx> {
                     Self::remove_const(&mut self.ecx, place.local);
                 }
             }
+        }
+    }
+
+    fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
+        trace!("visit_statement: {:?}", statement);
+        let source_info = statement.source_info;
+        self.source_info = Some(source_info);
+
+        // Recurse into statement before applying the assignment.
+        self.super_statement(statement, location);
+
+        match statement.kind {
             StatementKind::SetDiscriminant { ref place, .. } => {
                 match self.ecx.machine.can_const_prop[place.local] {
+                    // Do nothing if the place is indirect.
+                    _ if place.is_indirect() => {}
+                    ConstPropMode::NoPropagation => self.ensure_not_propagated(place.local),
                     ConstPropMode::FullConstProp | ConstPropMode::OnlyInsideOwnBlock => {
                         if self.use_ecx(source_info, |this| this.ecx.statement(statement)).is_some()
                         {
@@ -564,9 +573,6 @@ impl<'tcx> Visitor<'tcx> for ConstPropagator<'_, 'tcx> {
                             Self::remove_const(&mut self.ecx, place.local);
                         }
                     }
-                    ConstPropMode::NoPropagation => {
-                        Self::remove_const(&mut self.ecx, place.local);
-                    }
                 }
             }
             StatementKind::StorageLive(local) => {
@@ -682,19 +688,6 @@ impl<'tcx> Visitor<'tcx> for ConstPropagator<'_, 'tcx> {
     fn visit_basic_block_data(&mut self, block: BasicBlock, data: &BasicBlockData<'tcx>) {
         self.super_basic_block_data(block, data);
 
-        let ensure_not_propagated = |this: &mut Self, local: Local| {
-            if cfg!(debug_assertions) {
-                assert!(
-                    this.get_const(local.into()).is_none()
-                        || this
-                            .layout_of(this.local_decls[local].ty)
-                            .map_or(true, |layout| layout.is_zst()),
-                    "failed to remove values for `{local:?}`, value={:?}",
-                    this.get_const(local.into()),
-                )
-            }
-        };
-
         // We remove all Locals which are restricted in propagation to their containing blocks and
         // which were modified in the current block.
         // Take it out of the ecx so we can get a mutable reference to the ecx for `remove_const`.
@@ -702,10 +695,10 @@ impl<'tcx> Visitor<'tcx> for ConstPropagator<'_, 'tcx> {
         for (local, &mode) in can_const_prop.iter_enumerated() {
             match mode {
                 ConstPropMode::FullConstProp => {}
-                ConstPropMode::NoPropagation => ensure_not_propagated(self, local),
+                ConstPropMode::NoPropagation => self.ensure_not_propagated(local),
                 ConstPropMode::OnlyInsideOwnBlock => {
                     Self::remove_const(&mut self.ecx, local);
-                    ensure_not_propagated(self, local);
+                    self.ensure_not_propagated(local);
                 }
             }
         }