Just avoid creating too many opaques.

Author: cjgillot

compiler/rustc_mir_transform/src/gvn.rs

@@ -227,12 +227,13 @@ struct VnState<'body, 'a, 'tcx> {
     local_decls: &'body LocalDecls<'tcx>,
     is_coroutine: bool,
     /// Value stored in each local.
-    locals: IndexVec<Local, VnIndex>,
+    locals: IndexVec<Local, Option<VnIndex>>,
     /// Locals that are assigned that value.
     // This vector holds the locals that are SSA.
-    rev_locals_ssa: IndexVec<VnIndex, SmallVec<[(Local, Location); 1]>>,
-    // This vector holds the locals that are not SSA.
-    rev_locals_non_ssa: FxHashMap<VnIndex, SmallVec<[(Local, Location); 1]>>,
+    rev_locals_ssa: IndexVec<VnIndex, SmallVec<[(Local, BasicBlock); 1]>>,
+    // This map holds the locals that are not SSA. This map is cleared at the end of each block.
+    // Therefore, we do not need a location, the local always appears before the current location.
+    rev_locals_non_ssa: FxHashMap<VnIndex, SmallVec<[Local; 1]>>,
     values: FxIndexSet<(Value<'a, 'tcx>, Ty<'tcx>)>,
     /// Values evaluated as constants if possible.
     evaluated: IndexVec<VnIndex, Option<OpTy<'tcx>>>,
@@ -268,7 +269,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             ecx: InterpCx::new(tcx, DUMMY_SP, typing_env, DummyMachine),
             local_decls,
             is_coroutine: body.coroutine.is_some(),
-            locals: IndexVec::with_capacity(body.local_decls.len()),
+            locals: IndexVec::from_elem(None, &body.local_decls),
             rev_locals_ssa: IndexVec::with_capacity(num_values),
             rev_locals_non_ssa: FxHashMap::default(),
             values: FxIndexSet::with_capacity_and_hasher(num_values, Default::default()),
@@ -280,14 +281,12 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             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);
+        for local in body.args_iter() {
             if ssa.is_ssa(local) {
-                this.rev_locals_ssa[value].push((local, init_loc));
-            } else {
-                this.rev_locals_non_ssa.entry(value).or_default().push((local, init_loc));
+                let ty = body.local_decls[local].ty;
+                let value = this.new_opaque(ty);
+                this.rev_locals_ssa[value].push((local, START_BLOCK));
+                this.locals[local] = Some(value);
             }
         }
         this
@@ -344,7 +343,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.local(place.local);
             // Skip the initial `Deref`.
             projection.next();
             AddressBase::Deref(base)
@@ -356,7 +355,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             .arena
             .try_alloc_from_iter(
                 projection
-                    .map(|proj| proj.try_map(|value| Some(self.locals[value]), |ty| ty).ok_or(())),
+                    .map(|proj| proj.try_map(|value| Some(self.local(value)), |ty| ty).ok_or(())),
             )
             .ok()?;
         let value = Value::Address { base, projection, kind, provenance: self.next_opaque() };
@@ -376,12 +375,24 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
     /// Record that `local` is assigned `value`.
     #[instrument(level = "trace", skip(self))]
     fn assign(&mut self, local: Local, value: VnIndex, loc: Location) {
-        self.locals[local] = value;
+        self.locals[local] = Some(value);
         if self.ssa.is_ssa(local) {
-            self.rev_locals_ssa[value].push((local, loc));
+            self.rev_locals_ssa[value].push((local, loc.block));
         } else {
-            self.rev_locals_non_ssa.entry(value).or_default().push((local, loc));
+            self.rev_locals_non_ssa.entry(value).or_default().push(local);
+        }
+    }
+
+    /// Return the value assigned to a local, or assign an opaque value and return it.
+    #[instrument(level = "trace", skip(self), ret)]
+    fn local(&mut self, local: Local) -> VnIndex {
+        if let Some(value) = self.locals[local] {
+            return value;
         }
+        let value = self.new_opaque(self.local_decls[local].ty);
+        self.locals[local] = Some(value);
+        self.rev_locals_non_ssa.entry(value).or_default().push(local);
+        value
     }
 
     #[instrument(level = "trace", skip(self))]
@@ -390,9 +401,9 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             if this.ssa.is_ssa(local) {
                 return;
             }
