Use a VnIndex in Address projection.

Author: cjgillot

compiler/rustc_mir_transform/src/gvn.rs

@@ -163,6 +163,14 @@ enum AddressKind {
     Address(RawPtrKind),
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
+enum AddressBase {
+    /// This address is based on this local.
+    Local(Local),
+    /// This address is based on the deref of this pointer.
+    Deref(VnIndex),
+}
+
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
 enum Value<'a, 'tcx> {
     // Root values.
@@ -191,7 +199,10 @@ enum Value<'a, 'tcx> {
     Repeat(VnIndex, ty::Const<'tcx>),
     /// The address of a place.
     Address {
-        place: Place<'tcx>,
+        base: AddressBase,
+        // We do not use a plain `Place` as we want to be able to reason about indices.
+        // This does not contain any `Deref` projection.
+        projection: &'a [ProjectionElem<VnIndex, Ty<'tcx>>],
         kind: AddressKind,
         /// Give each borrow and pointer a different provenance, so we don't merge them.
         provenance: usize,
@@ -308,16 +319,38 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
 
     /// Create a new `Value::Address` distinct from all the others.
     #[instrument(level = "trace", skip(self), ret)]
-    fn new_pointer(&mut self, place: Place<'tcx>, kind: AddressKind) -> VnIndex {
+    fn new_pointer(
+        &mut self,
+        place: Place<'tcx>,
+        kind: AddressKind,
+        location: Location,
+    ) -> Option<VnIndex> {
         let pty = place.ty(self.local_decls, self.tcx).ty;
         let ty = match kind {
             AddressKind::Ref(bk) => {
                 Ty::new_ref(self.tcx, self.tcx.lifetimes.re_erased, pty, bk.to_mutbl_lossy())
             }
             AddressKind::Address(mutbl) => Ty::new_ptr(self.tcx, pty, mutbl.to_mutbl_lossy()),
         };
-        let value = Value::Address { place, kind, provenance: self.next_opaque() };
-        self.insert(ty, value)
+
+        let mut projection = place.projection.iter();
+        let base = if place.is_indirect_first_projection() {
+            let base = self.locals[place.local]?;
+            // Skip the initial `Deref`.
+            projection.next();
+            AddressBase::Deref(base)
+        } else {
+            AddressBase::Local(place.local)
+        };
+        // Do not try evaluating inside `Index`, this has been done by `simplify_place_value`.
+        let projection = self
+            .arena
+            .try_alloc_from_iter(
+                projection.map(|proj| proj.try_map(|value| self.locals[value], |ty| ty).ok_or(())),
+            )
+            .ok()?;
+        let value = Value::Address { base, projection, kind, provenance: self.next_opaque() };
+        Some(self.insert(ty, value))
     }
 
     #[inline]
@@ -458,14 +491,15 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
                 let elem = elem.try_map(|_| None, |()| ty.ty)?;
                 self.ecx.project(base, elem).discard_err()?
             }
-            Address { place, kind: _, provenance: _ } => {
-                if !place.is_indirect_first_projection() {
-                    return None;
-                }
-                let local = self.locals[place.local]?;
-                let pointer = self.evaluated[local].as_ref()?;
+            Address { base, projection, .. } => {
+                debug_assert!(!projection.contains(&ProjectionElem::Deref));
+                let pointer = match base {
+                    AddressBase::Deref(pointer) => self.evaluated[pointer].as_ref()?,
+                    // We have no stack to point to.
+                    AddressBase::Local(_) => return None,
+                };
                 let mut mplace = self.ecx.deref_pointer(pointer).discard_err()?;
-                for elem in place.projection.iter().skip(1) {
+                for elem in projection {
                     // `Index` by constants should have been replaces by `ConstantIndex` by
                     // `simplify_place_projection`.
                     let elem = elem.try_map(|_| None, |ty| ty)?;
@@ -589,19 +623,51 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         Some(op)
     }
 
+    /// Represent the *value* we obtain by dereferencing an `Address` value.
+    #[instrument(level = "trace", skip(self), ret)]
+    fn dereference_address(
+        &mut self,
+        base: AddressBase,
+        projection: &[ProjectionElem<VnIndex, Ty<'tcx>>],
+    ) -> Option<VnIndex> {
+        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 place_ty = PlaceTy::from_ty(self.ty(local));
+                (place_ty, local)
+            }
+            // The base is a pointer's deref, so we introduce the implicit deref.
+            AddressBase::Deref(reborrow) => {
+                let place_ty = PlaceTy::from_ty(self.ty(reborrow));
+                self.project(place_ty, reborrow, ProjectionElem::Deref)?
+            }
+        };
+        for &proj in projection {
+            (place_ty, value) = self.project(place_ty, value, proj)?;
+        }
+        Some(value)
+    }
+
+    #[instrument(level = "trace", skip(self), ret)]
     fn project(
         &mut self,
         place_ty: PlaceTy<'tcx>,
         value: VnIndex,
-        proj: PlaceElem<'tcx>,
-        from_non_ssa_index: &mut bool,
+        proj: ProjectionElem<VnIndex, Ty<'tcx>>,
     ) -> Option<(PlaceTy<'tcx>, VnIndex)> {
         let projection_ty = place_ty.projection_ty(self.tcx, proj);
         let proj = match proj {
             ProjectionElem::Deref => {
                 if let Some(Mutability::Not) = place_ty.ty.ref_mutability()
                     && projection_ty.ty.is_freeze(self.tcx, self.typing_env())
                 {
+                    if let Value::Address { base, projection, .. } = self.get(value)
+                        && let Some(value) = self.dereference_address(base, projection)
+                    {
+                        return Some((projection_ty, value));
+                    }
+
                     // An immutable borrow `_x` always points to the same value for the
                     // lifetime of the borrow, so we can merge all instances of `*_x`.
                     return Some((projection_ty, self.insert_deref(projection_ty.ty, value)));
@@ -638,10 +704,8 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             }
             ProjectionElem::Index(idx) => {
                 if let Value::Repeat(inner, _) = self.get(value) {
-                    *from_non_ssa_index |= self.locals[idx].is_none();
                     return Some((projection_ty, inner));
                 }
-                let idx = self.locals[idx]?;
                 ProjectionElem::Index(idx)
             }
             ProjectionElem::ConstantIndex { offset, min_length, from_end } => {
@@ -720,63 +784,71 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
     /// Represent the *value* which would be read from `place`, and point `place` to a preexisting
     /// place with the same value (if that already exists).
     #[instrument(level = "trace", skip(self), ret)]
-    fn simplify_place_value(
+    fn compute_place_value(
         &mut self,
-        place: &mut Place<'tcx>,
+        place: Place<'tcx>,
         location: Location,
-    ) -> Option<VnIndex> {
-        self.simplify_place_projection(place, location);
-
+    ) -> Result<VnIndex, PlaceRef<'tcx>> {
         // Invariant: `place` and `place_ref` point to the same value, even if they point to
         // different memory locations.
         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 Some(mut value) = self.locals[place.local] else { return Err(place_ref) };
         // 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);
-        let mut from_non_ssa_index = false;
         for (index, proj) in place.projection.iter().enumerate() {
-            if let Value::Projection(pointer, ProjectionElem::Deref) = self.get(value)
-                && let Value::Address { place: mut pointee, kind, .. } = self.get(pointer)
-                && let AddressKind::Ref(BorrowKind::Shared) = kind
-                && let Some(v) = self.simplify_place_value(&mut pointee, location)
-            {
-                value = v;
-                place_ref = pointee.project_deeper(&place.projection[index..], self.tcx).as_ref();
-            }
             if let Some(local) = self.try_as_local(value, location) {
                 // Both `local` and `Place { local: place.local, projection: projection[..index] }`
                 // hold the same value. Therefore, following place holds the value in the original
                 // `place`.
                 place_ref = PlaceRef { local, projection: &place.projection[index..] };
             }
 
-            (place_ty, value) = self.project(place_ty, value, proj, &mut from_non_ssa_index)?;
+            let Some(proj) = proj.try_map(|value| self.locals[value], |ty| ty) else {
+                return Err(place_ref);
+            };
+            let Some(ty_and_value) = self.project(place_ty, value, proj) else {
+                return Err(place_ref);
+            };
+            (place_ty, value) = ty_and_value;
         }
 
-        if let Value::Projection(pointer, ProjectionElem::Deref) = self.get(value)
-            && let Value::Address { place: mut pointee, kind, .. } = self.get(pointer)
-            && let AddressKind::Ref(BorrowKind::Shared) = kind
-            && let Some(v) = self.simplify_place_value(&mut pointee, location)
-        {
-            value = v;
-            place_ref = pointee.project_deeper(&[], self.tcx).as_ref();
-        }
-        if let Some(new_local) = self.try_as_local(value, location) {
-            place_ref = PlaceRef { local: new_local, projection: &[] };
-        } else if from_non_ssa_index {
-            // If access to non-SSA locals is unavoidable, bail out.
-            return None;
-        }
+        Ok(value)
+    }
 
-        if place_ref.local != place.local || place_ref.projection.len() < place.projection.len() {
-            // By the invariant on `place_ref`.
-            *place = place_ref.project_deeper(&[], self.tcx);
-            self.reused_locals.insert(place_ref.local);
-        }
+    /// Represent the *value* which would be read from `place`, and point `place` to a preexisting
+    /// place with the same value (if that already exists).
+    #[instrument(level = "trace", skip(self), ret)]
+    fn simplify_place_value(
+        &mut self,
+        place: &mut Place<'tcx>,
+        location: Location,
+    ) -> Option<VnIndex> {
+        self.simplify_place_projection(place, location);
 
-        Some(value)
+        match self.compute_place_value(*place, location) {
+            Ok(value) => {
+                if let Some(new_place) = self.try_as_place(value, location, true)
+                    && (new_place.local != place.local
+                        || new_place.projection.len() < place.projection.len())
+                {
+                    *place = new_place;
+                    self.reused_locals.insert(new_place.local);
+                }
+                Some(value)
+            }
+            Err(place_ref) => {
+                if place_ref.local != place.local
+                    || place_ref.projection.len() < place.projection.len()
+                {
+                    // By the invariant on `place_ref`.
+                    *place = place_ref.project_deeper(&[], self.tcx);
+                    self.reused_locals.insert(place_ref.local);
+                }
+                None
+            }
+        }
     }
 
     #[instrument(level = "trace", skip(self), ret)]
@@ -823,11 +895,11 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             Rvalue::Aggregate(..) => return self.simplify_aggregate(lhs, rvalue, location),
             Rvalue::Ref(_, borrow_kind, ref mut place) => {
                 self.simplify_place_projection(place, location);
-                return Some(self.new_pointer(*place, AddressKind::Ref(borrow_kind)));
+                return self.new_pointer(*place, AddressKind::Ref(borrow_kind), location);
             }
             Rvalue::RawPtr(mutbl, ref mut place) => {
                 self.simplify_place_projection(place, location);
-                return Some(self.new_pointer(*place, AddressKind::Address(mutbl)));
+                return self.new_pointer(*place, AddressKind::Address(mutbl), location);
             }
             Rvalue::WrapUnsafeBinder(ref mut op, _) => {
                 let value = self.simplify_operand(op, location)?;
@@ -1069,12 +1141,10 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
                     }
 
                     // `&mut *p`, `&raw *p`, etc don't change metadata.
-                    Value::Address { place, kind: _, provenance: _ }
-                        if let PlaceRef { local, projection: [PlaceElem::Deref] } =
-                            place.as_ref()
-                            && let Some(local_index) = self.locals[local] =>
+                    Value::Address { base: AddressBase::Deref(reborrowed), projection, .. }
+                        if projection.is_empty() =>
                     {
-                        arg_index = local_index;
+                        arg_index = reborrowed;
                         was_updated = true;
                         continue;
                     }
@@ -1420,9 +1490,9 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
 
         // The length information is stored in the wide pointer.
         // Reborrowing copies length information from one pointer to the other.
-        while let Value::Address { place: borrowed, .. } = self.get(inner)
-            && let [PlaceElem::Deref] = borrowed.projection[..]
-            && let Some(borrowed) = self.locals[borrowed.local]
+        while let Value::Address { base: AddressBase::Deref(borrowed), projection, .. } =
+            self.get(inner)
+            && projection.is_empty()
         {
             inner = borrowed;
         }

tests/mir-opt/gvn_repeat.index_place.GVN.diff

@@ -0,0 +1,15 @@
+- // MIR for `index_place` before GVN
++ // MIR for `index_place` after GVN
+  
+  fn index_place(_1: usize, _2: usize, _3: [i32; 5]) -> i32 {
+      let mut _0: i32;
+      let mut _4: &i32;
+  
+      bb0: {
+          _4 = &_3[_1];
+          _1 = copy _2;
+          _0 = copy (*_4);
+          return;
+      }
+  }
+  

tests/mir-opt/gvn_repeat.repeat_local.GVN.diff

@@ -10,8 +10,7 @@
           _4 = [copy _3; 5];
           _5 = &_4[_1];
           _1 = copy _2;
--         _0 = copy (*_5);
-+         _0 = copy _3;
+          _0 = copy (*_5);
           return;
       }
   }

tests/mir-opt/gvn_repeat.rs

@@ -6,6 +6,23 @@
 
 use std::intrinsics::mir::*;
 
+// EMIT_MIR gvn_repeat.index_place.GVN.diff
+#[custom_mir(dialect = "runtime")]
+pub fn index_place(mut idx1: usize, idx2: usize, array: [i32; 5]) -> i32 {
+    // CHECK-LABEL: fn index_place(
+    // CHECK: let mut [[ELEM:.*]]: &i32;
+    // CHECK: _0 = copy (*[[ELEM]])
+    mir! {
+        let elem;
+        {
+            elem = &array[idx1];
+            idx1 = idx2;
+            RET = *elem;
+            Return()
+        }
+    }
+}
+
 // EMIT_MIR gvn_repeat.repeat_place.GVN.diff
 #[custom_mir(dialect = "runtime")]
 pub fn repeat_place(mut idx1: usize, idx2: usize, val: &i32) -> i32 {
@@ -29,7 +46,8 @@ pub fn repeat_place(mut idx1: usize, idx2: usize, val: &i32) -> i32 {
 #[custom_mir(dialect = "runtime")]
 pub fn repeat_local(mut idx1: usize, idx2: usize, val: i32) -> i32 {
     // CHECK-LABEL: fn repeat_local(
-    // CHECK: _0 = copy _3
+    // CHECK: let mut [[ELEM:.*]]: &i32;
+    // CHECK: _0 = copy (*[[ELEM]]);
     mir! {
         let array;
         let elem;
@@ -44,6 +62,7 @@ pub fn repeat_local(mut idx1: usize, idx2: usize, val: i32) -> i32 {
 }
 
 fn main() {
+    assert_eq!(index_place(0, 5, [0; 5]), 0);
     assert_eq!(repeat_place(0, 5, &0), 0);
     assert_eq!(repeat_local(0, 5, 0), 0);
 }