Bound the size of created allocations.

Author: cjgillot

compiler/rustc_mir_transform/src/dataflow_const_prop.rs

@@ -556,20 +556,26 @@ impl<'tcx, 'locals> Collector<'tcx, 'locals> {
         let ty = place.ty(self.local_decls, self.patch.tcx).ty;
         let place = map.find(place.as_ref())?;
         let layout = ecx.layout_of(ty).ok()?;
-        if layout.abi.is_scalar() {
-            let value = propagatable_scalar(*ecx.tcx, place, state, map)?;
+        if layout.is_zst() {
+            Some(ConstantKind::Val(ConstValue::ZeroSized, ty))
+        } else if layout.abi.is_scalar()
+            && let Some(value) = propagatable_scalar(*ecx.tcx, place, state, map)
+        {
             Some(ConstantKind::Val(ConstValue::Scalar(value), ty))
-        } else {
+        } else if layout.is_sized() && layout.size <= 4 * ecx.tcx.data_layout.pointer_size {
             let alloc_id = ecx
                 .intern_with_temp_alloc(layout, |ecx, dest| {
                     try_write_constant(ecx, dest, place, ty, state, map)
                 })
                 .ok()?;
             Some(ConstantKind::Val(ConstValue::Indirect { alloc_id, offset: Size::ZERO }, ty))
+        } else {
+            None
         }
     }
 }
 
+#[instrument(level = "trace", skip(tcx, state, map), ret)]
 fn propagatable_scalar<'tcx>(
     tcx: TyCtxt<'tcx>,
     place: PlaceIndex,
@@ -602,17 +608,25 @@ fn try_write_constant<'tcx>(
     map: &Map,
 ) -> InterpResult<'tcx> {
     let layout = ecx.layout_of(ty)?;
+
+    // Fast path for ZSTs.
+    if layout.is_zst() {
+        return Ok(());
+    }
+
+    // Fast path for scalars.
+    if layout.abi.is_scalar()
+        && let Some(value) = propagatable_scalar(*ecx.tcx, place, state, map)
+    {
+        return ecx.write_immediate(Immediate::Scalar(value), dest);
+    }
+
     match ty.kind() {
         // ZSTs. Nothing to do.
         ty::FnDef(..) => {}
 
-        // Scalars.
-        _ if layout.abi.is_scalar() => {
-            let value = propagatable_scalar(*ecx.tcx, place, state, map).ok_or(err_inval!(ConstPropNonsense))?;
-            ecx.write_immediate(Immediate::Scalar(value), dest)?;
-        }
         // Those are scalars, must be handled above.
-        ty::Bool | ty::Int(_) | ty::Uint(_) | ty::Float(_) | ty::Char => bug!(),
+        ty::Bool | ty::Int(_) | ty::Uint(_) | ty::Float(_) | ty::Char => throw_inval!(ConstPropNonsense),
 
         ty::Tuple(elem_tys) => {
             for (i, elem) in elem_tys.iter().enumerate() {

tests/mir-opt/dataflow-const-prop/enum.statics.DataflowConstProp.32bit.diff

@@ -116,7 +116,7 @@
 
 - alloc2 (static: RC, size: 4, align: 4) {
 -     ╾─alloc14─╼                                     │ ╾──╼
-+ alloc20 (size: 8, align: 4) {
++ alloc28 (size: 8, align: 4) {
 +     00 00 00 00 00 00 00 00                         │ ........
   }
 

tests/mir-opt/dataflow-const-prop/enum.statics.DataflowConstProp.64bit.diff

@@ -116,7 +116,7 @@
 
 - alloc2 (static: RC, size: 8, align: 8) {
 -     ╾───────alloc14───────╼                         │ ╾──────╼
-+ alloc20 (size: 8, align: 4) {
++ alloc28 (size: 8, align: 4) {
 +     00 00 00 00 00 00 00 00                         │ ........
   }
 

tests/mir-opt/dataflow-const-prop/repr_transparent.main.DataflowConstProp.diff

@@ -17,7 +17,8 @@
 
       bb0: {
           StorageLive(_1);
-          _1 = I32(const 0_i32);
+-         _1 = I32(const 0_i32);
++         _1 = const I32(0_i32);
           StorageLive(_2);
           StorageLive(_3);
           StorageLive(_4);
@@ -31,12 +32,20 @@
           StorageDead(_5);
           StorageDead(_4);
 -         _2 = I32(move _3);
-+         _2 = I32(const 0_i32);
++         _2 = const I32(0_i32);
           StorageDead(_3);
           _0 = const ();
           StorageDead(_2);
           StorageDead(_1);
           return;
       }
++ }
++ 
++ alloc5 (size: 4, align: 4) {
++     00 00 00 00                                     │ ....
++ }
++ 
++ alloc4 (size: 4, align: 4) {
++     00 00 00 00                                     │ ....
   }