wip: So yeah, tests pass, but still eval_always (not far from disk caching though I think)

futile · futile · commit 7ae17c924fed · 2024-08-19T14:06:27.000+02:00
diff --git a/compiler/rustc_ast/src/tokenstream.rs b/compiler/rustc_ast/src/tokenstream.rs
@@ -14,12 +14,14 @@
 //! ownership of the original.
 
 use std::borrow::Cow;
+use std::hash::Hash;
 use std::{cmp, fmt, iter};
 
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::{self, Lrc};
 use rustc_macros::{Decodable, Encodable, HashStable_Generic};
-use rustc_serialize::{Decodable, Encodable};
+use rustc_serialize::{Decodable, Encodable, Encoder};
+use rustc_span::def_id::{CrateNum, DefIndex};
 use rustc_span::{sym, Span, SpanDecoder, SpanEncoder, Symbol, DUMMY_SP};
 
 use crate::ast::{AttrStyle, StmtKind};
@@ -139,8 +141,10 @@ impl fmt::Debug for LazyAttrTokenStream {
 }
 
 impl<S: SpanEncoder> Encodable<S> for LazyAttrTokenStream {
-    fn encode(&self, _s: &mut S) {
-        panic!("Attempted to encode LazyAttrTokenStream");
+    fn encode(&self, s: &mut S) {
+        // TODO: welp
+        // TODO: (also) `.flattened()` here?
+        self.to_attr_token_stream().encode(s)
     }
 }
 
@@ -296,6 +300,96 @@ pub struct AttrsTarget {
 #[derive(Clone, Debug, Default, Encodable, Decodable)]
 pub struct TokenStream(pub(crate) Lrc<Vec<TokenTree>>);
 
+struct HashEncoder<H: std::hash::Hasher> {
+    hasher: H,
+}
+
+impl<H: std::hash::Hasher> Encoder for HashEncoder<H> {
+    fn emit_usize(&mut self, v: usize) {
+        self.hasher.write_usize(v)
+    }
+
+    fn emit_u128(&mut self, v: u128) {
+        self.hasher.write_u128(v)
+    }
+
+    fn emit_u64(&mut self, v: u64) {
+        self.hasher.write_u64(v)
+    }
+
+    fn emit_u32(&mut self, v: u32) {
+        self.hasher.write_u32(v)
+    }
+
+    fn emit_u16(&mut self, v: u16) {
+        self.hasher.write_u16(v)
+    }
+
+    fn emit_u8(&mut self, v: u8) {
+        self.hasher.write_u8(v)
+    }
+
+    fn emit_isize(&mut self, v: isize) {
+        self.hasher.write_isize(v)
+    }
+
+    fn emit_i128(&mut self, v: i128) {
+        self.hasher.write_i128(v)
+    }
+
+    fn emit_i64(&mut self, v: i64) {
+        self.hasher.write_i64(v)
+    }
+
+    fn emit_i32(&mut self, v: i32) {
+        self.hasher.write_i32(v)
+    }
+
+    fn emit_i16(&mut self, v: i16) {
+        self.hasher.write_i16(v)
+    }
+
+    fn emit_raw_bytes(&mut self, s: &[u8]) {
+        self.hasher.write(s)
+    }
+}
+
+impl<H: std::hash::Hasher> SpanEncoder for HashEncoder<H> {
+    fn encode_span(&mut self, span: Span) {
+        span.hash(&mut self.hasher)
+    }
+
+    fn encode_symbol(&mut self, symbol: Symbol) {
+        symbol.hash(&mut self.hasher)
+    }
+
+    fn encode_expn_id(&mut self, expn_id: rustc_span::ExpnId) {
+        expn_id.hash(&mut self.hasher)
+    }
+
+    fn encode_syntax_context(&mut self, syntax_context: rustc_span::SyntaxContext) {
+        syntax_context.hash(&mut self.hasher)
+    }
+
+    fn encode_crate_num(&mut self, crate_num: CrateNum) {
+        crate_num.hash(&mut self.hasher)
+    }
+
+    fn encode_def_index(&mut self, def_index: DefIndex) {
+        def_index.hash(&mut self.hasher)
+    }
+
+    fn encode_def_id(&mut self, def_id: rustc_span::def_id::DefId) {
+        def_id.hash(&mut self.hasher)
+    }
+}
+
+impl Hash for TokenStream {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        Encodable::encode(self, &mut HashEncoder { hasher: state });
+    }
+}
+
 /// Indicates whether a token can join with the following token to form a
 /// compound token. Used for conversions to `proc_macro::Spacing`. Also used to
 /// guide pretty-printing, which is where the `JointHidden` value (which isn't
diff --git a/compiler/rustc_expand/src/base.rs b/compiler/rustc_expand/src/base.rs
@@ -1072,6 +1072,9 @@ pub trait ResolverExpand {
         trait_def_id: DefId,
         impl_def_id: LocalDefId,
     ) -> Result<Vec<(Ident, Option<Ident>)>, Indeterminate>;
+
+    fn register_proc_macro_invoc(&mut self, invoc_id: LocalExpnId, ext: Lrc<SyntaxExtension>);
+    fn unregister_proc_macro_invoc(&mut self, invoc_id: LocalExpnId);
 }
 
 pub trait LintStoreExpand {
diff --git a/compiler/rustc_expand/src/derive_macro_expansion.rs b/compiler/rustc_expand/src/derive_macro_expansion.rs
@@ -0,0 +1,113 @@
+// TODO: remove
+#![allow(dead_code)]
+
+use std::cell::Cell;
+use std::ptr;
+
+use rustc_ast::tokenstream::TokenStream;
+use rustc_middle::ty::TyCtxt;
+use rustc_span::profiling::SpannedEventArgRecorder;
+use rustc_span::LocalExpnId;
+
+use crate::base::ExtCtxt;
+use crate::errors;
+
+pub(super) fn expand<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    key: (LocalExpnId, &'tcx TokenStream),
+) -> Result<&'tcx TokenStream, ()> {
+    let (invoc_id, input) = key;
+
+    let res = with_context(|(ecx, client)| {
+        let span = invoc_id.expn_data().call_site;
+        let _timer =
+            ecx.sess.prof.generic_activity_with_arg_recorder("expand_proc_macro", |recorder| {
+                recorder.record_arg_with_span(ecx.sess.source_map(), ecx.expansion_descr(), span);
+            });
+        let proc_macro_backtrace = ecx.ecfg.proc_macro_backtrace;
+        let strategy = crate::proc_macro::exec_strategy(ecx);
+        let server = crate::proc_macro_server::Rustc::new(ecx);
+        let res = match client.run(&strategy, server, input.clone(), proc_macro_backtrace) {
+            // TODO: without flattened some (weird) tests fail, but no idea if it's correct/enough
+            Ok(stream) => Ok(tcx.arena.alloc(stream.flattened()) as &TokenStream),
+            Err(e) => {
+                ecx.dcx().emit_err({
+                    errors::ProcMacroDerivePanicked {
+                        span,
+                        message: e.as_str().map(|message| errors::ProcMacroDerivePanickedHelp {
+                            message: message.into(),
+                        }),
+                    }
+                });
+                Err(())
+            }
+        };
+        res
+    });
+
+    res
+}
+
+type CLIENT = pm::bridge::client::Client<pm::TokenStream, pm::TokenStream>;
+
+// based on rust/compiler/rustc_middle/src/ty/context/tls.rs
+// #[cfg(not(parallel_compiler))]
+thread_local! {
+    /// A thread local variable that stores a pointer to the current `CONTEXT`.
+    static TLV: Cell<(*mut (), Option<CLIENT>)> = const { Cell::new((ptr::null_mut(), None)) };
+}
+
+#[inline]
+fn erase(context: &mut ExtCtxt<'_>) -> *mut () {
+    context as *mut _ as *mut ()
+}
+
+#[inline]
+unsafe fn downcast<'a>(context: *mut ()) -> &'a mut ExtCtxt<'a> {
+    unsafe { &mut *(context as *mut ExtCtxt<'a>) }
+}
+
+/// Sets `context` as the new current `CONTEXT` for the duration of the function `f`.
+#[inline]
+pub fn enter_context<'a, F, R>(context: (&mut ExtCtxt<'a>, CLIENT), f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    let (ectx, client) = context;
+    let erased = (erase(ectx), Some(client));
+    TLV.with(|tlv| {
+        let old = tlv.replace(erased);
+        let _reset = rustc_data_structures::defer(move || tlv.set(old));
+        f()
+    })
+}
+
+/// Allows access to the current `CONTEXT` in a closure if one is available.
+#[inline]
+#[track_caller]
+pub fn with_context_opt<F, R>(f: F) -> R
+where
+    F: for<'a, 'b> FnOnce(Option<&'b mut (&mut ExtCtxt<'a>, CLIENT)>) -> R,
+{
+    let (ectx, client_opt) = TLV.get();
+    if ectx.is_null() {
+        f(None)
+    } else {
+        // We could get an `CONTEXT` pointer from another thread.
+        // Ensure that `CONTEXT` is `DynSync`.
+        // TODO: we should not be able to?
+        // sync::assert_dyn_sync::<CONTEXT<'_>>();
+
+        unsafe { f(Some(&mut (downcast(ectx), client_opt.unwrap()))) }
+    }
+}
+
+/// Allows access to the current `CONTEXT`.
+/// Panics if there is no `CONTEXT` available.
+#[inline]
+pub fn with_context<F, R>(f: F) -> R
+where
+    F: for<'a, 'b> FnOnce(&'b mut (&mut ExtCtxt<'a>, CLIENT)) -> R,
+{
+    with_context_opt(|opt_context| f(opt_context.expect("no CONTEXT stored in tls")))
+}
diff --git a/compiler/rustc_expand/src/expand.rs b/compiler/rustc_expand/src/expand.rs
@@ -794,9 +794,14 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
                         span,
                         path,
                     };
