do promote array indexing if we know it is in-bounds

6 files changed, 96 insertions, 59 deletions

diff --git a/compiler/rustc_mir/src/transform/promote_consts.rs b/compiler/rustc_mir/src/transform/promote_consts.rs
index ce3e03e3c6d..d8758e04544 100644
--- a/compiler/rustc_mir/src/transform/promote_consts.rs
+++ b/compiler/rustc_mir/src/transform/promote_consts.rs
@@ -415,10 +415,11 @@ impl<'tcx> Validator<'_, 'tcx> {
     // FIXME(eddyb) maybe cache this?
     fn validate_local(&self, local: Local) -> Result<(), Unpromotable> {
         if let TempState::Defined { location: loc, .. } = self.temps[local] {
-            let num_stmts = self.body[loc.block].statements.len();
+            let block = &self.body[loc.block];
+            let num_stmts = block.statements.len();
 
             if loc.statement_index < num_stmts {
-                let statement = &self.body[loc.block].statements[loc.statement_index];
+                let statement = &block.statements[loc.statement_index];
                 match &statement.kind {
                     StatementKind::Assign(box (_, rhs)) => self.validate_rvalue(rhs),
                     _ => {
@@ -430,7 +431,7 @@ impl<'tcx> Validator<'_, 'tcx> {
                     }
                 }
             } else {
-                let terminator = self.body[loc.block].terminator();
+                let terminator = block.terminator();
                 match &terminator.kind {
                     TerminatorKind::Call { func, args, .. } => self.validate_call(func, args),
                     TerminatorKind::Yield { .. } => Err(Unpromotable),
@@ -452,22 +453,15 @@ impl<'tcx> Validator<'_, 'tcx> {
                 match elem {
                     ProjectionElem::Deref => {
                         let mut promotable = false;
-                        // The `is_empty` predicate is introduced to exclude the case
-                        // where the projection operations are [ .field, * ].
-                        // The reason is because promotion will be illegal if field
-                        // accesses precede the dereferencing.
+                        // We need to make sure this is a `Deref` of a local with no further projections.
                         // Discussion can be found at
                         // https://github.com/rust-lang/rust/pull/74945#discussion_r463063247
-                        // There may be opportunity for generalization, but this needs to be
-                        // accounted for.
-                        if place_base.projection.is_empty() {
+                        if let Some(local) = place_base.as_local() {
                             // This is a special treatment for cases like *&STATIC where STATIC is a
                             // global static variable.
                             // This pattern is generated only when global static variables are directly
                             // accessed and is qualified for promotion safely.
-                            if let TempState::Defined { location, .. } =
-                                self.temps[place_base.local]
-                            {
+                            if let TempState::Defined { location, .. } = self.temps[local] {
                                 let def_stmt = self.body[location.block]
                                     .statements
                                     .get(location.statement_index);
@@ -505,9 +499,49 @@ impl<'tcx> Validator<'_, 'tcx> {
                     ProjectionElem::ConstantIndex { .. } | ProjectionElem::Subslice { .. } => {}
 
                     ProjectionElem::Index(local) => {
-                        // This could be OOB, so reject for implicit promotion.
                         if !self.explicit {
-                            return Err(Unpromotable);
+                            let mut promotable = false;
+                            // Only accept if we can predict the index and are indexing an array.
+                            let val = if let TempState::Defined { location: loc, .. } =
+                                self.temps[local]
+                            {
+                                let block = &self.body[loc.block];
+                                if loc.statement_index < block.statements.len() {
+                                    let statement = &block.statements[loc.statement_index];
+                                    match &statement.kind {
+                                        StatementKind::Assign(box (
+                                            _,
+                                            Rvalue::Use(Operand::Constant(c)),
+                                        )) => c.literal.try_eval_usize(self.tcx, self.param_env),
+                                        _ => None,
+                                    }
+                                } else {
+                                    None
+                                }
+                            } else {
+                                None
+                            };
+                            if let Some(idx) = val {
+                                // Determine the type of the thing we are indexing.
+                                let ty = place_base.ty(self.body, self.tcx).ty;
+                                match ty.kind() {
+                                    ty::Array(_, len) => {
+                                        // It's an array; determine its length.
+                                        if let Some(len) =
+                                            len.try_eval_usize(self.tcx, self.param_env)
+                                        {
+                                            // If the index is in-bounds, go ahead.
+                                            if idx < len {
+                                                promotable = true;
+                                            }
+                                        }
+                                    }
+                                    _ => {}
+                                }
+                            }
+                            if !promotable {
+                                return Err(Unpromotable);
+                            }
                         }
                         self.validate_local(local)?;
                     }
diff --git a/src/test/mir-opt/const_allocation2.main.ConstProp.after.32bit.mir b/src/test/mir-opt/const_allocation2.main.ConstProp.after.32bit.mir
index 9fe3234f6ae..d2e764f856f 100644
--- a/src/test/mir-opt/const_allocation2.main.ConstProp.after.32bit.mir
+++ b/src/test/mir-opt/const_allocation2.main.ConstProp.after.32bit.mir
@@ -24,41 +24,41 @@ fn main() -> () {
 }
 
 alloc0 (static: FOO, size: 8, align: 4) {
-    ╾─alloc31─╼ 03 00 00 00                         │ ╾──╼....
+    ╾─alloc27─╼ 03 00 00 00                         │ ╾──╼....
 }
 
-alloc31 (size: 48, align: 4) {
-    0x00 │ 00 00 00 00 __ __ __ __ ╾─alloc8──╼ 00 00 00 00 │ ....░░░░╾──╼....
-    0x10 │ 00 00 00 00 __ __ __ __ ╾─alloc14─╼ 02 00 00 00 │ ....░░░░╾──╼....
-    0x20 │ 01 00 00 00 2a 00 00 00 ╾─alloc29─╼ 03 00 00 00 │ ....*...╾──╼....
+alloc27 (size: 48, align: 4) {
+    0x00 │ 00 00 00 00 __ __ __ __ ╾─alloc12─╼ 00 00 00 00 │ ....░░░░╾──╼....
+    0x10 │ 00 00 00 00 __ __ __ __ ╾─alloc17─╼ 02 00 00 00 │ ....░░░░╾──╼....
+    0x20 │ 01 00 00 00 2a 00 00 00 ╾─alloc25─╼ 03 00 00 00 │ ....*...╾──╼....
 }
 
-alloc8 (size: 0, align: 4) {}
+alloc12 (size: 0, align: 4) {}
 
-alloc14 (size: 8, align: 4) {
-    ╾─alloc12─╼ ╾─alloc13─╼                         │ ╾──╼╾──╼
+alloc17 (size: 8, align: 4) {
+    ╾─alloc15─╼ ╾─alloc16─╼                         │ ╾──╼╾──╼
 }
 
-alloc12 (size: 1, align: 1) {
+alloc15 (size: 1, align: 1) {
     05                                              │ .
 }
 
-alloc13 (size: 1, align: 1) {
+alloc16 (size: 1, align: 1) {
     06                                              │ .
 }
 
-alloc29 (size: 12, align: 4) {
-    ╾─a21+0x3─╼ ╾─alloc23─╼ ╾─a28+0x2─╼             │ ╾──╼╾──╼╾──╼
+alloc25 (size: 12, align: 4) {
+    ╾─a21+0x3─╼ ╾─alloc22─╼ ╾─a24+0x2─╼             │ ╾──╼╾──╼╾──╼
 }
 
 alloc21 (size: 4, align: 1) {
     2a 45 15 6f                                     │ *E.o
 }
 
-alloc23 (size: 1, align: 1) {
+alloc22 (size: 1, align: 1) {
     2a                                              │ *
 }
 
-alloc28 (size: 4, align: 1) {
+alloc24 (size: 4, align: 1) {
     2a 45 15 6f                                     │ *E.o
 }
diff --git a/src/test/mir-opt/const_allocation2.main.ConstProp.after.64bit.mir b/src/test/mir-opt/const_allocation2.main.ConstProp.after.64bit.mir
index 3039f8ffab7..00bf91f90bf 100644
--- a/src/test/mir-opt/const_allocation2.main.ConstProp.after.64bit.mir
+++ b/src/test/mir-opt/const_allocation2.main.ConstProp.after.64bit.mir
@@ -24,44 +24,44 @@ fn main() -> () {
 }
 
 alloc0 (static: FOO, size: 16, align: 8) {
-    ╾───────alloc31───────╼ 03 00 00 00 00 00 00 00 │ ╾──────╼........
+    ╾───────alloc27───────╼ 03 00 00 00 00 00 00 00 │ ╾──────╼........
 }
 
-alloc31 (size: 72, align: 8) {
-    0x00 │ 00 00 00 00 __ __ __ __ ╾───────alloc8────────╼ │ ....░░░░╾──────╼
+alloc27 (size: 72, align: 8) {
+    0x00 │ 00 00 00 00 __ __ __ __ ╾───────alloc12───────╼ │ ....░░░░╾──────╼
     0x10 │ 00 00 00 00 00 00 00 00 00 00 00 00 __ __ __ __ │ ............░░░░
-    0x20 │ ╾───────alloc14───────╼ 02 00 00 00 00 00 00 00 │ ╾──────╼........
-    0x30 │ 01 00 00 00 2a 00 00 00 ╾───────alloc29───────╼ │ ....*...╾──────╼
+    0x20 │ ╾───────alloc17───────╼ 02 00 00 00 00 00 00 00 │ ╾──────╼........
+    0x30 │ 01 00 00 00 2a 00 00 00 ╾───────alloc25───────╼ │ ....*...╾──────╼
     0x40 │ 03 00 00 00 00 00 00 00                         │ ........
 }
 
-alloc8 (size: 0, align: 8) {}
+alloc12 (size: 0, align: 8) {}
 
-alloc14 (size: 16, align: 8) {
-    ╾───────alloc12───────╼ ╾───────alloc13───────╼ │ ╾──────╼╾──────╼
+alloc17 (size: 16, align: 8) {
+    ╾───────alloc15───────╼ ╾───────alloc16───────╼ │ ╾──────╼╾──────╼
 }
 
-alloc12 (size: 1, align: 1) {
+alloc15 (size: 1, align: 1) {
     05                                              │ .
 }
 
-alloc13 (size: 1, align: 1) {
+alloc16 (size: 1, align: 1) {
     06                                              │ .
 }
 
-alloc29 (size: 24, align: 8) {
-    0x00 │ ╾─────alloc21+0x3─────╼ ╾───────alloc23───────╼ │ ╾──────╼╾──────╼
-    0x10 │ ╾─────alloc28+0x2─────╼                         │ ╾──────╼
+alloc25 (size: 24, align: 8) {
+    0x00 │ ╾─────alloc21+0x3─────╼ ╾───────alloc22───────╼ │ ╾──────╼╾──────╼
+    0x10 │ ╾─────alloc24+0x2─────╼                         │ ╾──────╼
 }
 
 alloc21 (size: 4, align: 1) {
     2a 45 15 6f                                     │ *E.o
 }
 
-alloc23 (size: 1, align: 1) {
+alloc22 (size: 1, align: 1) {
     2a                                              │ *
 }
 
-alloc28 (size: 4, align: 1) {
+alloc24 (size: 4, align: 1) {
     2a 45 15 6f                                     │ *E.o
 }
diff --git a/src/test/mir-opt/const_allocation3.main.ConstProp.after.32bit.mir b/src/test/mir-opt/const_allocation3.main.ConstProp.after.32bit.mir
index cf7d82ef15e..519002da392 100644
--- a/src/test/mir-opt/const_allocation3.main.ConstProp.after.32bit.mir
+++ b/src/test/mir-opt/const_allocation3.main.ConstProp.after.32bit.mir
@@ -24,30 +24,30 @@ fn main() -> () {
 }
 
 alloc0 (static: FOO, size: 4, align: 4) {
-    ╾─alloc11─╼                                     │ ╾──╼
+    ╾─alloc10─╼                                     │ ╾──╼
 }
 
-alloc11 (size: 168, align: 1) {
+alloc10 (size: 168, align: 1) {
     0x00 │ ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab │ ................
-    0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾─alloc4──╼ │ ............╾──╼
+    0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾─alloc5──╼ │ ............╾──╼
     0x20 │ 01 ef cd ab 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x30 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x40 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x50 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x60 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x70 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
-    0x80 │ 00 00 00 00 00 00 00 00 00 00 ╾─alloc6──╼ 00 00 │ ..........╾──╼..
-    0x90 │ ╾─a9+0x63─╼ 00 00 00 00 00 00 00 00 00 00 00 00 │ ╾──╼............
+    0x80 │ 00 00 00 00 00 00 00 00 00 00 ╾─alloc7──╼ 00 00 │ ..........╾──╼..
+    0x90 │ ╾─a8+0x63─╼ 00 00 00 00 00 00 00 00 00 00 00 00 │ ╾──╼............
     0xa0 │ 00 00 00 00 00 00 00 00                         │ ........
 }
 
-alloc4 (size: 4, align: 4) {
+alloc5 (size: 4, align: 4) {
     2a 00 00 00                                     │ *...
 }
 
-alloc6 (fn: main)
+alloc7 (fn: main)
 
-alloc9 (size: 100, align: 1) {
+alloc8 (size: 100, align: 1) {
     0x00 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x10 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x20 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
diff --git a/src/test/mir-opt/const_allocation3.main.ConstProp.after.64bit.mir b/src/test/mir-opt/const_allocation3.main.ConstProp.after.64bit.mir
index f2227c40fcb..73c4288c32e 100644
--- a/src/test/mir-opt/const_allocation3.main.ConstProp.after.64bit.mir
+++ b/src/test/mir-opt/const_allocation3.main.ConstProp.after.64bit.mir
@@ -24,12 +24,12 @@ fn main() -> () {
 }
 
 alloc0 (static: FOO, size: 8, align: 8) {
-    ╾───────alloc11───────╼                         │ ╾──────╼
+    ╾───────alloc10───────╼                         │ ╾──────╼
 }
 
-alloc11 (size: 180, align: 1) {
+alloc10 (size: 180, align: 1) {
     0x00 │ ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab │ ................
-    0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾──alloc4── │ ............╾───
+    0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾──alloc5── │ ............╾───
     0x20 │ ──────────╼ 01 ef cd ab 00 00 00 00 00 00 00 00 │ ───╼............
     0x30 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x40 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
@@ -37,18 +37,18 @@ alloc11 (size: 180, align: 1) {
     0x60 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x70 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x80 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 ╾──── │ ..............╾─
-    0x90 │ ─────alloc6─────╼ 00 00 ╾─────alloc9+0x63─────╼ │ ─────╼..╾──────╼
+    0x90 │ ─────alloc7─────╼ 00 00 ╾─────alloc8+0x63─────╼ │ ─────╼..╾──────╼
     0xa0 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0xb0 │ 00 00 00 00                                     │ ....
 }
 
-alloc4 (size: 4, align: 4) {
+alloc5 (size: 4, align: 4) {
     2a 00 00 00                                     │ *...
 }
 
-alloc6 (fn: main)
+alloc7 (fn: main)
 
-alloc9 (size: 100, align: 1) {
+alloc8 (size: 100, align: 1) {
     0x00 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x10 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x20 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
diff --git a/src/test/ui/consts/promotion.rs b/src/test/ui/consts/promotion.rs
index 3d952964213..e7726dd0a4b 100644
--- a/src/test/ui/consts/promotion.rs
+++ b/src/test/ui/consts/promotion.rs
@@ -30,6 +30,9 @@ fn main() {
     baz_i32(&(1/1));
     baz_i32(&(1%1));
 
+    // in-bounds array access is okay
+    baz_i32(&([1,2,3][0] + 1));
+
     // Top-level projections do not get promoted, so no error here.
     if false {
         #[allow(unconditional_panic)]