fix discriminant getting skipped

Author: folkertdev

compiler/rustc_abi/src/layout/ty.rs

@@ -309,7 +309,7 @@ impl<'a, Ty> TyAndLayout<'a, Ty> {
     /// These bytes contain data for at least one valid value, but are padding for at least one
     /// other valid value. In other words, whether such a byte is initialized depends on the
     /// concrete value, hence "value-dependent". This is exactly the padding that is *not*
-    /// already returned by [`Self::guaranteed_padding_ranges`].
+    /// already returned by [`Self::variant_independent_padding_ranges`].
     ///
     /// For example, `Option<i8>` has no guaranteed padding, but the byte holding the payload is
     /// value-dependent padding: it is data for `Some(_)` and padding for `None`.
@@ -374,6 +374,16 @@ impl<'a, Ty> TyAndLayout<'a, Ty> {
                 }
             },
             Variants::Multiple { variants, .. } => {
+                // The variants do not contain e.g. the discriminant or coroutine upvars.
+                let FieldsShape::Arbitrary { offsets, in_memory_order: _ } = &self.fields else {
+                    unreachable!("a multi-variant layout should have `Arbitrary` fields")
+                };
+
+                for (field, &offset) in offsets.iter_enumerated() {
+                    let field = self.field(cx, field.as_usize());
+                    field.add_data_ranges(cx, base_offset + offset, out);
+                }
+
                 for variant in variants.indices() {
                     let variant = self.for_variant(cx, variant);
                     variant.add_data_ranges(cx, base_offset, out);
@@ -441,14 +451,31 @@ impl<'a, Ty> TyAndLayout<'a, Ty> {
                 }
             },
             Variants::Multiple { variants, .. } => {
-                // A byte is data for every value only when it is data for every value of every
-                // variant.
-                out = variants
+                // The variants do not contain e.g. the discriminant or coroutine upvars.
+                let FieldsShape::Arbitrary { offsets, in_memory_order: _ } = &self.fields else {
+                    unreachable!("a multi-variant layout should have `Arbitrary` fields")
+                };
+
+                // The fields are variant-independent and always data.
+                for (field, &offset) in offsets.iter_enumerated() {
+                    let field = self.field(cx, field.as_usize());
+                    let field = field.always_data_ranges(cx, base_offset + offset);
+                    for &(offset, size) in field.0.iter() {
+                        out.add_range(offset, size);
+                    }
+                }
+
+                // Otherwise a byte is data for every value only when it is data for every variant.
+                if let Some(common) = variants
                     .indices()
                     .map(|variant| self.for_variant(cx, variant))
                     .map(|variant| variant.always_data_ranges(cx, base_offset))
                     .reduce(|acc, f| acc.intersection(&f))
-                    .unwrap_or(out);
+                {
+                    for &(offset, size) in common.0.iter() {
+                        out.add_range(offset, size);
+                    }
+                }
             }
         }
 

compiler/rustc_lint/src/cmse_uninitialized_leak.rs

@@ -37,7 +37,7 @@ declare_lint! {
     /// return values that cross the secure boundary can contain stale secure data in their
     /// padding bytes.
     ///
-    /// The compiler zeroes all variant-independent padding: bytes that are padding for all valid
+    /// The compiler clears all variant-independent padding: bytes that are padding for all valid
     /// values of the type. But enum and union values can contain variant-dependent padding: bytes
     /// that are padding for some but not all valid values of the type. For instance, `Option<u8>`
     /// has no padding when `Some(_)` but does have padding when `None`.

tests/assembly-llvm/cmse-clear-padding.rs

@@ -202,3 +202,80 @@ extern "cmse-nonsecure-entry" fn cmse_ret_with_wide_u8_uninit_tuple(
 ) -> (MaybeUninit<WideU8>, MaybeUninit<WideU8>) {
     unsafe { (mem::transmute(a), mem::transmute(b)) }
 }
+
+// Check that the tag does not get cleared.
+enum Enum {
+    A(u16),
+    B(u16),
+}
+
+// CHECK-LABEL: cmse_ret_enum:
+// CHECK: .fnstart
+// CHECK-NEXT: .save   {r7, lr}
+// CHECK-NEXT: push    {r7, lr}
+// CHECK-NEXT: .setfp  r7, sp
+// CHECK-NEXT: mov r7, sp
+// CHECK-NEXT: pop.w   {r7, lr}
+// CHECK-NEXT: mov r1, lr
+// CHECK-NEXT: mov r2, lr
+// CHECK-NEXT: mov r3, lr
+// CHECK-NEXT: mov r12, lr
+// CHECK-NEXT: msr apsr_nzcvq, lr
+// CHECK-NEXT: bxns    lr
+#[no_mangle]
+extern "cmse-nonsecure-entry" fn cmse_ret_enum(x: Enum) -> Enum {
+    x
+}
+
+// Check that variant-dependent padding is left alone.
+enum VariantPadding {
+    A(u8),
+    B(u16),
+}
+
+// CHECK-LABEL: cmse_ret_enum_with_variant_padding:
+// CHECK: .fnstart
+// CHECK-NEXT: .save   {r7, lr}
+// CHECK-NEXT: push    {r7, lr}
+// CHECK-NEXT: .setfp  r7, sp
+// CHECK-NEXT: mov r7, sp
+// CHECK-NEXT: pop.w   {r7, lr}
+// CHECK-NEXT: mov r1, lr
+// CHECK-NEXT: mov r2, lr
+// CHECK-NEXT: mov r3, lr
+// CHECK-NEXT: mov r12, lr
+// CHECK-NEXT: msr apsr_nzcvq, lr
+// CHECK-NEXT: bxns    lr
+#[no_mangle]
+extern "cmse-nonsecure-entry" fn cmse_ret_enum_with_variant_padding(
+    x: VariantPadding,
+) -> VariantPadding {
+    x
+}
+
+// Check that variant-independent padding does get cleared.
+enum IndependentPadding {
+    A(WideU8),
+    B(WideU8),
+}
+
+// CHECK-LABEL: cmse_ret_enum_with_independent_padding:
+// CHECK: .fnstart
+// CHECK-NEXT: .save   {r7, lr}
+// CHECK-NEXT: push    {r7, lr}
+// CHECK-NEXT: .setfp  r7, sp
+// CHECK-NEXT: mov r7, sp
+// CHECK-NEXT: bic r0, r0, #-16777216
+// CHECK-NEXT: pop.w   {r7, lr}
+// CHECK-NEXT: mov r1, lr
+// CHECK-NEXT: mov r2, lr
+// CHECK-NEXT: mov r3, lr
+// CHECK-NEXT: mov r12, lr
+// CHECK-NEXT: msr apsr_nzcvq, lr
+// CHECK-NEXT: bxns    lr
+#[no_mangle]
+extern "cmse-nonsecure-entry" fn cmse_ret_enum_with_independent_padding(
+    x: IndependentPadding,
+) -> IndependentPadding {
+    x
+}