add `naive_layout_of` query

5 files changed, 263 insertions, 20 deletions

diff --git a/compiler/rustc_middle/src/query/erase.rs b/compiler/rustc_middle/src/query/erase.rs
index 2c481745d98..a0cb23b5a4c 100644
--- a/compiler/rustc_middle/src/query/erase.rs
+++ b/compiler/rustc_middle/src/query/erase.rs
@@ -111,6 +111,11 @@ impl EraseType
     >()];
 }
 
+impl EraseType for Result<ty::layout::TyAndNaiveLayout<'_>, &ty::layout::LayoutError<'_>> {
+    type Result =
+        [u8; size_of::<Result<ty::layout::TyAndNaiveLayout<'_>, &ty::layout::LayoutError<'_>>>()];
+}
+
 impl EraseType for Result<ty::Const<'_>, mir::interpret::LitToConstError> {
     type Result = [u8; size_of::<Result<ty::Const<'static>, mir::interpret::LitToConstError>>()];
 }
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index b36f0df78f1..c728cc0b39f 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -1394,6 +1394,17 @@ rustc_queries! {
         desc { "computing layout of `{}`", key.value }
     }
 
+    /// Computes the naive layout estimate of a type. Note that this implicitly
+    /// executes in "reveal all" mode, and will normalize the input type.
+    ///
+    /// Unlike `layout_of`, this doesn't recurse behind reference types.
+    query naive_layout_of(
+        key: ty::ParamEnvAnd<'tcx, Ty<'tcx>>
+    ) -> Result<ty::layout::TyAndNaiveLayout<'tcx>, &'tcx ty::layout::LayoutError<'tcx>> {
+        depth_limit
+        desc { "computing layout (naive) of `{}`", key.value }
+    }
+
     /// Compute a `FnAbi` suitable for indirect calls, i.e. to `fn` pointers.
     ///
     /// NB: this doesn't handle virtual calls - those should use `fn_abi_of_instance`
diff --git a/compiler/rustc_middle/src/ty/layout.rs b/compiler/rustc_middle/src/ty/layout.rs
index 62805d1e8b5..31005bdd64a 100644
--- a/compiler/rustc_middle/src/ty/layout.rs
+++ b/compiler/rustc_middle/src/ty/layout.rs
@@ -621,6 +621,61 @@ impl<T, E> MaybeResult<T> for Result<T, E> {
 
 pub type TyAndLayout<'tcx> = rustc_target::abi::TyAndLayout<'tcx, Ty<'tcx>>;
 
+#[derive(Copy, Clone, Debug, HashStable)]
+pub struct TyAndNaiveLayout<'tcx> {
+    pub ty: Ty<'tcx>,
+    pub layout: NaiveLayout,
+}
+
+impl std::ops::Deref for TyAndNaiveLayout<'_> {
+    type Target = NaiveLayout;
+    fn deref(&self) -> &Self::Target {
+        &self.layout
+    }
+}
+
+impl std::ops::DerefMut for TyAndNaiveLayout<'_> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.layout
+    }
+}
+
+/// A naive underestimation of the layout of a type.
+#[derive(Copy, Clone, Debug, HashStable)]
+pub struct NaiveLayout {
+    pub min_size: Size,
+    pub min_align: Align,
+}
+
+impl NaiveLayout {
+    pub const EMPTY: Self = Self { min_size: Size::ZERO, min_align: Align::ONE };
+
+    pub fn is_underestimate_of(&self, layout: Layout<'_>) -> bool {
+        self.min_size <= layout.size() && self.min_align <= layout.align().abi
+    }
+
+    #[must_use]
+    pub fn pad_to_align(self) -> Self {
+        Self { min_size: self.min_size.align_to(self.min_align), min_align: self.min_align }
+    }
+
+    #[must_use]
+    pub fn concat<C: HasDataLayout>(&self, other: &Self, cx: &C) -> Option<Self> {
+        Some(Self {
+            min_size: self.min_size.checked_add(other.min_size, cx)?,
+            min_align: std::cmp::max(self.min_align, other.min_align),
+        })
+    }
+
+    #[must_use]
+    pub fn union(&self, other: &Self) -> Self {
+        Self {
+            min_size: std::cmp::max(self.min_size, other.min_size),
+            min_align: std::cmp::max(self.min_align, other.min_align),
+        }
+    }
+}
+
 /// Trait for contexts that want to be able to compute layouts of types.
 /// This automatically gives access to `LayoutOf`, through a blanket `impl`.
 pub trait LayoutOfHelpers<'tcx>: HasDataLayout + HasTyCtxt<'tcx> + HasParamEnv<'tcx> {