+                    self.cx
+                        .resolver
+                        .register_proc_macro_invoc(invoc.expansion_data.id, ext.clone());
+                    invoc.expansion_data.id.expn_data();
                     let items = match expander.expand(self.cx, span, &meta, item, is_const) {
                         ExpandResult::Ready(items) => items,
                         ExpandResult::Retry(item) => {
+                            self.cx.resolver.unregister_proc_macro_invoc(invoc.expansion_data.id);
                             // Reassemble the original invocation for retrying.
                             return ExpandResult::Retry(Invocation {
                                 kind: InvocationKind::Derive { path: meta.path, item, is_const },
diff --git a/compiler/rustc_expand/src/lib.rs b/compiler/rustc_expand/src/lib.rs
@@ -28,10 +28,15 @@ mod proc_macro_server;
 pub use mbe::macro_rules::compile_declarative_macro;
 pub mod base;
 pub mod config;
+pub(crate) mod derive_macro_expansion;
 pub mod expand;
 pub mod module;
 // FIXME(Nilstrieb) Translate proc_macro diagnostics
 #[allow(rustc::untranslatable_diagnostic)]
 pub mod proc_macro;
 
+pub fn provide(providers: &mut rustc_middle::util::Providers) {
+    providers.derive_macro_expansion = derive_macro_expansion::expand;
+}
+
 rustc_fluent_macro::fluent_messages! { "../messages.ftl" }
diff --git a/compiler/rustc_expand/src/proc_macro.rs b/compiler/rustc_expand/src/proc_macro.rs
@@ -2,6 +2,7 @@ use rustc_ast as ast;
 use rustc_ast::ptr::P;
 use rustc_ast::tokenstream::TokenStream;
 use rustc_errors::ErrorGuaranteed;
+use rustc_middle::ty;
 use rustc_parse::parser::{ForceCollect, Parser};
 use rustc_session::config::ProcMacroExecutionStrategy;
 use rustc_span::profiling::SpannedEventArgRecorder;
@@ -31,7 +32,7 @@ impl<T> pm::bridge::server::MessagePipe<T> for MessagePipe<T> {
     }
 }
 
-fn exec_strategy(ecx: &ExtCtxt<'_>) -> impl pm::bridge::server::ExecutionStrategy {
+pub fn exec_strategy(ecx: &ExtCtxt<'_>) -> impl pm::bridge::server::ExecutionStrategy {
     pm::bridge::server::MaybeCrossThread::<MessagePipe<_>>::new(
         ecx.sess.opts.unstable_opts.proc_macro_execution_strategy
             == ProcMacroExecutionStrategy::CrossThread,
@@ -124,36 +125,27 @@ impl MultiItemModifier for DeriveProcMacro {
         // altogether. See #73345.
         crate::base::ann_pretty_printing_compatibility_hack(&item, &ecx.sess);
         let input = item.to_tokens();
-        let stream = {
-            let _timer =
-                ecx.sess.prof.generic_activity_with_arg_recorder("expand_proc_macro", |recorder| {
-                    recorder.record_arg_with_span(
-                        ecx.sess.source_map(),
-                        ecx.expansion_descr(),
-                        span,
-                    );
-                });
-            let proc_macro_backtrace = ecx.ecfg.proc_macro_backtrace;
-            let strategy = exec_strategy(ecx);
-            let server = proc_macro_server::Rustc::new(ecx);
-            match self.client.run(&strategy, server, input, proc_macro_backtrace) {
-                Ok(stream) => stream,
-                Err(e) => {
-                    ecx.dcx().emit_err({
-                        errors::ProcMacroDerivePanicked {
-                            span,
-                            message: e.as_str().map(|message| {
-                                errors::ProcMacroDerivePanickedHelp { message: message.into() }
-                            }),
-                        }
-                    });
-                    return ExpandResult::Ready(vec![]);
-                }
-            }
+        let res = ty::tls::with(|tcx| {
+            // TODO: without flattened some (weird) tests fail, but no idea if it's correct/enough
+            let input = tcx.arena.alloc(input.flattened()) as &TokenStream;
+            let invoc_id = ecx.current_expansion.id;
+
+            assert_eq!(invoc_id.expn_data().call_site, span);
+
+            let res = crate::derive_macro_expansion::enter_context((ecx, self.client), move || {
+                let res = tcx.derive_macro_expansion((invoc_id, input)).cloned();
+                res
+            });
+
+            res
+        });
+        let Ok(output) = res else {
+            // error will already have been emitted
+            return ExpandResult::Ready(vec![]);
         };
 
         let error_count_before = ecx.dcx().err_count();
-        let mut parser = Parser::new(&ecx.sess.psess, stream, Some("proc-macro derive"));
+        let mut parser = Parser::new(&ecx.sess.psess, output, Some("proc-macro derive"));
         let mut items = vec![];
 
         loop {
diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
@@ -624,6 +624,7 @@ pub static DEFAULT_QUERY_PROVIDERS: LazyLock<Providers> = LazyLock::new(|| {
     providers.resolutions = |tcx, ()| tcx.resolver_for_lowering_raw(()).1;
     providers.early_lint_checks = early_lint_checks;
     proc_macro_decls::provide(providers);
+    rustc_expand::provide(providers);
     rustc_const_eval::provide(providers);
     rustc_middle::hir::provide(providers);
     rustc_borrowck::provide(providers);
diff --git a/compiler/rustc_middle/src/arena.rs b/compiler/rustc_middle/src/arena.rs
@@ -115,6 +115,7 @@ macro_rules! arena_types {
             [decode] specialization_graph: rustc_middle::traits::specialization_graph::Graph,
             [] crate_inherent_impls: rustc_middle::ty::CrateInherentImpls,
             [] hir_owner_nodes: rustc_hir::OwnerNodes<'tcx>,
+            [] token_stream: rustc_ast::tokenstream::TokenStream,
         ]);
     )
 }
diff --git a/compiler/rustc_middle/src/query/erase.rs b/compiler/rustc_middle/src/query/erase.rs
@@ -1,6 +1,8 @@
 use std::intrinsics::transmute_unchecked;
 use std::mem::MaybeUninit;
 
+use rustc_ast::tokenstream::TokenStream;
+
 use crate::query::CyclePlaceholder;
 use crate::ty::adjustment::CoerceUnsizedInfo;
 use crate::ty::{self, Ty};
@@ -172,6 +174,10 @@ impl EraseType for Result<ty::EarlyBinder<'_, Ty<'_>>, CyclePlaceholder> {
     type Result = [u8; size_of::<Result<ty::EarlyBinder<'static, Ty<'_>>, CyclePlaceholder>>()];
 }
 
+impl EraseType for Result<&'_ TokenStream, ()> {
+    type Result = [u8; size_of::<Result<&'static TokenStream, ()>>()];
+}
+
 impl<T> EraseType for Option<&'_ T> {
     type Result = [u8; size_of::<Option<&'static ()>>()];
 }
diff --git a/compiler/rustc_middle/src/query/keys.rs b/compiler/rustc_middle/src/query/keys.rs
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
diff --git a/compiler/rustc_resolve/src/macros.rs b/compiler/rustc_resolve/src/macros.rs
diff --git a/compiler/rustc_span/src/hygiene.rs b/compiler/rustc_span/src/hygiene.rs

Original file line number	Diff line number	Diff line change
`@@ -1072,6 +1072,9 @@ pub trait ResolverExpand {`
`1072`	`1072`	`trait_def_id: DefId,`
`1073`	`1073`	`impl_def_id: LocalDefId,`
`1074`	`1074`	`) -> Result<Vec<(Ident, Option<Ident>)>, Indeterminate>;`
	`1075`	`+`
	`1076`	`+ fn register_proc_macro_invoc(&mut self, invoc_id: LocalExpnId, ext: Lrc<SyntaxExtension>);`
	`1077`	`+ fn unregister_proc_macro_invoc(&mut self, invoc_id: LocalExpnId);`
`1075`	`1078`	`}`
`1076`	`1079`
`1077`	`1080`	`pub trait LintStoreExpand {`
Original file line number	Diff line number	Diff line change
`@@ -115,6 +115,7 @@ macro_rules! arena_types {`
`115`	`115`	`[decode] specialization_graph: rustc_middle::traits::specialization_graph::Graph,`
`116`	`116`	`[] crate_inherent_impls: rustc_middle::ty::CrateInherentImpls,`
`117`	`117`	`[] hir_owner_nodes: rustc_hir::OwnerNodes<'tcx>,`
	`118`	`+ [] token_stream: rustc_ast::tokenstream::TokenStream,`
`118`	`119`	`]);`
`119`	`120`	`)`
`120`	`121`	`}`