Make non-power-of-two alignments a validity error in Layout

Inspired by the zulip conversation about how `Layout` should better enforce `size < isize::MAX as usize`, this uses an N-variant enum on N-bit platforms to require at the validity level that the existing invariant of "must be a power of two" is upheld.

This was MIRI can catch it, and means there's a more-specific type for `Layout` to store than just `NonZeroUsize`.

library/core/src/alloc/layout.rs

@@ -1,7 +1,6 @@
 use crate::cmp;
 use crate::fmt;
-use crate::mem;
-use crate::num::NonZeroUsize;
+use crate::mem::{self, ValidAlign};
 use crate::ptr::NonNull;
 
 // While this function is used in one place and its implementation
@@ -40,7 +39,7 @@ pub struct Layout {
     //
     // (However, we do not analogously require `align >= sizeof(void*)`,
     //  even though that is *also* a requirement of `posix_memalign`.)
-    align_: NonZeroUsize,
+    align_: ValidAlign,
 }
 
 impl Layout {
@@ -97,8 +96,8 @@ impl Layout {
     #[must_use]
     #[inline]
     pub const unsafe fn from_size_align_unchecked(size: usize, align: usize) -> Self {
-        // SAFETY: the caller must ensure that `align` is greater than zero.
-        Layout { size_: size, align_: unsafe { NonZeroUsize::new_unchecked(align) } }
+        // SAFETY: the caller must ensure that `align` is a power of two.
+        Layout { size_: size, align_: unsafe { ValidAlign::new_unchecked(align) } }
     }
 
     /// The minimum size in bytes for a memory block of this layout.
@@ -117,7 +116,7 @@ impl Layout {
                   without modifying the layout"]
     #[inline]
     pub const fn align(&self) -> usize {
-        self.align_.get()
+        self.align_.as_nonzero().get()
     }
 
     /// Constructs a `Layout` suitable for holding a value of type `T`.

library/core/src/mem/mod.rs

@@ -21,6 +21,12 @@ mod maybe_uninit;
 #[stable(feature = "maybe_uninit", since = "1.36.0")]
 pub use maybe_uninit::MaybeUninit;
 
+mod valid_align;
+// For now this type is left crate-local.  It could potentially make sense to expose
+// it publicly, as it would be a nice parameter type for methods which need to take
+// alignment as a parameter, such as `Layout::padding_needed_for`.
+pub(crate) use valid_align::ValidAlign;
+
 #[stable(feature = "rust1", since = "1.0.0")]
 #[doc(inline)]
 pub use crate::intrinsics::transmute;

library/core/src/mem/valid_align.rs

@@ -0,0 +1,240 @@
+use crate::convert::TryFrom;
+use crate::num::NonZeroUsize;
+use crate::{cmp, fmt, mem, num};
+
+/// A type storing a `usize` which is a power of two, and thus
+/// represents a possible alignment in the rust abstract machine.
+///
+/// Note that particularly large alignments, while representable in this type,
+/// are likely not to be supported by actual allocators and linkers.
+#[derive(Copy, Clone)]
+#[repr(transparent)]
+pub(crate) struct ValidAlign(ValidAlignEnum);
+
+// ValidAlign is `repr(usize)`, but via extra steps.
+const _: () = assert!(mem::size_of::<ValidAlign>() == mem::size_of::<usize>());
+const _: () = assert!(mem::align_of::<ValidAlign>() == mem::align_of::<usize>());
+
+impl ValidAlign {
+    /// Creates a `ValidAlign` from a power-of-two `usize`.
+    ///
+    /// # Safety
+    ///
+    /// `align` must be a power of two.
+    ///
+    /// Equivalently, it must be `1 << exp` for some `exp` in `0..usize::BITS`.
+    /// It must *not* be zero.
+    #[inline]
+    pub(crate) const unsafe fn new_unchecked(align: usize) -> Self {
+        debug_assert!(align.is_power_of_two());
+
+        // SAFETY: By precondition, this must be a power of two, and
+        // our variants encompass all possible powers of two.
+        unsafe { mem::transmute::<usize, ValidAlign>(align) }
+    }
+
+    #[inline]
+    pub(crate) const fn as_nonzero(self) -> NonZeroUsize {
+        // SAFETY: All the discriminants are non-zero.
+        unsafe { NonZeroUsize::new_unchecked(self.0 as usize) }
+    }
+
+    /// Returns the base 2 logarithm of the alignment.
+    ///
+    /// This is always exact, as `self` represents a power of two.
+    #[inline]
+    pub(crate) fn log2(self) -> u32 {
+        self.as_nonzero().trailing_zeros()
+    }
+}
+
+impl fmt::Debug for ValidAlign {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{:?} (1 << {:?})", self.as_nonzero(), self.log2())
+    }
+}
+
+impl TryFrom<NonZeroUsize> for ValidAlign {
+    type Error = num::TryFromIntError;
+
+    #[inline]
+    fn try_from(align: NonZeroUsize) -> Result<ValidAlign, Self::Error> {
+        if align.is_power_of_two() {
+            // SAFETY: Just checked for power-of-two
+            unsafe { Ok(ValidAlign::new_unchecked(align.get())) }
+        } else {
+            Err(num::TryFromIntError(()))
+        }
+    }
+}
+
+impl TryFrom<usize> for ValidAlign {
+    type Error = num::TryFromIntError;
+
+    #[inline]
+    fn try_from(align: usize) -> Result<ValidAlign, Self::Error> {
+        if align.is_power_of_two() {
+            // SAFETY: Just checked for power-of-two
+            unsafe { Ok(ValidAlign::new_unchecked(align)) }
+        } else {
+            Err(num::TryFromIntError(()))
+        }
+    }
+}
+
+impl cmp::Eq for ValidAlign {}
+
+impl cmp::PartialEq for ValidAlign {
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        self.as_nonzero() == other.as_nonzero()
+    }
+}
+
+impl cmp::Ord for ValidAlign {
+    #[inline]
+    fn cmp(&self, other: &Self) -> cmp::Ordering {
+        self.as_nonzero().cmp(&other.as_nonzero())
+    }
+}
+
+impl cmp::PartialOrd for ValidAlign {
+    #[inline]
+    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+#[cfg(target_pointer_width = "16")]
+type ValidAlignEnum = ValidAlignEnum16;
+#[cfg(target_pointer_width = "32")]
+type ValidAlignEnum = ValidAlignEnum32;
+#[cfg(target_pointer_width = "64")]
+type ValidAlignEnum = ValidAlignEnum64;
+
+#[derive(Copy, Clone)]
+#[repr(u16)]
+enum ValidAlignEnum16 {
+    _Align1Shl0 = 1 << 0,
+    _Align1Shl1 = 1 << 1,
+    _Align1Shl2 = 1 << 2,
+    _Align1Shl3 = 1 << 3,
+    _Align1Shl4 = 1 << 4,
+    _Align1Shl5 = 1 << 5,
+    _Align1Shl6 = 1 << 6,
+    _Align1Shl7 = 1 << 7,
+    _Align1Shl8 = 1 << 8,
+    _Align1Shl9 = 1 << 9,
+    _Align1Shl10 = 1 << 10,
+    _Align1Shl11 = 1 << 11,
+    _Align1Shl12 = 1 << 12,
+    _Align1Shl13 = 1 << 13,
+    _Align1Shl14 = 1 << 14,
+    _Align1Shl15 = 1 << 15,
+}
+
+#[derive(Copy, Clone)]
+#[repr(u32)]
+enum ValidAlignEnum32 {
+    _Align1Shl0 = 1 << 0,
+    _Align1Shl1 = 1 << 1,
+    _Align1Shl2 = 1 << 2,
+    _Align1Shl3 = 1 << 3,
+    _Align1Shl4 = 1 << 4,
+    _Align1Shl5 = 1 << 5,
+    _Align1Shl6 = 1 << 6,
+    _Align1Shl7 = 1 << 7,
+    _Align1Shl8 = 1 << 8,
+    _Align1Shl9 = 1 << 9,
+    _Align1Shl10 = 1 << 10,
+    _Align1Shl11 = 1 << 11,
+    _Align1Shl12 = 1 << 12,
+    _Align1Shl13 = 1 << 13,
+    _Align1Shl14 = 1 << 14,
+    _Align1Shl15 = 1 << 15,
+    _Align1Shl16 = 1 << 16,
+    _Align1Shl17 = 1 << 17,
+    _Align1Shl18 = 1 << 18,
+    _Align1Shl19 = 1 << 19,
+    _Align1Shl20 = 1 << 20,
+    _Align1Shl21 = 1 << 21,
+    _Align1Shl22 = 1 << 22,
+    _Align1Shl23 = 1 << 23,
+    _Align1Shl24 = 1 << 24,
+    _Align1Shl25 = 1 << 25,
+    _Align1Shl26 = 1 << 26,
+    _Align1Shl27 = 1 << 27,
+    _Align1Shl28 = 1 << 28,
+    _Align1Shl29 = 1 << 29,
+    _Align1Shl30 = 1 << 30,
+    _Align1Shl31 = 1 << 31,
+}
+
+#[derive(Copy, Clone)]
+#[repr(u64)]
+enum ValidAlignEnum64 {
+    _Align1Shl0 = 1 << 0,
+    _Align1Shl1 = 1 << 1,
+    _Align1Shl2 = 1 << 2,
+    _Align1Shl3 = 1 << 3,
+    _Align1Shl4 = 1 << 4,
+    _Align1Shl5 = 1 << 5,
+    _Align1Shl6 = 1 << 6,
+    _Align1Shl7 = 1 << 7,
+    _Align1Shl8 = 1 << 8,
+    _Align1Shl9 = 1 << 9,
+    _Align1Shl10 = 1 << 10,
+    _Align1Shl11 = 1 << 11,
+    _Align1Shl12 = 1 << 12,
+    _Align1Shl13 = 1 << 13,
+    _Align1Shl14 = 1 << 14,
+    _Align1Shl15 = 1 << 15,
+    _Align1Shl16 = 1 << 16,
+    _Align1Shl17 = 1 << 17,
+    _Align1Shl18 = 1 << 18,
+    _Align1Shl19 = 1 << 19,
+    _Align1Shl20 = 1 << 20,
+    _Align1Shl21 = 1 << 21,
+    _Align1Shl22 = 1 << 22,
+    _Align1Shl23 = 1 << 23,
+    _Align1Shl24 = 1 << 24,
+    _Align1Shl25 = 1 << 25,
+    _Align1Shl26 = 1 << 26,
+    _Align1Shl27 = 1 << 27,
+    _Align1Shl28 = 1 << 28,
+    _Align1Shl29 = 1 << 29,
+    _Align1Shl30 = 1 << 30,
+    _Align1Shl31 = 1 << 31,
+    _Align1Shl32 = 1 << 32,
+    _Align1Shl33 = 1 << 33,
+    _Align1Shl34 = 1 << 34,
+    _Align1Shl35 = 1 << 35,
+    _Align1Shl36 = 1 << 36,
+    _Align1Shl37 = 1 << 37,
+    _Align1Shl38 = 1 << 38,
+    _Align1Shl39 = 1 << 39,
+    _Align1Shl40 = 1 << 40,
+    _Align1Shl41 = 1 << 41,
+    _Align1Shl42 = 1 << 42,
+    _Align1Shl43 = 1 << 43,
+    _Align1Shl44 = 1 << 44,
+    _Align1Shl45 = 1 << 45,
+    _Align1Shl46 = 1 << 46,
+    _Align1Shl47 = 1 << 47,
+    _Align1Shl48 = 1 << 48,
+    _Align1Shl49 = 1 << 49,
+    _Align1Shl50 = 1 << 50,
+    _Align1Shl51 = 1 << 51,
+    _Align1Shl52 = 1 << 52,
+    _Align1Shl53 = 1 << 53,
+    _Align1Shl54 = 1 << 54,
+    _Align1Shl55 = 1 << 55,
+    _Align1Shl56 = 1 << 56,
+    _Align1Shl57 = 1 << 57,
+    _Align1Shl58 = 1 << 58,
+    _Align1Shl59 = 1 << 59,
+    _Align1Shl60 = 1 << 60,
+    _Align1Shl61 = 1 << 61,
+    _Align1Shl62 = 1 << 62,
+    _Align1Shl63 = 1 << 63,
+}

library/core/tests/alloc.rs

@@ -11,3 +11,21 @@ fn const_unchecked_layout() {
     assert_eq!(LAYOUT.align(), ALIGN);
     assert_eq!(Some(DANGLING), NonNull::new(ALIGN as *mut u8));
 }
+
+#[test]
+fn layout_debug_shows_log2_of_alignment() {
+    // `Debug` is not stable, but here's what it does right now
+    let layout = Layout::from_size_align(24576, 8192).unwrap();
+    let s = format!("{:?}", layout);
+    assert_eq!(s, "Layout { size_: 24576, align_: 8192 (1 << 13) }");
+}
+
+// Running this normally doesn't do much, but it's also run in Miri, which
+// will double-check that these are allowed by the validity invariants.
+#[test]
+fn layout_accepts_all_valid_alignments() {
+    for align in 0..usize::BITS {
+        let layout = Layout::from_size_align(0, 1_usize << align).unwrap();
+        assert_eq!(layout.align(), 1_usize << align);
+    }
+}

src/test/ui/consts/std/alloc.32bit.stderr

@@ -1,14 +1,25 @@
 error[E0080]: it is undefined behavior to use this value
-  --> $DIR/alloc.rs:8:1
+  --> $DIR/alloc.rs:9:1
    |
-LL | const LAYOUT_INVALID: Layout = unsafe { Layout::from_size_align_unchecked(0x1000, 0x00) };
-   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type validation failed at .align_: encountered 0, but expected something greater or equal to 1
+LL | const LAYOUT_INVALID_ZERO: Layout = unsafe { Layout::from_size_align_unchecked(0x1000, 0x00) };
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type validation failed at .align_.0.<enum-tag>: encountered 0x00000000, but expected a valid enum tag
    |
    = note: The rules on what exactly is undefined behavior aren't clear, so this check might be overzealous. Please open an issue on the rustc repository if you believe it should not be considered undefined behavior.
    = note: the raw bytes of the constant (size: 8, align: 4) {
                00 10 00 00 00 00 00 00                         │ ........
            }
 
-error: aborting due to previous error
+error[E0080]: it is undefined behavior to use this value
+  --> $DIR/alloc.rs:13:1
+   |
+LL | const LAYOUT_INVALID_THREE: Layout = unsafe { Layout::from_size_align_unchecked(9, 3) };
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type validation failed at .align_.0.<enum-tag>: encountered 0x00000003, but expected a valid enum tag
+   |
+   = note: The rules on what exactly is undefined behavior aren't clear, so this check might be overzealous. Please open an issue on the rustc repository if you believe it should not be considered undefined behavior.
+   = note: the raw bytes of the constant (size: 8, align: 4) {
+               09 00 00 00 03 00 00 00                         │ ........
+           }
+
+error: aborting due to 2 previous errors
 
 For more information about this error, try `rustc --explain E0080`.

src/test/ui/consts/std/alloc.64bit.stderr

@@ -1,14 +1,25 @@
 error[E0080]: it is undefined behavior to use this value
-  --> $DIR/alloc.rs:8:1
+  --> $DIR/alloc.rs:9:1
    |
-LL | const LAYOUT_INVALID: Layout = unsafe { Layout::from_size_align_unchecked(0x1000, 0x00) };
-   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type validation failed at .align_: encountered 0, but expected something greater or equal to 1
+LL | const LAYOUT_INVALID_ZERO: Layout = unsafe { Layout::from_size_align_unchecked(0x1000, 0x00) };
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type validation failed at .align_.0.<enum-tag>: encountered 0x0000000000000000, but expected a valid enum tag
    |
    = note: The rules on what exactly is undefined behavior aren't clear, so this check might be overzealous. Please open an issue on the rustc repository if you believe it should not be considered undefined behavior.
    = note: the raw bytes of the constant (size: 16, align: 8) {
                00 10 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
            }
 
-error: aborting due to previous error
+error[E0080]: it is undefined behavior to use this value
+  --> $DIR/alloc.rs:13:1
+   |
+LL | const LAYOUT_INVALID_THREE: Layout = unsafe { Layout::from_size_align_unchecked(9, 3) };
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type validation failed at .align_.0.<enum-tag>: encountered 0x0000000000000003, but expected a valid enum tag
+   |
+   = note: The rules on what exactly is undefined behavior aren't clear, so this check might be overzealous. Please open an issue on the rustc repository if you believe it should not be considered undefined behavior.
+   = note: the raw bytes of the constant (size: 16, align: 8) {
+               09 00 00 00 00 00 00 00 03 00 00 00 00 00 00 00 │ ................
+           }
+
+error: aborting due to 2 previous errors
 
 For more information about this error, try `rustc --explain E0080`.

src/test/ui/consts/std/alloc.rs

@@ -1,11 +1,16 @@
 // stderr-per-bitwidth
+// ignore-debug (the debug assertions change the error)
 use std::alloc::Layout;
 
 // ok
 const LAYOUT_VALID: Layout = unsafe { Layout::from_size_align_unchecked(0x1000, 0x08) };
 
 // not ok, since alignment needs to be non-zero.
-const LAYOUT_INVALID: Layout = unsafe { Layout::from_size_align_unchecked(0x1000, 0x00) };
+const LAYOUT_INVALID_ZERO: Layout = unsafe { Layout::from_size_align_unchecked(0x1000, 0x00) };
+//~^ ERROR it is undefined behavior to use this value
+
+// not ok, since alignment needs to be a power of two.
+const LAYOUT_INVALID_THREE: Layout = unsafe { Layout::from_size_align_unchecked(9, 3) };
 //~^ ERROR it is undefined behavior to use this value
 
 fn main() {}