pointee_info_at: fix logic for recursing into enums

4 files changed, 61 insertions, 20 deletions

diff --git a/compiler/rustc_abi/src/lib.rs b/compiler/rustc_abi/src/lib.rs
index ec758785173..7b6abdf1ea9 100644
--- a/compiler/rustc_abi/src/lib.rs
+++ b/compiler/rustc_abi/src/lib.rs
@@ -1743,15 +1743,23 @@ pub enum PointerKind {
     Box { unpin: bool, global: bool },
 }
 
-/// Note that this information is advisory only, and backends are free to ignore it.
-/// It can only be used to encode potential optimizations, but no critical information.
+/// Encodes extra information we have about a pointer.
+/// Note that this information is advisory only, and backends are free to ignore it:
+/// if the information is wrong, that can cause UB, but if the information is absent,
+/// that must always be okay.
 #[derive(Copy, Clone, Debug)]
 pub struct PointeeInfo {
-    pub size: Size,
-    pub align: Align,
     /// If this is `None`, then this is a raw pointer, so size and alignment are not guaranteed to
     /// be reliable.
     pub safe: Option<PointerKind>,
+    /// If `safe` is `Some`, then the pointer is either null or dereferenceable for this many bytes.
+    /// On a function argument, "dereferenceable" here means "dereferenceable for the entire duration
+    /// of this function call", i.e. it is UB for the memory that this pointer points to to be freed
+    /// while this function is still running.
+    /// The size can be zero if the pointer is not dereferenceable.
+    pub size: Size,
+    /// If `safe` is `Some`, then the pointer is aligned as indicated.
+    pub align: Align,
 }
 
 impl<FieldIdx: Idx, VariantIdx: Idx> LayoutData<FieldIdx, VariantIdx> {
diff --git a/compiler/rustc_middle/src/ty/layout.rs b/compiler/rustc_middle/src/ty/layout.rs
index 76e3183fcbb..a0eb9029319 100644
--- a/compiler/rustc_middle/src/ty/layout.rs
+++ b/compiler/rustc_middle/src/ty/layout.rs
@@ -1011,25 +1011,41 @@ where
             }
 
             _ => {
-                let mut data_variant = match this.variants {
+                let mut data_variant = match &this.variants {
                     // Within the discriminant field, only the niche itself is
                     // always initialized, so we only check for a pointer at its
                     // offset.
                     //
-                    // If the niche is a pointer, it's either valid (according
-                    // to its type), or null (which the niche field's scalar
-                    // validity range encodes). This allows using
-                    // `dereferenceable_or_null` for e.g., `Option<&T>`, and
-                    // this will continue to work as long as we don't start
-                    // using more niches than just null (e.g., the first page of
-                    // the address space, or unaligned pointers).
+                    // Our goal here is to check whether this represents a
+                    // "dereferenceable or null" pointer, so we need to ensure
+                    // that there is only one other variant, and it must be null.
+                    // Below, we will then check whether the pointer is indeed
+                    // dereferenceable.
                     Variants::Multiple {
-                        tag_encoding: TagEncoding::Niche { untagged_variant, .. },
+                        tag_encoding:
+                            TagEncoding::Niche { untagged_variant, niche_variants, niche_start },
                         tag_field,
+                        variants,
                         ..
-                    } if this.fields.offset(tag_field) == offset => {
-                        Some(this.for_variant(cx, untagged_variant))
+                    } if variants.len() == 2 && this.fields.offset(*tag_field) == offset => {
+                        let tagged_variant = if untagged_variant.as_u32() == 0 {
+                            VariantIdx::from_u32(1)
+                        } else {
+                            VariantIdx::from_u32(0)
+                        };
+                        assert_eq!(tagged_variant, *niche_variants.start());
+                        if *niche_start == 0 {
+                            // The other variant is encoded as "null", so we can recurse searching for
+                            // a pointer here. This relies on the fact that the codegen backend
+                            // only adds "dereferenceable" if there's also a "nonnull" proof,
+                            // and that null is aligned for all alignments so it's okay to forward
+                            // the pointer's alignment.
+                            Some(this.for_variant(cx, *untagged_variant))
+                        } else {
+                            None
+                        }
                     }
+                    Variants::Multiple { .. } => None,
                     _ => Some(this),
                 };
 
diff --git a/compiler/rustc_target/src/callconv/mod.rs b/compiler/rustc_target/src/callconv/mod.rs
index 8c3df9c426b..aa639f1624f 100644
--- a/compiler/rustc_target/src/callconv/mod.rs
+++ b/compiler/rustc_target/src/callconv/mod.rs
@@ -143,7 +143,8 @@ pub struct ArgAttributes {
     pub regular: ArgAttribute,
     pub arg_ext: ArgExtension,
     /// The minimum size of the pointee, guaranteed to be valid for the duration of the whole call
-    /// (corresponding to LLVM's dereferenceable and dereferenceable_or_null attributes).
+    /// (corresponding to LLVM's dereferenceable_or_null attributes, i.e., it is okay for this to be
+    /// set on a null pointer, but all non-null pointers must be dereferenceable).
     pub pointee_size: Size,
     pub pointee_align: Option<Align>,
 }
diff --git a/tests/codegen/function-arguments.rs b/tests/codegen/function-arguments.rs
index 7fa1d659885..503799d3ed2 100644
--- a/tests/codegen/function-arguments.rs
+++ b/tests/codegen/function-arguments.rs
@@ -1,5 +1,6 @@
 //@ compile-flags: -O -C no-prepopulate-passes
 #![crate_type = "lib"]
+#![feature(rustc_attrs)]
 #![feature(dyn_star)]
 #![feature(allocator_api)]
 
@@ -143,13 +144,28 @@ pub fn indirect_struct(_: S) {}
 #[no_mangle]
 pub fn borrowed_struct(_: &S) {}
 
-// CHECK: @option_borrow(ptr noalias noundef readonly align 4 dereferenceable_or_null(4) %x)
+// CHECK: @option_borrow(ptr noalias noundef readonly align 4 dereferenceable_or_null(4) %_x)
 #[no_mangle]
-pub fn option_borrow(x: Option<&i32>) {}
+pub fn option_borrow(_x: Option<&i32>) {}
 
-// CHECK: @option_borrow_mut(ptr noalias noundef align 4 dereferenceable_or_null(4) %x)
+// CHECK: @option_borrow_mut(ptr noalias noundef align 4 dereferenceable_or_null(4) %_x)
 #[no_mangle]
-pub fn option_borrow_mut(x: Option<&mut i32>) {}
+pub fn option_borrow_mut(_x: Option<&mut i32>) {}
+
+// Function that must NOT have `dereferenceable` or `align`.
+#[rustc_layout_scalar_valid_range_start(16)]
+pub struct RestrictedAddress(&'static i16);
+enum E {
+    A(RestrictedAddress),
+    B,
+    C,
+}
+// If the `nonnull` ever goes missing, you might have to tweak the
+// scalar_valid_range on `RestrictedAddress` to get it back. You
+// might even have to add a `rustc_layout_scalar_valid_range_end`.
+// CHECK: @nonnull_and_nondereferenceable(ptr noundef nonnull %_x)
+#[no_mangle]
+pub fn nonnull_and_nondereferenceable(_x: E) {}
 
 // CHECK: @raw_struct(ptr noundef %_1)
 #[no_mangle]