Unify validity checks into a single query

Previously, there were two queries to check whether a type allows the
0x01 or zeroed bitpattern.

I am planning on adding a further initness to check, truly uninit for
MaybeUninit, which would make this three queries. This seems overkill
for such a small feature, so this PR unifies them into one.

compiler/rustc_abi/src/lib.rs

@@ -1505,14 +1505,6 @@ pub struct PointeeInfo {
     pub safe: Option<PointerKind>,
 }
 
-/// Used in `might_permit_raw_init` to indicate the kind of initialisation
-/// that is checked to be valid
-#[derive(Copy, Clone, Debug, PartialEq, Eq)]
-pub enum InitKind {
-    Zero,
-    UninitMitigated0x01Fill,
-}
-
 impl LayoutS {
     /// Returns `true` if the layout corresponds to an unsized type.
     pub fn is_unsized(&self) -> bool {

compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs

@@ -21,7 +21,8 @@ mod simd;
 pub(crate) use cpuid::codegen_cpuid_call;
 pub(crate) use llvm::codegen_llvm_intrinsic_call;
 
-use rustc_middle::ty::layout::HasParamEnv;
+use rustc_middle::ty;
+use rustc_middle::ty::layout::{HasParamEnv, InitKind};
 use rustc_middle::ty::print::with_no_trimmed_paths;
 use rustc_middle::ty::subst::SubstsRef;
 use rustc_span::symbol::{kw, sym, Symbol};
@@ -642,7 +643,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
             if intrinsic == sym::assert_zero_valid
                 && !fx
                     .tcx
-                    .permits_zero_init(fx.param_env().and(ty))
+                    .check_validity_of_init((InitKind::Zero, fx.param_env().and(ty)))
                     .expect("expected to have layout during codegen")
             {
                 with_no_trimmed_paths!({
@@ -661,7 +662,10 @@ fn codegen_regular_intrinsic_call<'tcx>(
             if intrinsic == sym::assert_mem_uninitialized_valid
                 && !fx
                     .tcx
-                    .permits_uninit_init(fx.param_env().and(ty))
+                    .check_validity_of_init((
+                        InitKind::UninitMitigated0x01Fill,
+                        fx.param_env().and(ty),
+                    ))
                     .expect("expected to have layout during codegen")
             {
                 with_no_trimmed_paths!({

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -14,7 +14,7 @@ use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
 use rustc_hir::lang_items::LangItem;
 use rustc_index::vec::Idx;
 use rustc_middle::mir::{self, AssertKind, SwitchTargets};
-use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf};
+use rustc_middle::ty::layout::{HasTyCtxt, InitKind, LayoutOf};
 use rustc_middle::ty::print::{with_no_trimmed_paths, with_no_visible_paths};
 use rustc_middle::ty::{self, Instance, Ty, TypeVisitable};
 use rustc_session::config::OptLevel;
@@ -676,11 +676,14 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 Inhabited => layout.abi.is_uninhabited(),
                 ZeroValid => !bx
                     .tcx()
-                    .permits_zero_init(bx.param_env().and(ty))
+                    .check_validity_of_init((InitKind::Zero, bx.param_env().and(ty)))
                     .expect("expected to have layout during codegen"),
                 MemUninitializedValid => !bx
                     .tcx()
-                    .permits_uninit_init(bx.param_env().and(ty))
+                    .check_validity_of_init((
+                        InitKind::UninitMitigated0x01Fill,
+                        bx.param_env().and(ty),
+                    ))
                     .expect("expected to have layout during codegen"),
             };
             Some(if do_panic {

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -11,7 +11,7 @@ use rustc_middle::mir::{
     BinOp, NonDivergingIntrinsic,
 };
 use rustc_middle::ty;
-use rustc_middle::ty::layout::LayoutOf as _;
+use rustc_middle::ty::layout::{InitKind, LayoutOf as _};
 use rustc_middle::ty::subst::SubstsRef;
 use rustc_middle::ty::{Ty, TyCtxt};
 use rustc_span::symbol::{sym, Symbol};
@@ -437,7 +437,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 if intrinsic_name == sym::assert_zero_valid {
                     let should_panic = !self
                         .tcx
-                        .permits_zero_init(self.param_env.and(ty))
+                        .check_validity_of_init((InitKind::Zero, self.param_env.and(ty)))
                         .map_err(|_| err_inval!(TooGeneric))?;
 
                     if should_panic {
@@ -454,7 +454,10 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 if intrinsic_name == sym::assert_mem_uninitialized_valid {
                     let should_panic = !self
                         .tcx
-                        .permits_uninit_init(self.param_env.and(ty))
+                        .check_validity_of_init((
+                            InitKind::UninitMitigated0x01Fill,
+                            self.param_env.and(ty),
+                        ))
                         .map_err(|_| err_inval!(TooGeneric))?;
 
                     if should_panic {

compiler/rustc_const_eval/src/lib.rs

@@ -38,7 +38,6 @@ use rustc_errors::{DiagnosticMessage, SubdiagnosticMessage};
 use rustc_macros::fluent_messages;
 use rustc_middle::ty;
 use rustc_middle::ty::query::Providers;
-use rustc_target::abi::InitKind;
 
 fluent_messages! { "../locales/en-US.ftl" }
 
@@ -62,9 +61,7 @@ pub fn provide(providers: &mut Providers) {
         let (param_env, value) = param_env_and_value.into_parts();
         const_eval::deref_mir_constant(tcx, param_env, value)
     };
-    providers.permits_uninit_init = |tcx, param_env_and_ty| {
-        util::might_permit_raw_init(tcx, param_env_and_ty, InitKind::UninitMitigated0x01Fill)
+    providers.check_validity_of_init = |tcx, (init_kind, param_env_and_ty)| {
+        util::might_permit_raw_init(tcx, init_kind, param_env_and_ty)
     };
-    providers.permits_zero_init =
-        |tcx, param_env_and_ty| util::might_permit_raw_init(tcx, param_env_and_ty, InitKind::Zero);
 }

compiler/rustc_const_eval/src/util/might_permit_raw_init.rs

@@ -1,7 +1,7 @@
-use rustc_middle::ty::layout::{LayoutCx, LayoutError, LayoutOf, TyAndLayout};
+use rustc_middle::ty::layout::{InitKind, LayoutCx, LayoutError, LayoutOf, TyAndLayout};
 use rustc_middle::ty::{ParamEnv, ParamEnvAnd, Ty, TyCtxt};
 use rustc_session::Limit;
-use rustc_target::abi::{Abi, FieldsShape, InitKind, Scalar, Variants};
+use rustc_target::abi::{Abi, FieldsShape, Scalar, Variants};
 
 use crate::const_eval::{CheckAlignment, CompileTimeInterpreter};
 use crate::interpret::{InterpCx, MemoryKind, OpTy};
@@ -20,8 +20,8 @@ use crate::interpret::{InterpCx, MemoryKind, OpTy};
 /// to the full uninit check).
 pub fn might_permit_raw_init<'tcx>(
     tcx: TyCtxt<'tcx>,
-    param_env_and_ty: ParamEnvAnd<'tcx, Ty<'tcx>>,
     kind: InitKind,
+    param_env_and_ty: ParamEnvAnd<'tcx, Ty<'tcx>>,
 ) -> Result<bool, LayoutError<'tcx>> {
     if tcx.sess.opts.unstable_opts.strict_init_checks {
         might_permit_raw_init_strict(tcx.layout_of(param_env_and_ty)?, tcx, kind)

compiler/rustc_middle/src/query/keys.rs

@@ -4,8 +4,9 @@ use crate::infer::canonical::Canonical;
 use crate::mir;
 use crate::traits;
 use crate::ty::fast_reject::SimplifiedType;
+use crate::ty::layout::{InitKind, TyAndLayout};
 use crate::ty::subst::{GenericArg, SubstsRef};
-use crate::ty::{self, layout::TyAndLayout, Ty, TyCtxt};
+use crate::ty::{self, Ty, TyCtxt};
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId, LOCAL_CRATE};
 use rustc_hir::hir_id::{HirId, OwnerId};
 use rustc_query_system::query::{DefaultCacheSelector, SingleCacheSelector, VecCacheSelector};
@@ -696,3 +697,24 @@ impl Key for HirId {
         None
     }
 }
+
+impl<'tcx> Key for (InitKind, ty::ParamEnvAnd<'tcx, Ty<'tcx>>) {
+    type CacheSelector = DefaultCacheSelector<Self>;
+
+    // Just forward to `Ty<'tcx>`
+    #[inline(always)]
+    fn query_crate_is_local(&self) -> bool {
+        true
+    }
+
+    fn default_span(&self, _: TyCtxt<'_>) -> Span {
+        DUMMY_SP
+    }
+
+    fn ty_adt_id(&self) -> Option<DefId> {
+        match self.1.value.kind() {
+            ty::Adt(adt, _) => Some(adt.did()),
+            _ => None,
+        }
+    }
+}

compiler/rustc_middle/src/query/mod.rs

@@ -2173,12 +2173,8 @@ rustc_queries! {
         separate_provide_extern
     }
 
-    query permits_uninit_init(key: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> Result<bool, ty::layout::LayoutError<'tcx>> {
-        desc { "checking to see if `{}` permits being left uninit", key.value }
-    }
-
-    query permits_zero_init(key: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> Result<bool, ty::layout::LayoutError<'tcx>> {
-        desc { "checking to see if `{}` permits being left zeroed", key.value }
+    query check_validity_of_init(key: (InitKind, ty::ParamEnvAnd<'tcx, Ty<'tcx>>)) -> Result<bool, ty::layout::LayoutError<'tcx>> {
+        desc { "checking to see if `{}` permits being left {}", key.1.value, key.0 }
     }
 
     query compare_impl_const(

compiler/rustc_middle/src/ty/layout.rs

@@ -169,6 +169,23 @@ pub const FAT_PTR_EXTRA: usize = 1;
 /// * Cranelift stores the base-2 log of the lane count in a 4 bit integer.
 pub const MAX_SIMD_LANES: u64 = 1 << 0xF;
 
+/// Used in `might_permit_raw_init` to indicate the kind of initialisation
+/// that is checked to be valid
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, HashStable)]
+pub enum InitKind {
+    Zero,
+    UninitMitigated0x01Fill,
+}
+
+impl fmt::Display for InitKind {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Zero => f.write_str("zeroed"),
+            Self::UninitMitigated0x01Fill => f.write_str("filled with 0x01"),
+        }
+    }
+}
+
 #[derive(Copy, Clone, Debug, HashStable, TyEncodable, TyDecodable)]
 pub enum LayoutError<'tcx> {
     Unknown(Ty<'tcx>),

compiler/rustc_middle/src/ty/query.rs

@@ -32,6 +32,7 @@ use crate::traits::specialization_graph;
 use crate::traits::{self, ImplSource};
 use crate::ty::context::TyCtxtFeed;
 use crate::ty::fast_reject::SimplifiedType;
+use crate::ty::layout::InitKind;
 use crate::ty::subst::{GenericArg, SubstsRef};
 use crate::ty::util::AlwaysRequiresDrop;
 use crate::ty::GeneratorDiagnosticData;

compiler/rustc_mir_transform/src/instcombine.rs

@@ -6,8 +6,8 @@ use rustc_middle::mir::{
     BinOp, Body, Constant, ConstantKind, LocalDecls, Operand, Place, ProjectionElem, Rvalue,
     SourceInfo, Statement, StatementKind, Terminator, TerminatorKind, UnOp,
 };
-use rustc_middle::ty::layout::LayoutError;
-use rustc_middle::ty::{self, ParamEnv, ParamEnvAnd, SubstsRef, Ty, TyCtxt};
+use rustc_middle::ty::layout::InitKind;
+use rustc_middle::ty::{self, ParamEnv, SubstsRef, Ty, TyCtxt};
 use rustc_span::symbol::{sym, Symbol};
 
 pub struct InstCombine;
@@ -234,16 +234,15 @@ impl<'tcx> InstCombineContext<'tcx, '_> {
         }
         let ty = substs.type_at(0);
 
-        // Check this is a foldable intrinsic before we query the layout of our generic parameter
-        let Some(assert_panics) = intrinsic_assert_panics(intrinsic_name) else { return; };
-        match assert_panics(self.tcx, self.param_env.and(ty)) {
-            // We don't know the layout, don't touch the assertion
-            Err(_) => {}
-            Ok(true) => {
+        let known_is_valid = intrinsic_assert_panics(self.tcx, self.param_env, ty, intrinsic_name);
+        match known_is_valid {
+            // We don't know the layout or it's not validity assertion at all, don't touch it
+            None => {}
+            Some(true) => {
                 // If we know the assert panics, indicate to later opts that the call diverges
                 *target = None;
             }
-            Ok(false) => {
+            Some(false) => {
                 // If we know the assert does not panic, turn the call into a Goto
                 terminator.kind = TerminatorKind::Goto { target: *target_block };
             }
@@ -252,33 +251,21 @@ impl<'tcx> InstCombineContext<'tcx, '_> {
 }
 
 fn intrinsic_assert_panics<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    param_env: ty::ParamEnv<'tcx>,
+    ty: Ty<'tcx>,
     intrinsic_name: Symbol,
-) -> Option<fn(TyCtxt<'tcx>, ParamEnvAnd<'tcx, Ty<'tcx>>) -> Result<bool, LayoutError<'tcx>>> {
-    fn inhabited_predicate<'tcx>(
-        tcx: TyCtxt<'tcx>,
-        param_env_and_ty: ParamEnvAnd<'tcx, Ty<'tcx>>,
-    ) -> Result<bool, LayoutError<'tcx>> {
-        Ok(tcx.layout_of(param_env_and_ty)?.abi.is_uninhabited())
-    }
-    fn zero_valid_predicate<'tcx>(
-        tcx: TyCtxt<'tcx>,
-        param_env_and_ty: ParamEnvAnd<'tcx, Ty<'tcx>>,
-    ) -> Result<bool, LayoutError<'tcx>> {
-        Ok(!tcx.permits_zero_init(param_env_and_ty)?)
-    }
-    fn mem_uninitialized_valid_predicate<'tcx>(
-        tcx: TyCtxt<'tcx>,
-        param_env_and_ty: ParamEnvAnd<'tcx, Ty<'tcx>>,
-    ) -> Result<bool, LayoutError<'tcx>> {
-        Ok(!tcx.permits_uninit_init(param_env_and_ty)?)
-    }
-
-    match intrinsic_name {
-        sym::assert_inhabited => Some(inhabited_predicate),
-        sym::assert_zero_valid => Some(zero_valid_predicate),
-        sym::assert_mem_uninitialized_valid => Some(mem_uninitialized_valid_predicate),
-        _ => None,
-    }
+) -> Option<bool> {
+    Some(match intrinsic_name {
+        sym::assert_inhabited => tcx.layout_of(param_env.and(ty)).ok()?.abi.is_uninhabited(),
+        sym::assert_zero_valid => {
+            !tcx.check_validity_of_init((InitKind::Zero, param_env.and(ty))).ok()?
+        }
+        sym::assert_mem_uninitialized_valid => !tcx
+            .check_validity_of_init((InitKind::UninitMitigated0x01Fill, param_env.and(ty)))
+            .ok()?,
+        _ => return None,
+    })
 }
 
 fn resolve_rust_intrinsic<'tcx>(