rustc_mir_transform: Add a local value numbering pass, off by default.

This commit adds a local [value numbering] pass, which is intended to
eventually subsume constant and copy propagation. It's designed to detect
whether a value has been computed multiple times and replace such duplicate
computations with a reference to the originally-computed value. LLVM performs
this optimization, but it's limited because it doesn't have the ability to
reason about immutability of memory the way Rust MIR can. For example, this
pass optimizes [the following code]:

    let s1: &S = ...;
    let s2: &mut S = ...;
    let mut sum = 0;
    sum += s1.a + s1.b;
    s2.a = 1;
    sum += s1.a + s1.b;

into:

    let s1: &S = ...;
    let s2: &mut S = ...;
    let mut sum = 0;
    sum += s1.a + s1.b;
    s2.a = 1;
    sum += sum;

LLVM won't do this optimization because it can't prove that the assignment to a
field of `s2` doesn't mutate the fields of `s1`.

Because this pass is likely to have bugs and may be slow, I've turned it off by
default for now. I didn't notice any show-stopping bugs in the test suite, but
out of an abundance of caution I'm keeping it off.

Closes #91688.

[value numbering]: https://en.wikipedia.org/wiki/Value_numbering

[the following code]: https://rust.godbolt.org/z/7f6Gff7Ms

Author: pcwalton

compiler/rustc_mir_transform/src/const_goto.rs

@@ -22,7 +22,7 @@ use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
 use rustc_middle::{mir::visit::Visitor, ty::ParamEnv};
 
-use super::simplify::{simplify_cfg, simplify_locals};
+use super::simplify::{remove_unused_locals, simplify_cfg};
 
 pub struct ConstGoto;
 
@@ -51,7 +51,7 @@ impl<'tcx> MirPass<'tcx> for ConstGoto {
         // make it easier for further passes
         if should_simplify {
             simplify_cfg(tcx, body);
-            simplify_locals(body, tcx);
+            remove_unused_locals(body, tcx);
         }
     }
 }

compiler/rustc_mir_transform/src/lib.rs

@@ -87,6 +87,7 @@ mod simplify_comparison_integral;
 mod simplify_try;
 mod uninhabited_enum_branching;
 mod unreachable_prop;
+mod value_numbering;
 
 use rustc_const_eval::transform::check_consts::{self, ConstCx};
 use rustc_const_eval::transform::promote_consts;
@@ -480,9 +481,11 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &separate_const_switch::SeparateConstSwitch,
             //
             // FIXME(#70073): This pass is responsible for both optimization as well as some lints.
-            &const_prop::ConstProp,
             //
             // Const-prop runs unconditionally, but doesn't mutate the MIR at mir-opt-level=0.
+            &const_prop::ConstProp,
+            // Experimental: only runs if `-Z number-values` is on.
+            &value_numbering::ValueNumbering,
             &o1(simplify_branches::SimplifyConstCondition::new("after-const-prop")),
             &early_otherwise_branch::EarlyOtherwiseBranch,
             &simplify_comparison_integral::SimplifyComparisonIntegral,

compiler/rustc_mir_transform/src/simplify.rs

@@ -372,11 +372,20 @@ impl<'tcx> MirPass<'tcx> for SimplifyLocals {
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         trace!("running SimplifyLocals on {:?}", body.source);
-        simplify_locals(body, tcx);
+        // FIXME(pcwalton): Maybe we should call `promote_mutable_locals_to_immutable()` here.
+        // Right now, the only pass that really benefits from this is ValueNumbering, which calls it
+        // itself.
+        remove_unused_locals(body, tcx);
     }
 }
 
-pub fn simplify_locals<'tcx>(body: &mut Body<'tcx>, tcx: TyCtxt<'tcx>) {
+pub fn promote_mutable_locals_to_immutable<'tcx>(body: &mut Body<'tcx>, _: TyCtxt<'tcx>) {
+    let mut mutable_locals = PromoteMutableLocals::new(body);
+    mutable_locals.visit_body(body);
+    mutable_locals.update_locals(body);
+}
+
+pub fn remove_unused_locals<'tcx>(body: &mut Body<'tcx>, tcx: TyCtxt<'tcx>) {
     // First, we're going to get a count of *actual* uses for every `Local`.
     let mut used_locals = UsedLocals::new(body);
 
@@ -424,6 +433,39 @@ fn make_local_map<V>(
     map
 }
 
+/// Promotes mutable locals to immutable.
+struct PromoteMutableLocals {
+    mutating_use_count: IndexVec<Local, u32>,
+}
+
+impl PromoteMutableLocals {
+    fn new<'tcx>(body: &Body<'tcx>) -> PromoteMutableLocals {
+        PromoteMutableLocals { mutating_use_count: IndexVec::from_elem(0, &body.local_decls) }
+    }
+
+    fn update_locals<'tcx>(&self, body: &mut Body<'tcx>) {
+        for (local, local_decl) in body.local_decls.iter_mut().enumerate() {
+            let local = Local::from_usize(local);
+            if local.as_usize() == 0 {
+                // Return value must always be mutable.
+                continue;
+            }
+            if local_decl.mutability == Mutability::Mut && self.mutating_use_count[local] <= 1 {
+                local_decl.mutability = Mutability::Not;
+            }
+        }
+    }
+}
+
+impl<'tcx> Visitor<'tcx> for PromoteMutableLocals {
+    fn visit_local(&mut self, local: &Local, place_context: PlaceContext, _: Location) {
+        match place_context {
+            PlaceContext::MutatingUse(_) => self.mutating_use_count[*local] += 1,
+            PlaceContext::NonMutatingUse(_) | PlaceContext::NonUse(_) => {}
+        }
+    }
+}
+
 /// Keeps track of used & unused locals.
 struct UsedLocals {
     increment: bool,