Provide a way for custom derives to know if they were invoked via #[derive_const]

compiler/rustc_expand/src/proc_macro.rs

@@ -59,7 +59,7 @@ impl base::BangProcMacro for BangProcMacro {
 
         let proc_macro_backtrace = ecx.ecfg.proc_macro_backtrace;
         let strategy = exec_strategy(ecx);
-        let server = proc_macro_server::Rustc::new(ecx);
+        let server = proc_macro_server::Rustc::new(ecx, Default::default());
         self.client.run(&strategy, server, input, proc_macro_backtrace).map_err(|e| {
             ecx.dcx().emit_err(errors::ProcMacroPanicked {
                 span,
@@ -90,7 +90,7 @@ impl base::AttrProcMacro for AttrProcMacro {
 
         let proc_macro_backtrace = ecx.ecfg.proc_macro_backtrace;
         let strategy = exec_strategy(ecx);
-        let server = proc_macro_server::Rustc::new(ecx);
+        let server = proc_macro_server::Rustc::new(ecx, Default::default());
         self.client.run(&strategy, server, annotation, annotated, proc_macro_backtrace).map_err(
             |e| {
                 let mut err = ecx.dcx().struct_span_err(span, "custom attribute panicked");
@@ -114,7 +114,7 @@ impl MultiItemModifier for DeriveProcMacro {
         span: Span,
         _meta_item: &ast::MetaItem,
         item: Annotatable,
-        _is_derive_const: bool,
+        is_derive_const: bool,
     ) -> ExpandResult<Vec<Annotatable>, Annotatable> {
         // We need special handling for statement items
         // (e.g. `fn foo() { #[derive(Debug)] struct Bar; }`)
@@ -142,7 +142,10 @@ impl MultiItemModifier for DeriveProcMacro {
                 });
             let proc_macro_backtrace = ecx.ecfg.proc_macro_backtrace;
             let strategy = exec_strategy(ecx);
-            let server = proc_macro_server::Rustc::new(ecx);
+            let server = proc_macro_server::Rustc::new(
+                ecx,
+                proc_macro_server::ExpansionOptions { is_derive_const },
+            );
             match self.client.run(&strategy, server, input, proc_macro_backtrace) {
                 Ok(stream) => stream,
                 Err(e) => {

compiler/rustc_expand/src/proc_macro_server.rs

@@ -397,17 +397,19 @@ pub(crate) struct Rustc<'a, 'b> {
     mixed_site: Span,
     krate: CrateNum,
     rebased_spans: FxHashMap<usize, Span>,
+    options: ExpansionOptions,
 }
 
 impl<'a, 'b> Rustc<'a, 'b> {
-    pub fn new(ecx: &'a mut ExtCtxt<'b>) -> Self {
+    pub fn new(ecx: &'a mut ExtCtxt<'b>, options: ExpansionOptions) -> Self {
         let expn_data = ecx.current_expansion.id.expn_data();
         Rustc {
             def_site: ecx.with_def_site_ctxt(expn_data.def_site),
             call_site: ecx.with_call_site_ctxt(expn_data.call_site),
             mixed_site: ecx.with_mixed_site_ctxt(expn_data.call_site),
             krate: expn_data.macro_def_id.unwrap().krate,
             rebased_spans: FxHashMap::default(),
+            options,
             ecx,
         }
     }
@@ -417,6 +419,11 @@ impl<'a, 'b> Rustc<'a, 'b> {
     }
 }
 
+#[derive(Default)]
+pub(crate) struct ExpansionOptions {
+    pub(crate) is_derive_const: bool,
+}
+
 impl server::Types for Rustc<'_, '_> {
     type FreeFunctions = FreeFunctions;
     type TokenStream = TokenStream;
@@ -503,6 +510,10 @@ impl server::FreeFunctions for Rustc<'_, '_> {
         }
         self.sess().dcx.emit_diagnostic(diag);
     }
+
+    fn is_derive_const(&mut self) -> bool {
+        self.options.is_derive_const
+    }
 }
 
 impl server::TokenStream for Rustc<'_, '_> {

compiler/rustc_passes/src/check_attr.rs

@@ -2293,10 +2293,17 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
         }
 
         let tcx = self.tcx;
-        let Some(token_stream_def_id) = tcx.get_diagnostic_item(sym::TokenStream) else {
+
+        let Some(token_stream) = tcx
+            .get_diagnostic_item(sym::TokenStream)
+            .and_then(|did| tcx.type_of(did).no_bound_vars())
+        else {
             return;
         };
-        let Some(token_stream) = tcx.type_of(token_stream_def_id).no_bound_vars() else {
+        let Some(derive_expansion_options) = tcx
+            .get_diagnostic_item(sym::DeriveExpansionOptions)
+            .and_then(|did| tcx.type_of(did).no_bound_vars())
+        else {
             return;
         };
 
@@ -2332,8 +2339,24 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
             Unsafety::Normal,
             Abi::Rust,
         );
+        let expected_options_sig = tcx.mk_fn_sig(
+            [token_stream, derive_expansion_options],
+            token_stream,
+            false,
+            Unsafety::Normal,
+            Abi::Rust,
+        );
 
-        if let Err(terr) = ocx.eq(&cause, param_env, expected_sig, sig) {
+        let mut result = infcx.probe(|_| ocx.eq(&cause, param_env, expected_sig, sig));
+        if result.is_err()
+            && let ProcMacroKind::Derive = kind
+        {
+            if infcx.probe(|_| ocx.eq(&cause, param_env, expected_options_sig, sig)).is_ok() {
+                result = Ok(());
+            }
+        }
+
+        if let Err(terr) = result {
             let mut diag = tcx.dcx().create_err(errors::ProcMacroBadSig { span, kind });
 
             let hir_sig = tcx.hir().fn_sig_by_hir_id(hir_id);
@@ -2368,18 +2391,33 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                 }
             }
 
-            infcx.err_ctxt().note_type_err(
-                &mut diag,
-                &cause,
-                None,
-                Some(ValuePairs::PolySigs(ExpectedFound {
-                    expected: ty::Binder::dummy(expected_sig),
-                    found: ty::Binder::dummy(sig),
-                })),
-                terr,
-                false,
-                false,
-            );
+            let mut note_expected_found = |expected_sig| {
+                infcx.err_ctxt().note_type_err(
+                    &mut diag,
+                    &cause,
+                    None,
+                    Some(ValuePairs::PolySigs(ExpectedFound {
+                        expected: ty::Binder::dummy(expected_sig),
+                        found: ty::Binder::dummy(sig),
+                    })),
+                    terr,
+                    false,
+                    false,
+                )
+            };
+
+            note_expected_found(expected_sig);
+
+            if let ProcMacroKind::Derive = kind
+                && tcx
+                    .features()
+                    .declared_lib_features
+                    .iter()
+                    .any(|&(feature, _)| feature == sym::derive_const)
+            {
+                note_expected_found(expected_options_sig);
+            }
+
             diag.emit();
             self.abort.set(true);
         }

compiler/rustc_span/src/symbol.rs

@@ -182,6 +182,7 @@ symbols! {
         DecorateLint,
         Default,
         Deref,
+        DeriveExpansionOptions,
         DiagnosticMessage,
         DirBuilder,
         Display,

library/proc_macro/src/bridge/client.rs

@@ -484,12 +484,16 @@ impl ProcMacro {
         }
     }
 
-    pub const fn custom_derive(
+    pub const fn custom_derive<Discriminant>(
         trait_name: &'static str,
         attributes: &'static [&'static str],
-        expand: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy,
+        expand: impl ~const ExpandCustomDerive<Discriminant>,
     ) -> Self {
-        ProcMacro::CustomDerive { trait_name, attributes, client: Client::expand1(expand) }
+        ProcMacro::CustomDerive {
+            trait_name,
+            attributes,
+            client: Client::expand1(expand.into_fn()),
+        }
     }
 
     pub const fn attr(
@@ -506,3 +510,26 @@ impl ProcMacro {
         ProcMacro::Bang { name, client: Client::expand1(expand) }
     }
 }
+
+#[const_trait]
+pub trait ExpandCustomDerive<Discriminant> {
+    fn into_fn(self) -> impl Fn(crate::TokenStream) -> crate::TokenStream + Copy;
+}
+
+impl<F> const ExpandCustomDerive<()> for F
+where
+    F: Fn(crate::TokenStream) -> crate::TokenStream + Copy,
+{
+    fn into_fn(self) -> impl Fn(crate::TokenStream) -> crate::TokenStream + Copy {
+        self
+    }
+}
+
+impl<F> const ExpandCustomDerive<crate::DeriveExpansionOptions> for F
+where
+    F: Fn(crate::TokenStream, crate::DeriveExpansionOptions) -> crate::TokenStream + Copy,
+{
+    fn into_fn(self) -> impl Fn(crate::TokenStream) -> crate::TokenStream + Copy {
+        move |input| self(input, Default::default())
+    }
+}

library/proc_macro/src/bridge/mod.rs

@@ -60,6 +60,7 @@ macro_rules! with_api {
                 fn track_path(path: &str);
                 fn literal_from_str(s: &str) -> Result<Literal<$S::Span, $S::Symbol>, ()>;
                 fn emit_diagnostic(diagnostic: Diagnostic<$S::Span>);
+                fn is_derive_const() -> bool;
             },
             TokenStream {
                 fn drop($self: $S::TokenStream);

library/proc_macro/src/lib.rs

@@ -25,7 +25,9 @@
 #![feature(rustc_allow_const_fn_unstable)]
 #![feature(staged_api)]
 #![feature(allow_internal_unstable)]
+#![feature(const_trait_impl)]
 #![feature(decl_macro)]
+#![feature(effects)]
 #![feature(maybe_uninit_write_slice)]
 #![feature(negative_impls)]
 #![feature(new_uninit)]
@@ -85,6 +87,32 @@ impl !Send for TokenStream {}
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 impl !Sync for TokenStream {}
 
+/// Derive expansion options.
+#[rustc_diagnostic_item = "DeriveExpansionOptions"]
+#[unstable(feature = "derive_const", issue = "none")]
+#[derive(Default, Clone)]
+#[non_exhaustive]
+pub struct DeriveExpansionOptions;
+
+impl DeriveExpansionOptions {
+    /// Returns the default options.
+    #[unstable(feature = "derive_const", issue = "none")]
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Whether this is a `#[derive_const]` or a `#[derive]`.
+    #[unstable(feature = "derive_const", issue = "none")]
+    pub fn is_const(&self) -> bool {
+        bridge::client::FreeFunctions::is_derive_const()
+    }
+}
+
+#[unstable(feature = "derive_const", issue = "none")]
+impl !Send for DeriveExpansionOptions {}
+#[unstable(feature = "derive_const", issue = "none")]
+impl !Sync for DeriveExpansionOptions {}
+
 /// Error returned from `TokenStream::from_str`.
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 #[non_exhaustive]

src/tools/rust-analyzer/crates/proc-macro-srv/src/server/rust_analyzer_span.rs

@@ -82,6 +82,11 @@ impl server::FreeFunctions for RaSpanServer {
     fn emit_diagnostic(&mut self, _: bridge::Diagnostic<Self::Span>) {
         // FIXME handle diagnostic
     }
+
+    fn is_derive_const(&mut self) -> bool {
+        // FIXME: pass the correct information
+        false
+    }
 }
 
 impl server::TokenStream for RaSpanServer {

src/tools/rust-analyzer/crates/proc-macro-srv/src/server/token_id.rs

@@ -72,6 +72,11 @@ impl server::FreeFunctions for TokenIdServer {
     }
 
     fn emit_diagnostic(&mut self, _: bridge::Diagnostic<Self::Span>) {}
+
+    fn is_derive_const(&mut self) -> bool {
+        // FIXME: pass the correct information
+        false
+    }
 }
 
 impl server::TokenStream for TokenIdServer {

tests/ui/rfcs/rfc-2632-const-trait-impl/const_derives/auxiliary/is-derive-const.rs

@@ -0,0 +1,13 @@
+// force-host
+// no-prefer-dynamic
+#![crate_type = "proc-macro"]
+#![feature(derive_const)]
+
+extern crate proc_macro;
+
+use proc_macro::{TokenStream, DeriveExpansionOptions};
+
+#[proc_macro_derive(IsDeriveConst)]
+pub fn is_derive_const(_: TokenStream, options: DeriveExpansionOptions) -> TokenStream {
+    format!("const IS_DERIVE_CONST: bool = {};", options.is_const()).parse().unwrap()
+}

tests/ui/rfcs/rfc-2632-const-trait-impl/const_derives/derive-const-custom.rs

@@ -0,0 +1,20 @@
+// check-pass
+// edition: 2021
+// aux-crate:is_derive_const=is-derive-const.rs
+#![feature(derive_const)]
+
+const _: () = {
+    #[derive(is_derive_const::IsDeriveConst)]
+    struct _Type;
+
+    assert!(!IS_DERIVE_CONST);
+};
+
+const _: () = {
+    #[derive_const(is_derive_const::IsDeriveConst)]
+    struct _Type;
+
+    assert!(IS_DERIVE_CONST);
+};
+
+fn main() {}

tests/ui/rfcs/rfc-2632-const-trait-impl/const_derives/signature-proc-macro-derive.rs

@@ -0,0 +1,19 @@
+// Check that we suggest *both* possible signatures of derive proc macros, namely
+//     fn(TokenStream) -> TokenStream
+// and
+//     fn(TokenStream, DeriveExpansionOptions) -> TokenStream
+// provided libs feature `derive_const` is enabled.
+
+// force-host
+// no-prefer-dynamic
+
+#![crate_type = "proc-macro"]
+#![feature(derive_const)]
+
+extern crate proc_macro;
+
+#[proc_macro_derive(Blah)]
+pub fn bad_input() -> proc_macro::TokenStream {
+    //~^ ERROR derive proc macro has incorrect signature
+    Default::default()
+}

tests/ui/rfcs/rfc-2632-const-trait-impl/const_derives/signature-proc-macro-derive.stderr

@@ -0,0 +1,16 @@
+error: derive proc macro has incorrect signature
+  --> $DIR/signature-proc-macro-derive.rs:16:1
+   |
+LL | pub fn bad_input() -> proc_macro::TokenStream {
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   | |
+   | incorrect number of function parameters
+   | incorrect number of function parameters
+   |
+   = note: expected signature `fn(proc_macro::TokenStream) -> proc_macro::TokenStream`
+              found signature `fn() -> proc_macro::TokenStream`
+   = note: expected signature `fn(proc_macro::TokenStream, DeriveExpansionOptions) -> proc_macro::TokenStream`
+              found signature `fn() -> proc_macro::TokenStream`
+
+error: aborting due to 1 previous error
+