@@ -673,6 +728,18 @@ pub trait LayoutOf<'tcx>: LayoutOfHelpers<'tcx> {
                 .map_err(|err| self.handle_layout_err(*err, span, ty)),
         )
     }
+
+    /// Computes the naive layout estimate of a type. Note that this implicitly
+    /// executes in "reveal all" mode, and will normalize the input type.
+    ///
+    /// Unlike `layout_of`, this doesn't recurse behind reference types.
+    #[inline]
+    fn naive_layout_of(
+        &self,
+        ty: Ty<'tcx>,
+    ) -> Result<TyAndNaiveLayout<'tcx>, &'tcx LayoutError<'tcx>> {
+        self.tcx().naive_layout_of(self.param_env().and(ty))
+    }
 }
 
 impl<'tcx, C: LayoutOfHelpers<'tcx>> LayoutOf<'tcx> for C {}
diff --git a/compiler/rustc_query_system/src/query/job.rs b/compiler/rustc_query_system/src/query/job.rs
index d2140161f1d..a53d1fcc69e 100644
--- a/compiler/rustc_query_system/src/query/job.rs
+++ b/compiler/rustc_query_system/src/query/job.rs
@@ -176,7 +176,8 @@ impl QueryJobId {
         while let Some(id) = current_id {
             let info = query_map.get(&id).unwrap();
             // FIXME: This string comparison should probably not be done.
-            if format!("{:?}", info.query.dep_kind) == "layout_of" {
+            let query_name = format!("{:?}", info.query.dep_kind);
+            if query_name == "layout_of" || query_name == "naive_layout_of" {
                 depth += 1;
                 last_layout = Some((info.clone(), depth));
             }
diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs
index b840ff184e0..b7e0a3a53a0 100644
--- a/compiler/rustc_ty_utils/src/layout.rs
+++ b/compiler/rustc_ty_utils/src/layout.rs
@@ -5,7 +5,8 @@ use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::mir::{GeneratorLayout, GeneratorSavedLocal};
 use rustc_middle::query::Providers;
 use rustc_middle::ty::layout::{
-    IntegerExt, LayoutCx, LayoutError, LayoutOf, TyAndLayout, MAX_SIMD_LANES,
+    IntegerExt, LayoutCx, LayoutError, LayoutOf, NaiveLayout, TyAndLayout, TyAndNaiveLayout,
+    MAX_SIMD_LANES,
 };
 use rustc_middle::ty::{
     self, AdtDef, EarlyBinder, GenericArgsRef, ReprOptions, Ty, TyCtxt, TypeVisitableExt,
@@ -24,14 +25,14 @@ use crate::errors::{
 use crate::layout_sanity_check::sanity_check_layout;
 
 pub fn provide(providers: &mut Providers) {
-    *providers = Providers { layout_of, ..*providers };
+    *providers = Providers { layout_of, naive_layout_of, ..*providers };
 }
 
 #[instrument(skip(tcx, query), level = "debug")]
-fn layout_of<'tcx>(
+fn naive_layout_of<'tcx>(
     tcx: TyCtxt<'tcx>,
     query: ty::ParamEnvAnd<'tcx, Ty<'tcx>>,
-) -> Result<TyAndLayout<'tcx>, &'tcx LayoutError<'tcx>> {
+) -> Result<TyAndNaiveLayout<'tcx>, &'tcx LayoutError<'tcx>> {
     let (param_env, ty) = query.into_parts();
     debug!(?ty);
 
@@ -53,16 +54,43 @@ fn layout_of<'tcx>(
 
     if ty != unnormalized_ty {
         // Ensure this layout is also cached for the normalized type.
-        return tcx.layout_of(param_env.and(ty));
+        return tcx.naive_layout_of(param_env.and(ty));
     }
 
     let cx = LayoutCx { tcx, param_env };
+    let layout = naive_layout_of_uncached(&cx, ty)?;
+    Ok(TyAndNaiveLayout { ty, layout })
+}
 
+#[instrument(skip(tcx, query), level = "debug")]
+fn layout_of<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    query: ty::ParamEnvAnd<'tcx, Ty<'tcx>>,
+) -> Result<TyAndLayout<'tcx>, &'tcx LayoutError<'tcx>> {
+    let (param_env, unnormalized_ty) = query.into_parts();
+    let param_env = param_env.with_reveal_all_normalized(tcx);
+    // `naive_layout_of` takes care of normalizing the type.
+    let naive = tcx.naive_layout_of(query)?;
+    let ty = naive.ty;
+
+    if ty != unnormalized_ty {
+        // Ensure this layout is also cached for the normalized type.
+        return tcx.layout_of(param_env.and(ty));
+    }
+
+    let cx = LayoutCx { tcx, param_env };
     let layout = layout_of_uncached(&cx, ty)?;
-    let layout = TyAndLayout { ty, layout };
 
-    record_layout_for_printing(&cx, layout);
+    if !naive.is_underestimate_of(layout) {
+        bug!(
+            "the estimated naive layout is bigger than the actual layout:\n{:#?}\n{:#?}",
+            naive,
+            layout,
+        );
+    }
 
+    let layout = TyAndLayout { ty, layout };
+    record_layout_for_printing(&cx, layout);
     sanity_check_layout(&cx, &layout);
 
     Ok(layout)
@@ -75,6 +103,132 @@ fn error<'tcx>(
     cx.tcx.arena.alloc(err)
 }
 
