CopyProp: Propagate moves of constants

Author: clubby789

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -1,3 +1,4 @@
+use rustc_data_structures::fx::FxIndexMap;
 use rustc_index::bit_set::BitSet;
 use rustc_index::IndexSlice;
 use rustc_middle::mir::visit::*;
@@ -37,6 +38,8 @@ fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
 
     let fully_moved = fully_moved_locals(&ssa, body);
     debug!(?fully_moved);
+    let const_locals = const_locals(&ssa, body);
+    debug!(?const_locals);
 
     let mut storage_to_remove = BitSet::new_empty(fully_moved.domain_size());
     for (local, &head) in ssa.copy_classes().iter_enumerated() {
@@ -51,6 +54,7 @@ fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         tcx,
         copy_classes: ssa.copy_classes(),
         fully_moved,
+        const_locals,
         borrowed_locals,
         storage_to_remove,
     }
@@ -96,10 +100,28 @@ fn fully_moved_locals(ssa: &SsaLocals, body: &Body<'_>) -> BitSet<Local> {
     fully_moved
 }
 
+/// Finds all locals that are only assigned to once with a deterministic constant
+#[instrument(level = "trace", skip(ssa, body))]
+fn const_locals<'body, 'tcx>(
+    ssa: &'body SsaLocals,
+    body: &'body Body<'tcx>,
+) -> FxIndexMap<Local, ConstOperand<'tcx>> {
+    let mut const_locals = FxIndexMap::default();
+    for (local, rvalue, _) in ssa.assignments(body) {
+        let Rvalue::Use(Operand::Constant(val)) = rvalue else { continue };
+        if !val.const_.is_deterministic() {
+            continue;
+        }
+        const_locals.insert(local, **val);
+    }
+    const_locals
+}
+
 /// Utility to help performing substitution of `*pattern` by `target`.
 struct Replacer<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     fully_moved: BitSet<Local>,
+    const_locals: FxIndexMap<Local, ConstOperand<'tcx>>,
     storage_to_remove: BitSet<Local>,
     borrowed_locals: BitSet<Local>,
     copy_classes: &'a IndexSlice<Local, Local>,
@@ -154,9 +176,14 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
         if let Operand::Move(place) = *operand
             // A move out of a projection of a copy is equivalent to a copy of the original projection.
             && !place.is_indirect_first_projection()
-            && !self.fully_moved.contains(place.local)
         {
-            *operand = Operand::Copy(place);
+            if !self.fully_moved.contains(place.local) {
+                *operand = Operand::Copy(place);
+            } else if let Some(local) = place.as_local()
+                && let Some(val) = self.const_locals.get(&local)
+            {
+                *operand = Operand::Constant(Box::new(*val))
+            }
         }
         self.super_operand(operand, loc);
     }