Forward the ABI of the non-zero sized fields of an union if they have the same ABI

This is supposed to fix the performence regression of using MaybeUninit in
https://github.com/rust-lang/rust/pull/54668

3 files changed, 120 insertions, 1 deletions

diff --git a/src/librustc/ty/layout.rs b/src/librustc/ty/layout.rs
index 877bd5a82e6..0b00cfc2d7f 100644
--- a/src/librustc/ty/layout.rs
+++ b/src/librustc/ty/layout.rs
@@ -697,7 +697,9 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                             Align::from_bytes(repr_align, repr_align).unwrap());
                     }
 
+                    let optimize = !def.repr.inhibit_union_abi_opt();
                     let mut size = Size::ZERO;
+                    let mut abi = Abi::Aggregate { sized: true };
                     let index = VariantIdx::new(0);
                     for field in &variants[index] {
                         assert!(!field.is_unsized());
@@ -708,13 +710,44 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                         } else {
                             align = align.max(field.align);
                         }
+
+                        // If all non-ZST fields have the same ABI, forward this ABI
+                        if optimize && !field.is_zst() {
+                            // Normalize scalar_unit to the maximal valid range
+                            let field_abi = match &field.abi {
+                                Abi::Scalar(x) => Abi::Scalar(scalar_unit(x.value)),
+                                Abi::ScalarPair(x, y) => {
+                                    Abi::ScalarPair(
+                                        scalar_unit(x.value),
+                                        scalar_unit(y.value),
+                                    )
+                                }
+                                Abi::Vector { element: x, count } => {
+                                    Abi::Vector {
+                                        element: scalar_unit(x.value),
+                                        count: *count,
+                                    }
+                                }
+                                Abi::Uninhabited |
+                                Abi::Aggregate { .. }  => Abi::Aggregate { sized: true },
+                            };
+
+                            if size == Size::ZERO {
+                                // first non ZST: initialize 'abi'
+                                abi = field_abi;
+                            } else if abi != field_abi  {
+                                // different fields have different ABI: reset to Aggregate
+                                abi = Abi::Aggregate { sized: true };
+                            }
+                        }
+
                         size = cmp::max(size, field.size);
                     }
 
                     return Ok(tcx.intern_layout(LayoutDetails {
                         variants: Variants::Single { index },
                         fields: FieldPlacement::Union(variants[index].len()),
-                        abi: Abi::Aggregate { sized: true },
+                        abi,
                         align,
                         size: size.abi_align(align)
                     }));
diff --git a/src/librustc/ty/mod.rs b/src/librustc/ty/mod.rs
index 83feadd50d7..9553d124483 100644
--- a/src/librustc/ty/mod.rs
+++ b/src/librustc/ty/mod.rs
@@ -1994,6 +1994,12 @@ impl ReprOptions {
     pub fn inhibit_struct_field_reordering_opt(&self) -> bool {
         !(self.flags & ReprFlags::IS_UNOPTIMISABLE).is_empty() || (self.pack == 1)
     }
+
+    /// Returns true if this `#[repr()]` should inhibit union abi optimisations
+    pub fn inhibit_union_abi_opt(&self) -> bool {
+        self.c()
+    }
+
 }
 
 impl<'a, 'gcx, 'tcx> AdtDef {
diff --git a/src/test/codegen/union-abi.rs b/src/test/codegen/union-abi.rs
new file mode 100644
index 00000000000..0fa06fa777b
--- /dev/null
+++ b/src/test/codegen/union-abi.rs
@@ -0,0 +1,80 @@
+// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// compile-flags: -C no-prepopulate-passes
+
+// This test that using union forward the abi of the inner type, as
+// discussed in #54668
+
+#![crate_type="lib"]
+#![feature(repr_simd)]
+
+#[derive(Copy, Clone)]
+pub enum Unhab {}
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+pub struct i64x4(i64, i64, i64, i64);
+
+#[derive(Copy, Clone)]
+pub union UnionI64x4{ a:(), b: i64x4 }
+
+// CHECK: define <4 x i64> @test_UnionI64x4(<4 x i64> %arg0)
+#[no_mangle]
+pub extern fn test_UnionI64x4(_: UnionI64x4) -> UnionI64x4 { loop {} }
+
+pub union UnionI64x4_{ a: i64x4, b: (), c:i64x4, d: Unhab, e: ((),()), f: UnionI64x4 }
+
+// CHECK: define <4 x i64> @test_UnionI64x4_(<4 x i64> %arg0)
+#[no_mangle]
+pub extern fn test_UnionI64x4_(_: UnionI64x4_) -> UnionI64x4_ { loop {} }
+
+pub union UnionI64x4I64{ a: i64x4, b: i64 }
+
+// CHECK: define void @test_UnionI64x4I64(%UnionI64x4I64* {{.*}} %arg0)
+#[no_mangle]
+pub fn test_UnionI64x4I64(_: UnionI64x4I64) { loop {} }
+
+pub union UnionI64x4Tuple{ a: i64x4, b: (i64, i64, i64, i64) }
+
+// CHECK: define void @test_UnionI64x4Tuple(%UnionI64x4Tuple* {{.*}} %arg0)
+#[no_mangle]
+pub fn test_UnionI64x4Tuple(_: UnionI64x4Tuple) { loop {} }
+
+
+pub union UnionF32{a:f32}
+
+// CHECK: define float @test_UnionF32(float %arg0)
+#[no_mangle]
+pub extern fn test_UnionF32(_: UnionF32) -> UnionF32 { loop {} }
+
+pub union UnionF32F32{a:f32, b:f32}
+
+// CHECK: define float @test_UnionF32F32(float %arg0)
+#[no_mangle]
+pub extern fn test_UnionF32F32(_: UnionF32F32) -> UnionF32F32 { loop {} }
+
+pub union UnionF32U32{a:f32, b:u32}
+
+// CHECK: define i32 @test_UnionF32U32(i32)
+#[no_mangle]
+pub extern fn test_UnionF32U32(_: UnionF32U32) -> UnionF32U32 { loop {} }
+
+pub union UnionU128{a:u128}
+// CHECK: define i128 @test_UnionU128(i128 %arg0)
+#[no_mangle]
+pub fn test_UnionU128(_: UnionU128) -> UnionU128 { loop {} }
+
+#[repr(C)]
+pub union CUnionU128{a:u128}
+// CHECK: define void @test_CUnionU128(%CUnionU128* {{.*}} %arg0)
+#[no_mangle]
+pub fn test_CUnionU128(_: CUnionU128) { loop {} }
+