Add const-stability helpers

4 files changed, 45 insertions, 10 deletions

diff --git a/compiler/rustc_mir/src/transform/check_consts/mod.rs b/compiler/rustc_mir/src/transform/check_consts/mod.rs
index c1b4cb5f1a8..b49945f3f9d 100644
--- a/compiler/rustc_mir/src/transform/check_consts/mod.rs
+++ b/compiler/rustc_mir/src/transform/check_consts/mod.rs
@@ -51,6 +51,12 @@ impl ConstCx<'mir, 'tcx> {
     pub fn const_kind(&self) -> hir::ConstContext {
         self.const_kind.expect("`const_kind` must not be called on a non-const fn")
     }
+
+    pub fn is_const_stable_const_fn(&self) -> bool {
+        self.const_kind == Some(hir::ConstContext::ConstFn)
+            && self.tcx.features().staged_api
+            && is_const_stable(self.tcx, self.def_id.to_def_id())
+    }
 }
 
 /// Returns `true` if this `DefId` points to one of the official `panic` lang items.
@@ -63,3 +69,35 @@ pub fn allow_internal_unstable(tcx: TyCtxt<'tcx>, def_id: DefId, feature_gate: S
     attr::allow_internal_unstable(&tcx.sess, attrs)
         .map_or(false, |mut features| features.any(|name| name == feature_gate))
 }
+
+// Returns `true` if the given `const fn` is "const-stable".
+//
+// Const-stability is only relevant for `const fn` within a `staged_api` crate. Only "const-stable"
+// functions can be called in a const-context by users of the stable compiler. "const-stable"
+// functions are subject to more stringent restrictions than "const-unstable" functions: They
+// cannot use unstable features and can only call other "const-stable" functions.
+pub fn is_const_stable(tcx: TyCtxt<'tcx>, def_id: DefId) -> bool {
+    use attr::{ConstStability, Stability, StabilityLevel};
+
+    // Const-stability is only relevant for `const fn`.
+    assert!(tcx.is_const_fn_raw(def_id));
+
+    // Functions with `#[rustc_const_unstable]` are const-unstable.
+    match tcx.lookup_const_stability(def_id) {
+        Some(ConstStability { level: StabilityLevel::Unstable { .. }, .. }) => return false,
+        Some(ConstStability { level: StabilityLevel::Stable { .. }, .. }) => return true,
+        None => {}
+    }
+
+    // Functions with `#[unstable]` are const-unstable.
+    //
+    // FIXME(ecstaticmorse): We should keep const-stability attributes wholly separate from normal stability
+    // attributes. `#[unstable]` should be irrelevant.
+    if let Some(Stability { level: StabilityLevel::Unstable { .. }, .. }) =
+        tcx.lookup_stability(def_id)
+    {
+        return false;
+    }
+
+    true
+}
diff --git a/compiler/rustc_mir/src/transform/check_consts/ops.rs b/compiler/rustc_mir/src/transform/check_consts/ops.rs
index 61171c87036..2dfd7e2171c 100644
--- a/compiler/rustc_mir/src/transform/check_consts/ops.rs
+++ b/compiler/rustc_mir/src/transform/check_consts/ops.rs
@@ -18,9 +18,7 @@ pub fn non_const<O: NonConstOp>(ccx: &ConstCx<'_, '_>, op: O, span: Span) {
         Status::Allowed => return,
 
         Status::Unstable(gate) if ccx.tcx.features().enabled(gate) => {
-            let unstable_in_stable = ccx.const_kind() == hir::ConstContext::ConstFn
-                && ccx.tcx.features().enabled(sym::staged_api)
-                && !ccx.tcx.has_attr(ccx.def_id.to_def_id(), sym::rustc_const_unstable)
+            let unstable_in_stable = ccx.is_const_stable_const_fn()
                 && !super::allow_internal_unstable(ccx.tcx, ccx.def_id.to_def_id(), gate);
 
             if unstable_in_stable {
diff --git a/compiler/rustc_mir/src/transform/check_consts/post_drop_elaboration.rs b/compiler/rustc_mir/src/transform/check_consts/post_drop_elaboration.rs
index 0228f2d7de0..e34511211c6 100644
--- a/compiler/rustc_mir/src/transform/check_consts/post_drop_elaboration.rs
+++ b/compiler/rustc_mir/src/transform/check_consts/post_drop_elaboration.rs
@@ -11,13 +11,13 @@ use super::ConstCx;
 
 /// Returns `true` if we should use the more precise live drop checker that runs after drop
 /// elaboration.
-pub fn checking_enabled(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> bool {
+pub fn checking_enabled(ccx: &ConstCx<'_, '_>) -> bool {
     // Const-stable functions must always use the stable live drop checker.
-    if tcx.features().staged_api && !tcx.has_attr(def_id.to_def_id(), sym::rustc_const_unstable) {
+    if ccx.is_const_stable_const_fn() {
         return false;
     }
 
-    tcx.features().const_precise_live_drops
+    ccx.tcx.features().const_precise_live_drops
 }
 
 /// Look for live drops in a const context.
@@ -30,12 +30,11 @@ pub fn check_live_drops(tcx: TyCtxt<'tcx>, def_id: LocalDefId, body: &mir::Body<
         return;
     }
 
-    if !checking_enabled(tcx, def_id) {
+    let ccx = ConstCx { body, tcx, def_id, const_kind, param_env: tcx.param_env(def_id) };
+    if !checking_enabled(&ccx) {
         return;
     }
 
-    let ccx = ConstCx { body, tcx, def_id, const_kind, param_env: tcx.param_env(def_id) };
-
     let mut visitor = CheckLiveDrops { ccx: &ccx, qualifs: Qualifs::default() };
 
     visitor.visit_body(body);
diff --git a/compiler/rustc_mir/src/transform/check_consts/validation.rs b/compiler/rustc_mir/src/transform/check_consts/validation.rs
index 0501302b761..af9d7cc1aa5 100644
--- a/compiler/rustc_mir/src/transform/check_consts/validation.rs
+++ b/compiler/rustc_mir/src/transform/check_consts/validation.rs
@@ -551,7 +551,7 @@ impl Visitor<'tcx> for Validator<'mir, 'tcx> {
             | TerminatorKind::DropAndReplace { place: dropped_place, .. } => {
                 // If we are checking live drops after drop-elaboration, don't emit duplicate
                 // errors here.
-                if super::post_drop_elaboration::checking_enabled(self.tcx, self.def_id) {
+                if super::post_drop_elaboration::checking_enabled(self.ccx) {
                     return;
                 }