Introduce proc_macro::DeriveExpansionOptions

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.sess.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.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

@@ -370,17 +370,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,
         }
     }
@@ -390,6 +392,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;
@@ -477,6 +484,10 @@ impl server::FreeFunctions for Rustc<'_, '_> {
         }
         self.sess().span_diagnostic.emit_diagnostic(&mut 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

@@ -2285,10 +2285,17 @@ impl CheckAttrVisitor<'_> {
         }
 
         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;
         };
 
@@ -2324,8 +2331,22 @@ impl CheckAttrVisitor<'_> {
             Unsafety::Normal,
             Abi::Rust,
         );
+        let expected_options_sig = tcx.mk_fn_sig(
+            [token_stream, derive_expansion_options],
+            token_stream,
+            false,
+            Unsafety::Normal,
+            Abi::Rust,
+        );
+
+        let mut result = infcx.probe(|_| ocx.eq(&cause, param_env, expected_sig, sig));
+        if result.is_err()
+            && let ProcMacroKind::Derive = kind
+        {
+            result = infcx.probe(|_| ocx.eq(&cause, param_env, expected_options_sig, sig));
+        }
 
-        if let Err(terr) = ocx.eq(&cause, param_env, expected_sig, sig) {
+        if let Err(terr) = result {
             let mut diag = tcx.sess.create_err(errors::ProcMacroBadSig { span, kind });
 
             let hir_sig = tcx.hir().fn_sig_by_hir_id(hir_id);
@@ -2360,18 +2381,33 @@ impl CheckAttrVisitor<'_> {
                 }
             }
 
-            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

@@ -183,6 +183,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

@@ -59,6 +59,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]

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
+