restrict the valid range of references if `-Z reference-niches` is set

Note that this doesn't actually work at all, as many places in rustc
assume that references only have the null niche.

2 files changed, 64 insertions, 11 deletions

diff --git a/compiler/rustc_abi/src/lib.rs b/compiler/rustc_abi/src/lib.rs
index 29b3c44dcd6..0839e95723a 100644
--- a/compiler/rustc_abi/src/lib.rs
+++ b/compiler/rustc_abi/src/lib.rs
@@ -354,6 +354,31 @@ impl TargetDataLayout {
         }
     }
 
+    /// Returns the theoretical maximum address.
+    ///
+    /// Note that this doesn't take into account target-specific limitations.
+    #[inline]
+    pub fn max_address(&self) -> u64 {
+        match self.pointer_size.bits() {
+            16 => u16::MAX.into(),
+            32 => u32::MAX.into(),
+            64 => u64::MAX,
+            bits => panic!("max_address: unknown pointer bit size {}", bits),
+        }
+    }
+
+    /// Returns the (inclusive) range of possible addresses for an allocation with
+    /// the given size and alignment.
+    ///
+    /// Note that this doesn't take into account target-specific limitations.
+    #[inline]
+    pub fn address_range_for(&self, size: Size, align: Align) -> (u64, u64) {
+        let end = Size::from_bytes(self.max_address());
+        let min = align.bytes();
+        let max = (end - size).align_down_to(align).bytes();
+        (min, max)
+    }
+
     #[inline]
     pub fn vector_align(&self, vec_size: Size) -> AbiAndPrefAlign {
         for &(size, align) in &self.vector_align {
@@ -482,6 +507,12 @@ impl Size {
     }
 
     #[inline]
+    pub fn align_down_to(self, align: Align) -> Size {
+        let mask = align.bytes() - 1;
+        Size::from_bytes(self.bytes() & !mask)
+    }
+
+    #[inline]
     pub fn is_aligned(self, align: Align) -> bool {
         let mask = align.bytes() - 1;
         self.bytes() & mask == 0
diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs
index 19545a1b135..c81e76e4471 100644
--- a/compiler/rustc_ty_utils/src/layout.rs
+++ b/compiler/rustc_ty_utils/src/layout.rs
@@ -30,11 +30,13 @@ pub fn provide(providers: &mut Providers) {
 
 #[instrument(skip(tcx), level = "debug")]
 fn reference_niches_policy<'tcx>(tcx: TyCtxt<'tcx>, _: LocalCrate) -> ReferenceNichePolicy {
-    const DEFAULT: ReferenceNichePolicy = ReferenceNichePolicy { size: false, align: false };
-
-    tcx.sess.opts.unstable_opts.reference_niches.unwrap_or(DEFAULT)
+    tcx.sess.opts.unstable_opts.reference_niches.unwrap_or(DEFAULT_REF_NICHES)
 }
 
+/// The reference niche policy for builtin types, and for types in
+/// crates not specifying `-Z reference-niches`.
+const DEFAULT_REF_NICHES: ReferenceNichePolicy = ReferenceNichePolicy { size: false, align: false };
+
 #[instrument(skip(tcx, query), level = "debug")]
 fn naive_layout_of<'tcx>(
     tcx: TyCtxt<'tcx>,
@@ -163,7 +165,6 @@ fn naive_layout_of_uncached<'tcx>(
         // Potentially-wide pointers.
         ty::Ref(_, pointee, _) | ty::RawPtr(ty::TypeAndMut { ty: pointee, .. }) => {
             let data_ptr = scalar(Pointer(AddressSpace::DATA));
-
             if let Some(metadata) = ptr_metadata_scalar(cx, pointee)? {
                 // Effectively a (ptr, meta) tuple.
                 data_ptr
@@ -322,15 +323,36 @@ fn layout_of_uncached<'tcx>(
         ty::Ref(_, pointee, _) | ty::RawPtr(ty::TypeAndMut { ty: pointee, .. }) => {
             let mut data_ptr = scalar_unit(Pointer(AddressSpace::DATA));
             if !ty.is_unsafe_ptr() {
-                match cx.naive_layout_of(pointee) {
-                    // TODO(reference_niches): actually use the naive layout to set
-                    // reference niches; the query is still kept to for testing purposes.
-                    Ok(_) => (),
+                // Calling `layout_of` here would cause a query cycle for recursive types;
+                // so use a conservative estimate that doesn't look past references.
+                let naive = match cx.naive_layout_of(pointee) {
+                    Ok(n) => n.layout,
                     // This can happen when computing the `SizeSkeleton` of a generic type.
-                    Err(LayoutError::Unknown(_)) => (),
+                    Err(LayoutError::Unknown(_)) => {
+                        // TODO(reference_niches): this is *very* incorrect, but we can't
+                        // return an error here; this would break transmute checks.
+                        // We need some other solution.
+                        NaiveLayout::EMPTY
+                    }
                     Err(err) => return Err(err),
-                }
-                data_ptr.valid_range_mut().start = 1;
+                };
+
+                let niches = match *pointee.kind() {
+                    ty::FnDef(def, ..)
+                    | ty::Foreign(def)
+                    | ty::Generator(def, ..)
+                    | ty::Closure(def, ..) => tcx.reference_niches_policy(def.krate),
+                    ty::Adt(def, _) => tcx.reference_niches_policy(def.did().krate),
+                    _ => DEFAULT_REF_NICHES,
+                };
+
+                let (min_addr, max_addr) = dl.address_range_for(
+                    if niches.size { naive.min_size } else { Size::ZERO },
+                    if niches.align { naive.min_align } else { Align::ONE },
+                );
+
+                *data_ptr.valid_range_mut() =
+                    WrappingRange { start: min_addr.into(), end: max_addr.into() };
             }
 
             if let Some(metadata) = ptr_metadata_scalar(cx, pointee)? {