+fn naive_layout_of_uncached<'tcx>(
+    cx: &LayoutCx<'tcx, TyCtxt<'tcx>>,
+    ty: Ty<'tcx>,
+) -> Result<NaiveLayout, &'tcx LayoutError<'tcx>> {
+    let tcx = cx.tcx;
+    let dl = cx.data_layout();
+
+    let scalar =
+        |value: Primitive| NaiveLayout { min_size: value.size(dl), min_align: value.align(dl).abi };
+
+    let univariant = |fields: &mut dyn Iterator<Item = Ty<'tcx>>,
+                      repr: &ReprOptions|
+     -> Result<NaiveLayout, &'tcx LayoutError<'tcx>> {
+        // For simplicity, ignore inter-field padding; this may underestimate the size.
+        // FIXME(reference_niches): Be smarter and implement something closer to the real layout logic.
+        let mut layout = NaiveLayout::EMPTY;
+        for field in fields {
+            let field = cx.naive_layout_of(field)?;
+            layout = layout
+                .concat(&field, cx)
+                .ok_or_else(|| error(cx, LayoutError::SizeOverflow(ty)))?;
+        }
+
+        if let Some(align) = repr.align {
+            layout.min_align = std::cmp::max(layout.min_align, align);
+        }
+        if let Some(pack) = repr.pack {
+            layout.min_align = std::cmp::min(layout.min_align, pack);
+        }
+
+        Ok(layout.pad_to_align())
+    };
+
+    debug_assert!(!ty.has_non_region_infer());
+
+    Ok(match *ty.kind() {
+        // Basic scalars
+        ty::Bool => scalar(Int(I8, false)),
+        ty::Char => scalar(Int(I32, false)),
+        ty::Int(ity) => scalar(Int(Integer::from_int_ty(dl, ity), true)),
+        ty::Uint(ity) => scalar(Int(Integer::from_uint_ty(dl, ity), false)),
+        ty::Float(fty) => scalar(match fty {
+            ty::FloatTy::F32 => F32,
+            ty::FloatTy::F64 => F64,
+        }),
+        ty::FnPtr(_) => scalar(Pointer(dl.instruction_address_space)),
+
+        // The never type.
+        ty::Never => NaiveLayout::EMPTY,
+
+        // Potentially-wide pointers.
+        ty::Ref(_, _, _) | ty::RawPtr(_) => {
+            // TODO(reference_niches): handle wide pointers
+            scalar(Pointer(AddressSpace::DATA))
+        }
+
+        ty::Dynamic(_, _, ty::DynStar) => {
+            let ptr = scalar(Pointer(AddressSpace::DATA));
+            ptr.concat(&ptr, cx).ok_or_else(|| error(cx, LayoutError::SizeOverflow(ty)))?
+        }
+
+        // Arrays and slices.
+        ty::Array(element, _count) => {
+            let element = cx.naive_layout_of(element)?;
+            NaiveLayout {
+                min_size: Size::ZERO, // TODO(reference_niches): proper array size
+                min_align: element.min_align,
+            }
+        }
+        ty::Slice(element) => {
+            NaiveLayout { min_size: Size::ZERO, min_align: cx.naive_layout_of(element)?.min_align }
+        }
+        ty::Str => NaiveLayout::EMPTY,
+
+        // Odd unit types.
+        ty::FnDef(..) | ty::Dynamic(_, _, ty::Dyn) | ty::Foreign(..) => NaiveLayout::EMPTY,
+
+        // FIXME(reference_niches): try to actually compute a reasonable layout estimate,
+        // without duplicating too much code from `generator_layout`.
+        ty::Generator(..) => NaiveLayout::EMPTY,
+
+        ty::Closure(_, ref substs) => {
+            univariant(&mut substs.as_closure().upvar_tys(), &ReprOptions::default())?
+        }
+
+        ty::Tuple(tys) => univariant(&mut tys.iter(), &ReprOptions::default())?,
+
+        ty::Adt(def, substs) if def.is_union() => {
+            let repr = def.repr();
+            let only_variant = &def.variants()[FIRST_VARIANT];
+            only_variant.fields.iter().try_fold(NaiveLayout::EMPTY, |layout, f| {
+                let mut fields = std::iter::once(f.ty(tcx, substs));
+                univariant(&mut fields, &repr).map(|l| layout.union(&l))
+            })?
+        }
+
+        ty::Adt(def, substs) => {
+            // For simplicity, assume that any discriminant field (if it exists)
+            // gets niched inside one of the variants; this will underestimate the size
+            // (and sometimes alignment) of enums.
+            // FIXME(reference_niches): Be smarter and actually take into accoount the discriminant.
+            let repr = def.repr();
+            def.variants().iter().try_fold(NaiveLayout::EMPTY, |layout, v| {
+                let mut fields = v.fields.iter().map(|f| f.ty(tcx, substs));
+                let vlayout = univariant(&mut fields, &repr)?;
+                Ok(layout.union(&vlayout))
+            })?
+        }
+
+        // Types with no meaningful known layout.
+        ty::Alias(..) => {
+            // NOTE(eddyb) `layout_of` query should've normalized these away,
+            // if that was possible, so there's no reason to try again here.
+            return Err(error(cx, LayoutError::Unknown(ty)));
+        }
+
+        ty::Bound(..) | ty::GeneratorWitness(..) | ty::GeneratorWitnessMIR(..) | ty::Infer(_) => {
+            bug!("Layout::compute: unexpected type `{}`", ty)
+        }
+
+        ty::Placeholder(..) | ty::Param(_) | ty::Error(_) => {
+            return Err(error(cx, LayoutError::Unknown(ty)));
+        }
+    })
+}
+
 fn univariant_uninterned<'tcx>(
     cx: &LayoutCx<'tcx, TyCtxt<'tcx>>,
     ty: Ty<'tcx>,
