Check comptime fns

oli-obk · oli-obk · commit 47a255dac34c · 2025-11-25T09:08:48.000Z
diff --git a/compiler/rustc_ast_lowering/messages.ftl b/compiler/rustc_ast_lowering/messages.ftl
@@ -50,6 +50,11 @@ ast_lowering_clobber_abi_not_supported =
 
 ast_lowering_closure_cannot_be_static = closures cannot be static
 
+ast_lowering_const_comptime_fn = a function cannot be both `comptime` and `const`
+    .suggestion = remove the `const`
+    .label = `comptime` because of this
+    .note = `const` implies the function can be called at runtime, too
+
 ast_lowering_coroutine_too_many_parameters =
     too many parameters for a coroutine (expected 0 or 1 parameters)
 
diff --git a/compiler/rustc_ast_lowering/src/errors.rs b/compiler/rustc_ast_lowering/src/errors.rs
@@ -3,6 +3,16 @@ use rustc_errors::codes::*;
 use rustc_macros::{Diagnostic, Subdiagnostic};
 use rustc_span::{Ident, Span, Symbol};
 
+#[derive(Diagnostic)]
+#[diag(ast_lowering_const_comptime_fn)]
+pub(crate) struct ConstComptimeFn {
+    #[primary_span]
+    #[suggestion(applicability = "machine-applicable", code = "")]
+    pub span: Span,
+    #[label]
+    pub attr_span: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag(ast_lowering_generic_type_with_parentheses, code = E0214)]
 pub(crate) struct GenericTypeWithParentheses {
diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
@@ -23,6 +23,7 @@ use super::{
     AstOwner, FnDeclKind, ImplTraitContext, ImplTraitPosition, LoweringContext, ParamMode,
     RelaxedBoundForbiddenReason, RelaxedBoundPolicy, ResolverAstLoweringExt,
 };
+use crate::errors::ConstComptimeFn;
 
 pub(super) struct ItemLowerer<'a, 'hir> {
     pub(super) tcx: TyCtxt<'hir>,
@@ -1537,12 +1538,23 @@ impl<'hir> LoweringContext<'_, 'hir> {
             safety.into()
         };
 
-        hir::FnHeader {
-            safety,
-            asyncness,
-            constness: self.lower_constness(h.constness),
-            abi: self.lower_extern(h.ext),
+        let mut constness = self.lower_constness(h.constness);
+        if let Some(&attr_span) = find_attr!(attrs, AttributeKind::Comptime(span) => span) {
+            match std::mem::replace(&mut constness, rustc_hir::Constness::Comptime) {
+                rustc_hir::Constness::Comptime => {
+                    unreachable!("lower_constness cannot produce comptime")
+                }
+                // A function can't be `const` and `comptime` at the same time
+                rustc_hir::Constness::Const => {
+                    let Const::Yes(span) = h.constness else { unreachable!() };
+                    self.dcx().emit_err(ConstComptimeFn { span, attr_span });
+                }
+                // Good
+                rustc_hir::Constness::NotConst => {}
+            }
         }
+
+        hir::FnHeader { safety, asyncness, constness, abi: self.lower_extern(h.ext) }
     }
 
     pub(super) fn lower_abi(&mut self, abi_str: StrLit) -> ExternAbi {
diff --git a/compiler/rustc_attr_parsing/src/attributes/semantics.rs b/compiler/rustc_attr_parsing/src/attributes/semantics.rs
@@ -7,3 +7,14 @@ impl<S: Stage> NoArgsAttributeParser<S> for MayDangleParser {
     const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(ALL_TARGETS); //FIXME Still checked fully in `check_attr.rs`
     const CREATE: fn(span: Span) -> AttributeKind = AttributeKind::MayDangle;
 }
+
+pub(crate) struct ComptimeParser;
+impl<S: Stage> NoArgsAttributeParser<S> for ComptimeParser {
+    const PATH: &[Symbol] = &[sym::rustc_comptime];
+    const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Error;
+    const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(&[
+        Allow(Target::Method(MethodKind::Inherent)),
+        Allow(Target::Fn),
+    ]);
+    const CREATE: fn(Span) -> AttributeKind = AttributeKind::Comptime;
+}
diff --git a/compiler/rustc_attr_parsing/src/context.rs b/compiler/rustc_attr_parsing/src/context.rs
@@ -58,7 +58,7 @@ use crate::attributes::rustc_internal::{
     RustcLayoutScalarValidRangeEndParser, RustcLayoutScalarValidRangeStartParser, RustcMainParser,
     RustcObjectLifetimeDefaultParser, RustcSimdMonomorphizeLaneLimitParser,
 };
-use crate::attributes::semantics::MayDangleParser;
+use crate::attributes::semantics::{ComptimeParser, MayDangleParser};
 use crate::attributes::stability::{
     BodyStabilityParser, ConstStabilityIndirectParser, ConstStabilityParser, StabilityParser,
 };
@@ -216,6 +216,7 @@ attribute_parsers!(
         Single<WithoutArgs<AutomaticallyDerivedParser>>,
         Single<WithoutArgs<CoinductiveParser>>,
         Single<WithoutArgs<ColdParser>>,
+        Single<WithoutArgs<ComptimeParser>>,
         Single<WithoutArgs<ConstContinueParser>>,
         Single<WithoutArgs<ConstStabilityIndirectParser>>,
         Single<WithoutArgs<CoroutineParser>>,
diff --git a/compiler/rustc_const_eval/src/const_eval/fn_queries.rs b/compiler/rustc_const_eval/src/const_eval/fn_queries.rs
@@ -34,8 +34,8 @@ fn constness(tcx: TyCtxt<'_>, def_id: LocalDefId) -> hir::Constness {
         hir::Node::Item(i) if let hir::ItemKind::Impl(impl_) = i.kind => impl_.constness,
         _ => {
             if let Some(fn_kind) = node.fn_kind() {
-                if fn_kind.constness() == hir::Constness::Const {
-                    return hir::Constness::Const;
+                if fn_kind.constness() != hir::Constness::NotConst {
+                    return fn_kind.constness();
                 }
 
                 // If the function itself is not annotated with `const`, it may still be a `const fn`
diff --git a/compiler/rustc_feature/src/builtin_attrs.rs b/compiler/rustc_feature/src/builtin_attrs.rs
@@ -1407,6 +1407,11 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         rustc_force_inline, Normal, template!(Word, NameValueStr: "reason"), WarnFollowing, EncodeCrossCrate::Yes,
         "`#[rustc_force_inline]` forces a free function to be inlined"
     ),
+    rustc_attr!(
+        rustc_comptime, Normal, template!(Word), ErrorFollowing,
+        EncodeCrossCrate::No,
+        "the `#[rustc_comptime]` attribute is just used to avoid adding syntax for `comptime fn`"
+    ),
 
     // ==========================================================================
     // Internal attributes, Testing:
diff --git a/compiler/rustc_hir/src/attrs/data_structures.rs b/compiler/rustc_hir/src/attrs/data_structures.rs
@@ -491,6 +491,9 @@ pub enum AttributeKind {
     /// Represents `#[cold]`.
     Cold(Span),
 
+    /// Represents `#[rustc_comptime]`
+    Comptime(Span),
+
     /// Represents `#[rustc_confusables]`.
     Confusables {
         symbols: ThinVec<Symbol>,
diff --git a/compiler/rustc_hir/src/attrs/encode_cross_crate.rs b/compiler/rustc_hir/src/attrs/encode_cross_crate.rs
@@ -28,6 +28,7 @@ impl AttributeKind {
             BodyStability { .. } => No,
             Coinductive(..) => No,
             Cold(..) => No,
+            Comptime(..) => No,
             Confusables { .. } => Yes,
             ConstContinue(..) => No,
             ConstStability { .. } => Yes,
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
@@ -2137,7 +2137,10 @@ impl<'tcx> TyCtxt<'tcx> {
         matches!(
             self.def_kind(def_id),
             DefKind::Fn | DefKind::AssocFn | DefKind::Ctor(_, CtorKind::Fn) | DefKind::Closure
-        ) && self.constness(def_id) == hir::Constness::Const
+        ) && match self.constness(def_id) {
+            hir::Constness::Comptime | hir::Constness::Const => true,
+            hir::Constness::NotConst => false,
+        }
     }
 
     /// Whether this item is conditionally constant for the purposes of the
diff --git a/compiler/rustc_passes/src/check_attr.rs b/compiler/rustc_passes/src/check_attr.rs
@@ -232,6 +232,7 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                     | AttributeKind::Marker(..)
                     | AttributeKind::SkipDuringMethodDispatch { .. }
                     | AttributeKind::Coinductive(..)
+                    | AttributeKind::Comptime(..)
                     | AttributeKind::DenyExplicitImpl(..)
                     | AttributeKind::DoNotImplementViaObject(..)
                     | AttributeKind::SpecializationTrait(..)
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
@@ -1908,6 +1908,7 @@ symbols! {
         rustc_clean,
         rustc_coherence_is_core,
         rustc_coinductive,
+        rustc_comptime,
         rustc_confusables,
         rustc_const_stable,
         rustc_const_stable_indirect,
diff --git a/tests/ui/comptime/const_comptime.rs b/tests/ui/comptime/const_comptime.rs
@@ -0,0 +1,7 @@
+#![feature(rustc_attrs)]
+
+#[rustc_comptime]
+const fn foo() {}
+//~^ ERROR a function cannot be both `comptime` and `const`
+
+fn main() {}
diff --git a/tests/ui/comptime/const_comptime.stderr b/tests/ui/comptime/const_comptime.stderr
@@ -0,0 +1,10 @@
+error: a function cannot be both `comptime` and `const`
+  --> $DIR/const_comptime.rs:4:1
+   |
+LL | #[rustc_comptime]
+   | ----------------- `comptime` because of this
+LL | const fn foo() {}
+   | ^^^^^ help: remove the `const`
+
+error: aborting due to 1 previous error
+
diff --git a/tests/ui/comptime/not_callable.rs b/tests/ui/comptime/not_callable.rs
@@ -0,0 +1,25 @@
+#![feature(rustc_attrs)]
+
+//TODO: should fail
+//@check-pass
+
+#[rustc_comptime]
+fn foo() {}
+
+fn main() {
+    // Ok
+    const { foo() };
+    // TODO: Not ok
+    foo();
+}
+
+const fn bar() {
+    // TODO: Not ok
+    foo();
+}
+
+#[rustc_comptime]
+fn baz() {
+    // Should be allowed
+    foo();
+}
diff --git a/tests/ui/comptime/trait_comptime.rs b/tests/ui/comptime/trait_comptime.rs
@@ -0,0 +1,26 @@
+#![feature(rustc_attrs)]
+
+trait Foo {
+    #[rustc_comptime]
+    //~^ ERROR: cannot be used on required trait methods
+    fn foo();
+
+    #[rustc_comptime]
+    //~^ ERROR: cannot be used on provided trait methods
+    fn bar() {}
+}
+
+struct Bar;
+
+impl Bar {
+    #[rustc_comptime]
+    fn foo() {}
+}
+
+impl Foo for Bar {
+    #[rustc_comptime]
+    //~^ ERROR: cannot be used on trait methods
+    fn foo() {}
+}
+
+fn main() {}
diff --git a/tests/ui/comptime/trait_comptime.stderr b/tests/ui/comptime/trait_comptime.stderr
@@ -0,0 +1,26 @@
+error: `#[rustc_comptime]` attribute cannot be used on required trait methods
+  --> $DIR/trait_comptime.rs:4:5
+   |
+LL |     #[rustc_comptime]
+   |     ^^^^^^^^^^^^^^^^^
+   |
+   = help: `#[rustc_comptime]` can be applied to functions and inherent methods
+
+error: `#[rustc_comptime]` attribute cannot be used on provided trait methods
+  --> $DIR/trait_comptime.rs:8:5
+   |
+LL |     #[rustc_comptime]
+   |     ^^^^^^^^^^^^^^^^^
+   |
+   = help: `#[rustc_comptime]` can be applied to functions and inherent methods
+
+error: `#[rustc_comptime]` attribute cannot be used on trait methods in impl blocks
+  --> $DIR/trait_comptime.rs:21:5
+   |
+LL |     #[rustc_comptime]
+   |     ^^^^^^^^^^^^^^^^^
+   |
+   = help: `#[rustc_comptime]` can be applied to functions and inherent methods
+
+error: aborting due to 3 previous errors
+