-            let value = this.locals[local];
-            this.rev_locals_non_ssa.entry(value).or_default().retain(|(l, _)| *l != local);
-            this.locals[local] = this.new_opaque(this.ty(value));
+            if let Some(value) = this.locals[local].take() {
+                this.rev_locals_non_ssa.entry(value).or_default().retain(|l| *l != local);
+            }
         };
         if place.is_indirect_first_projection() {
             // Non-local mutation maybe invalidate deref.
@@ -668,7 +679,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.local(local);
                 let place_ty = PlaceTy::from_ty(self.ty(local));
                 (place_ty, local)
             }
@@ -774,7 +785,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 base = self.locals[place.local]
+            && let base = self.local(place.local)
             && let Some(new_local) = self.try_as_local(base, location)
             && place.local != new_local
         {
@@ -787,7 +798,7 @@ 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 idx = self.locals[idx_local];
+                let idx = self.local(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)
@@ -823,7 +834,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 mut value = self.locals[place.local];
+        let mut value = self.local(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() {
@@ -834,7 +845,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
                 place_ref = PlaceRef { local, projection: &place.projection[index..] };
             }
 
-            let Some(proj) = proj.try_map(|value| Some(self.locals[value]), |ty| ty) else {
+            let Some(proj) = proj.try_map(|value| Some(self.local(value)), |ty| ty) else {
                 return Err(place_ref);
             };
             let Some(ty_and_value) = self.project(place_ty, value, proj) else {
@@ -1770,15 +1781,15 @@ impl<'tcx> VnState<'_, '_, 'tcx> {
     #[instrument(level = "trace", skip(self), ret)]
     fn try_as_local(&mut self, index: VnIndex, loc: Location) -> Option<Local> {
         let ssa = self.rev_locals_ssa.get(index)?;
-        let non_ssa = self.rev_locals_non_ssa.entry(index).or_default();
-        ssa.iter()
-            .chain(non_ssa.iter())
-            .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)
+        if let Some((other, _)) = ssa.iter().find(|&&(other, assign_block)| {
+            self.ssa.assignment_dominates(&self.dominators, other, loc) || assign_block == loc.block
+        }) {
+            Some(*other)
+        } else if let Some(non_ssa) = self.rev_locals_non_ssa.get(&index) {
+            non_ssa.first().copied()
+        } else {
+            None
+        }
     }
 }
 
@@ -1791,13 +1802,7 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
         self.rev_locals_non_ssa.clear();
         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_non_ssa
-                    .entry(new)
-                    .or_default()
-                    .push((local, Location { block, statement_index: 0 }));
+                self.locals[local] = None;
             }
         }
         self.super_basic_block_data(block, bbdata);

tests/mir-opt/const_prop/discriminant.main.GVN.32bit.diff

@@ -15,7 +15,8 @@
 
       bb0: {
           StorageLive(_1);
-          StorageLive(_2);
+-         StorageLive(_2);
++         nop;
           StorageLive(_3);
 -         _3 = Option::<bool>::Some(const true);
 -         _4 = discriminant(_3);
@@ -41,8 +42,10 @@
       }
 
       bb4: {
-          _1 = Add(move _2, const 0_i32);
-          StorageDead(_2);
+-         _1 = Add(move _2, const 0_i32);
+-         StorageDead(_2);
++         _1 = copy _2;
++         nop;
           StorageDead(_3);
           _0 = const ();
           StorageDead(_1);

tests/mir-opt/const_prop/discriminant.main.GVN.64bit.diff

@@ -15,7 +15,8 @@
 
       bb0: {
           StorageLive(_1);
-          StorageLive(_2);
+-         StorageLive(_2);
++         nop;
           StorageLive(_3);
 -         _3 = Option::<bool>::Some(const true);
 -         _4 = discriminant(_3);
@@ -41,8 +42,10 @@
       }
 
       bb4: {
-          _1 = Add(move _2, const 0_i32);
-          StorageDead(_2);
+-         _1 = Add(move _2, const 0_i32);
+-         StorageDead(_2);
++         _1 = copy _2;
++         nop;
           StorageDead(_3);
           _0 = const ();
           StorageDead(_1);

tests/mir-opt/const_prop/discriminant.rs

@@ -20,6 +20,6 @@ fn main() {
     // CHECK: [[tmp]] = const 10_i32;
     // CHECK: goto -> bb4;
     // CHECK: bb4: {
-    // CHECK: {{_.*}} = Add(move [[tmp]], const 0_i32);
+    // CHECK: {{_.*}} = copy [[tmp]];
     let x = (if let Some(true) = Some(true) { 42 } else { 10 }) + 0;
 }

tests/mir-opt/dont_reset_cast_kind_without_updating_operand.test.GVN.32bit.panic-abort.diff

@@ -70,8 +70,7 @@
       bb0: {
           StorageLive(_1);
           StorageLive(_2);
--         StorageLive(_3);
-+         nop;
+          StorageLive(_3);
           StorageLive(_11);
           StorageLive(_12);
           StorageLive(_13);
@@ -89,8 +88,7 @@
       }
 
       bb1: {
--         StorageDead(_3);
-+         nop;
+          StorageDead(_3);
           StorageDead(_1);
           return;
       }
@@ -123,9 +121,8 @@
           StorageDead(_2);
           StorageLive(_4);
 -         _9 = deref_copy _3;
--         _10 = copy ((_9.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *const () (Transmute);
 +         _9 = copy _3;
-+         _10 = copy ((_3.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *const () (Transmute);
+          _10 = copy ((_9.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *const () (Transmute);
           _4 = copy _10;
 -         StorageLive(_5);
 +         nop;
@@ -142,7 +139,7 @@
           StorageLive(_8);
           _8 = copy _5;
 -         _7 = copy _8 as *mut () (PtrToPtr);
-+         _7 = copy ((_3.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *mut () (Transmute);
++         _7 = copy ((_9.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *mut () (Transmute);
           StorageDead(_8);
           StorageDead(_7);
 -         StorageDead(_5);

tests/mir-opt/dont_reset_cast_kind_without_updating_operand.test.GVN.32bit.panic-unwind.diff

@@ -36,8 +36,7 @@
       bb0: {
           StorageLive(_1);
           StorageLive(_2);
--         StorageLive(_3);
-+         nop;
+          StorageLive(_3);
           _3 = Box::<()>::new(const ()) -> [return: bb1, unwind continue];
       }
 
@@ -47,9 +46,8 @@
           StorageDead(_2);
           StorageLive(_4);
 -         _9 = deref_copy _3;
--         _10 = copy ((_9.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *const () (Transmute);
 +         _9 = copy _3;
-+         _10 = copy ((_3.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *const () (Transmute);
+          _10 = copy ((_9.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *const () (Transmute);
           _4 = copy _10;
 -         StorageLive(_5);
 +         nop;
@@ -66,7 +64,7 @@
           StorageLive(_8);
           _8 = copy _5;
 -         _7 = copy _8 as *mut () (PtrToPtr);
-+         _7 = copy ((_3.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *mut () (Transmute);
++         _7 = copy ((_9.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *mut () (Transmute);
           StorageDead(_8);
           StorageDead(_7);
 -         StorageDead(_5);
@@ -76,8 +74,7 @@
       }
 
       bb2: {
--         StorageDead(_3);
-+         nop;
+          StorageDead(_3);
           StorageDead(_1);
           return;
       }

tests/mir-opt/dont_reset_cast_kind_without_updating_operand.test.GVN.64bit.panic-abort.diff

@@ -70,8 +70,7 @@
       bb0: {
           StorageLive(_1);
           StorageLive(_2);
--         StorageLive(_3);
-+         nop;
+          StorageLive(_3);
           StorageLive(_11);
           StorageLive(_12);
           StorageLive(_13);
@@ -89,8 +88,7 @@
       }
 
       bb1: {
--         StorageDead(_3);
-+         nop;
+          StorageDead(_3);
           StorageDead(_1);
           return;
       }
@@ -123,9 +121,8 @@
           StorageDead(_2);
           StorageLive(_4);
 -         _9 = deref_copy _3;
--         _10 = copy ((_9.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *const () (Transmute);
 +         _9 = copy _3;
-+         _10 = copy ((_3.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *const () (Transmute);
+          _10 = copy ((_9.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *const () (Transmute);
           _4 = copy _10;
 -         StorageLive(_5);
 +         nop;
@@ -142,7 +139,7 @@
           StorageLive(_8);
           _8 = copy _5;
 -         _7 = copy _8 as *mut () (PtrToPtr);
-+         _7 = copy ((_3.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *mut () (Transmute);
++         _7 = copy ((_9.0: std::ptr::Unique<()>).0: std::ptr::NonNull<()>) as *mut () (Transmute);
           StorageDead(_8);
           StorageDead(_7);
 -         StorageDead(_5);