Use ScalarPair for tagged enums

4 files changed, 90 insertions, 14 deletions

diff --git a/src/librustc/ty/layout.rs b/src/librustc/ty/layout.rs
index 77e2e9447f1..cfed0839acb 100644
--- a/src/librustc/ty/layout.rs
+++ b/src/librustc/ty/layout.rs
@@ -1622,7 +1622,7 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                 }
 
                 // Create the set of structs that represent each variant.
-                let mut variants = variants.into_iter().enumerate().map(|(i, field_layouts)| {
+                let mut layout_variants = variants.iter().enumerate().map(|(i, field_layouts)| {
                     let mut st = univariant_uninterned(&field_layouts,
                         &def.repr, StructKind::Prefixed(min_ity.size(), prefix_align))?;
                     st.variants = Variants::Single { index: i };
@@ -1683,7 +1683,7 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                     // Patch up the variants' first few fields.
                     let old_ity_size = min_ity.size();
                     let new_ity_size = ity.size();
-                    for variant in &mut variants {
+                    for variant in &mut layout_variants {
                         if variant.abi == Abi::Uninhabited {
                             continue;
                         }
@@ -1710,15 +1710,80 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                     value: Int(ity, signed),
                     valid_range: (min as u128 & tag_mask)..=(max as u128 & tag_mask),
                 };
-                let abi = if tag.value.size(dl) == size {
-                    Abi::Scalar(tag.clone())
-                } else {
-                    Abi::Aggregate { sized: true }
-                };
+                let mut abi = Abi::Aggregate { sized: true };
+                if tag.value.size(dl) == size {
+                    abi = Abi::Scalar(tag.clone());
+                } else if !tag.is_bool() {
+                    // HACK(nox): Blindly using ScalarPair for all tagged enums
+                    // where applicable leads to Option<u8> being handled as {i1, i8},
+                    // which later confuses SROA and some loop optimisations,
+                    // ultimately leading to the repeat-trusted-len test
+                    // failing. We make the trade-off of using ScalarPair only
+                    // for types where the tag isn't a boolean.
+                    let mut common_prim = None;
+                    for (field_layouts, layout_variant) in variants.iter().zip(&layout_variants) {
+                        let offsets = match layout_variant.fields {
+                            FieldPlacement::Arbitrary { ref offsets, .. } => offsets,
+                            _ => bug!(),
+                        };
+                        let mut fields = field_layouts
+                            .iter()
+                            .zip(offsets)
+                            .filter(|p| !p.0.is_zst());
+                        let (field, offset) = match (fields.next(), fields.next()) {
+                            (None, None) => continue,
+                            (Some(pair), None) => pair,
+                            _ => {
+                                common_prim = None;
+                                break;
+                            }
+                        };
+                        let prim = match field.details.abi {
+                            Abi::Scalar(ref scalar) => scalar.value,
+                            _ => {
+                                common_prim = None;
+                                break;
+                            }
+                        };
+                        if let Some(pair) = common_prim {
+                            // This is pretty conservative. We could go fancier
+                            // by conflating things like i32 and u32, or even
+                            // realising that (u8, u8) could just cohabit with
+                            // u16 or even u32.
+                            if pair != (prim, offset) {
+                                common_prim = None;
+                                break;
+                            }
+                        } else {
+                            common_prim = Some((prim, offset));
+                        }
+                    }
+                    if let Some((prim, offset)) = common_prim {
+                        let pair = scalar_pair(tag.clone(), scalar_unit(prim));
+                        let pair_offsets = match pair.fields {
+                            FieldPlacement::Arbitrary {
+                                ref offsets,
+                                ref memory_index
+                            } => {
+                                assert_eq!(memory_index, &[0, 1]);
+                                offsets
+                            }
+                            _ => bug!()
+                        };
+                        if pair_offsets[0] == Size::from_bytes(0) &&
+                            pair_offsets[1] == *offset &&
+                            align == pair.align &&
+                            size == pair.size {
+                            // We can use `ScalarPair` only when it matches our
+                            // already computed layout (including `#[repr(C)]`).
+                            abi = pair.abi;
+                        }
+                    }
+                }
                 tcx.intern_layout(LayoutDetails {
                     variants: Variants::Tagged {
                         discr: tag,
-                        variants
+                        variants: layout_variants,
                     },
                     fields: FieldPlacement::Arbitrary {
                         offsets: vec![Size::from_bytes(0)],
diff --git a/src/test/codegen/align-struct.rs b/src/test/codegen/align-struct.rs
index 155319cb154..f306608f432 100644
--- a/src/test/codegen/align-struct.rs
+++ b/src/test/codegen/align-struct.rs
@@ -29,7 +29,6 @@ pub enum Enum4 {
     A(i32),
     B(i32),
 }
-// CHECK: %Enum4 = type { [0 x i32], i32, [1 x i32] }
 // CHECK: %"Enum4::A" = type { [1 x i32], i32, [0 x i32] }
 
 pub enum Enum64 {
@@ -59,7 +58,7 @@ pub fn nested64(a: Align64, b: i32, c: i32, d: i8) -> Nested64 {
 // CHECK-LABEL: @enum4
 #[no_mangle]
 pub fn enum4(a: i32) -> Enum4 {
-// CHECK: %e4 = alloca %Enum4, align 4
+// CHECK: %e4 = alloca { i32, i32 }, align 4
     let e4 = Enum4::A(a);
     e4
 }
diff --git a/src/test/codegen/function-arguments.rs b/src/test/codegen/function-arguments.rs
index de302c69056..40a9ea5a181 100644
--- a/src/test/codegen/function-arguments.rs
+++ b/src/test/codegen/function-arguments.rs
@@ -145,6 +145,18 @@ pub fn return_slice(x: &[u16]) -> &[u16] {
   x
 }
 
+// CHECK: { i16, i16 } @enum_id_1(i16 %x.0, i16 %x.1)
+#[no_mangle]
+pub fn enum_id_1(x: Option<Result<u16, u16>>) -> Option<Result<u16, u16>> {
+  x
+}
+
+// CHECK: i16 @enum_id_2(i16)
+#[no_mangle]
+pub fn enum_id_2(x: Option<u8>) -> Option<u8> {
+  x
+}
+
 // CHECK: noalias i8* @allocator()
 #[no_mangle]
 #[allocator]
diff --git a/src/test/codegen/lifetime_start_end.rs b/src/test/codegen/lifetime_start_end.rs
index 62aa93398ac..ea3f0de5d08 100644
--- a/src/test/codegen/lifetime_start_end.rs
+++ b/src/test/codegen/lifetime_start_end.rs
@@ -25,16 +25,16 @@ pub fn test() {
         let b = &Some(a);
         &b; // keep variable in an alloca
 
-// CHECK: [[S_b:%[0-9]+]] = bitcast %"core::option::Option<i32>"** %b to i8*
+// CHECK: [[S_b:%[0-9]+]] = bitcast { i32, i32 }** %b to i8*
 // CHECK: call void @llvm.lifetime.start{{.*}}(i{{[0-9 ]+}}, i8* [[S_b]])
 
-// CHECK: [[S__4:%[0-9]+]] = bitcast %"core::option::Option<i32>"* %_4 to i8*
+// CHECK: [[S__4:%[0-9]+]] = bitcast { i32, i32 }* %_4 to i8*
 // CHECK: call void @llvm.lifetime.start{{.*}}(i{{[0-9 ]+}}, i8* [[S__4]])
 
-// CHECK: [[E_b:%[0-9]+]] = bitcast %"core::option::Option<i32>"** %b to i8*
+// CHECK: [[E_b:%[0-9]+]] = bitcast { i32, i32 }** %b to i8*
 // CHECK: call void @llvm.lifetime.end{{.*}}(i{{[0-9 ]+}}, i8* [[E_b]])
 
-// CHECK: [[E__4:%[0-9]+]] = bitcast %"core::option::Option<i32>"* %_4 to i8*
+// CHECK: [[E__4:%[0-9]+]] = bitcast { i32, i32 }* %_4 to i8*
 // CHECK: call void @llvm.lifetime.end{{.*}}(i{{[0-9 ]+}}, i8* [[E__4]])
     }