Add an InstSimplify for repetitive array expressions

Author: saethlin

compiler/rustc_mir_transform/src/instsimplify.rs

@@ -48,6 +48,7 @@ impl<'tcx> crate::MirPass<'tcx> for InstSimplify {
                         ctx.simplify_ref_deref(rvalue);
                         ctx.simplify_ptr_aggregate(rvalue);
                         ctx.simplify_cast(rvalue);
+                        ctx.simplify_repeated_aggregate(rvalue);
                     }
                     _ => {}
                 }
@@ -68,6 +69,35 @@ struct InstSimplifyContext<'a, 'tcx> {
 }
 
 impl<'tcx> InstSimplifyContext<'_, 'tcx> {
+    /// Transform aggregates like [0, 0, 0, 0, 0] into [0; 5].
+    /// GVN can also do this optimization, but GVN is only run at mir-opt-level 2 so having this in
+    /// InstSimplify helps unoptimized builds.
+    fn simplify_repeated_aggregate(&self, rvalue: &mut Rvalue<'tcx>) {
+        let Rvalue::Aggregate(box AggregateKind::Array(_), fields) = rvalue else {
+            return;
+        };
+        if fields.len() < 5 {
+            return;
+        }
+        let first = &fields[rustc_abi::FieldIdx::ZERO];
+        let Operand::Constant(first) = first else {
+            return;
+        };
+        let Ok(first_val) = first.const_.eval(self.tcx, self.typing_env, first.span) else {
+            return;
+        };
+        if fields.iter().all(|field| {
+            let Operand::Constant(field) = field else {
+                return false;
+            };
+            let field = field.const_.eval(self.tcx, self.typing_env, field.span);
+            field == Ok(first_val)
+        }) {
+            let len = ty::Const::from_target_usize(self.tcx, fields.len().try_into().unwrap());
+            *rvalue = Rvalue::Repeat(Operand::Constant(first.clone()), len);
+        }
+    }
+
     /// Transform boolean comparisons into logical operations.
     fn simplify_bool_cmp(&self, rvalue: &mut Rvalue<'tcx>) {
         match rvalue {

tests/mir-opt/instsimplify/aggregate_array.make_array.InstSimplify-after-simplifycfg.diff

@@ -0,0 +1,13 @@
+- // MIR for `make_array` before InstSimplify-after-simplifycfg
++ // MIR for `make_array` after InstSimplify-after-simplifycfg
+  
+  fn make_array() -> [u8; 5] {
+      let mut _0: [u8; 5];
+  
+      bb0: {
+-         _0 = [const 0_u8, const 0_u8, const 0_u8, const 0_u8, const 0_u8];
++         _0 = [const 0_u8; 5];
+          return;
+      }
+  }
+  

tests/mir-opt/instsimplify/aggregate_array.rs

@@ -0,0 +1,11 @@
+//@ test-mir-pass: InstSimplify-after-simplifycfg
+#![crate_type = "lib"]
+
+// EMIT_MIR aggregate_array.make_array.InstSimplify-after-simplifycfg.diff
+pub fn make_array() -> [u8; 5] {
+    // CHECK-LABEL: fn make_array(
+    // CHECK-NOT: _0 = [const 0_u8, const 0_u8, const 0_u8, const 0_u8, const 0_u8];
+    // CHECK: _0 = [const 0_u8; 5];
+
+    [0, 0, 0, 0, 0]
+}