Introduce proc_macro::DeriveExpansionOptions

Author: fmease

Diff for: 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) => {

Diff for: 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<'_, '_> {

Diff for: compiler/rustc_passes/src/check_attr.rs

@@ -2280,6 +2280,13 @@ impl CheckAttrVisitor<'_> {
     ///
     /// If this best effort goes wrong, it will just emit a worse error later (see #102923)
     fn check_proc_macro(&self, hir_id: HirId, target: Target, kind: ProcMacroKind) {
+        // FIXME(fmease): This is temporary. Update this function to permit the signature of derive
+        // macros to be EITHER `fn(TS) -> TS` OR `fn(TS, DEO) -> TS`. Should be straightforward to
+        // implement.
+        if true {
+            return;
+        }
+
         if target != Target::Fn {
             return;
         }

Diff for: 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())
+    }
+}

Diff for: 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);

Diff for: 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,31 @@ impl !Send for TokenStream {}
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 impl !Sync for TokenStream {}
 
+/// Derive expansion options.
+#[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]

Diff for: 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()
+}

Diff for: 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() {}