@@ -146,6 +300,14 @@ 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(_) => (),
+                    // This can happen when computing the `SizeSkeleton` of a generic type.
+                    Err(LayoutError::Unknown(_)) => (),
+                    Err(err) => return Err(err),
+                }
                 data_ptr.valid_range_mut().start = 1;
             }
 
@@ -558,18 +720,15 @@ fn layout_of_uncached<'tcx>(
         }
 
         // Types with no meaningful known layout.
-        ty::Alias(..) => {
-            // NOTE(eddyb) `layout_of` query should've normalized these away,
-            // if that was possible, so there's no reason to try again here.
-            return Err(error(cx, LayoutError::Unknown(ty)));
-        }
-
-        ty::Bound(..) | ty::GeneratorWitness(..) | ty::GeneratorWitnessMIR(..) | ty::Infer(_) => {
-            bug!("Layout::compute: unexpected type `{}`", ty)
-        }
-
-        ty::Placeholder(..) | ty::Param(_) | ty::Error(_) => {
-            return Err(error(cx, LayoutError::Unknown(ty)));
+        ty::Alias(..)
+        | ty::Bound(..)
+        | ty::GeneratorWitness(..)
+        | ty::GeneratorWitnessMIR(..)
+        | ty::Infer(_)
+        | ty::Placeholder(..)
+        | ty::Param(_)
+        | ty::Error(_) => {
+            unreachable!("already rejected by `naive_layout_of`");
         }
     })
 }