rustc_expand_queries/
derive.rs

1use rustc_ast::tokenstream::TokenStream;
2use rustc_expand::base::ExtCtxt;
3use rustc_middle::ty::{TyCtxt, tls};
4use rustc_proc_macro as pm;
5use rustc_span::LocalExpnId;
6
7type DeriveClient = pm::bridge::client::Client;
8
9/// Stores the context necessary to expand a derive proc macro via a query.
10struct QueryDeriveExpandCtx {
11    /// Type-erased version of `&mut ExtCtxt`
12    expansion_ctx: *mut (),
13    client: DeriveClient,
14}
15
16impl QueryDeriveExpandCtx {
17    /// Store the extension context and the client into the thread local value.
18    /// It will be accessible via the `with` method while `f` is active.
19    fn enter<F, R>(ecx: &mut ExtCtxt<'_>, client: DeriveClient, f: F) -> R
20    where
21        F: FnOnce() -> R,
22    {
23        // We need erasure to get rid of the lifetime
24        let ctx = Self { expansion_ctx: ecx as *mut _ as *mut (), client };
25        DERIVE_EXPAND_CTX.set(&ctx, f)
26    }
27
28    /// Accesses the thread local value of the derive expansion context.
29    /// Must be called while the `enter` function is active.
30    fn with<F, R>(f: F) -> R
31    where
32        F: for<'a, 'b> FnOnce(&'b mut ExtCtxt<'a>, DeriveClient) -> R,
33    {
34        DERIVE_EXPAND_CTX.with(|ctx| {
35            let ectx = {
36                let casted = ctx.expansion_ctx.cast::<ExtCtxt<'_>>();
37                // SAFETY: We can only get the value from `with` while the `enter` function
38                // is active (on the callstack), and that function's signature ensures that the
39                // lifetime is valid.
40                // If `with` is called at some other time, it will panic due to usage of
41                // `scoped_tls::with`.
42                unsafe { casted.as_mut().unwrap() }
43            };
44
45            f(ectx, ctx.client)
46        })
47    }
48}
49
50// When we invoke a query to expand a derive proc macro, we need to provide it with the expansion
51// context and derive Client. We do that using a thread-local.
52static DERIVE_EXPAND_CTX: ::scoped_tls::ScopedKey<QueryDeriveExpandCtx> =
    ::scoped_tls::ScopedKey {
        inner: {
            const FOO: ::std::thread::LocalKey<::std::cell::Cell<*const ()>> =
                {
                    const __RUST_STD_INTERNAL_INIT: ::std::cell::Cell<*const ()>
                        =
                        { ::std::cell::Cell::new(::std::ptr::null()) };
                    unsafe {
                        ::std::thread::LocalKey::new(const {
                                    if ::std::mem::needs_drop::<::std::cell::Cell<*const ()>>()
                                        {
                                        |_|
                                            {
                                                #[thread_local]
                                                static __RUST_STD_INTERNAL_VAL:
                                                    ::std::thread::local_impl::EagerStorage<::std::cell::Cell<*const ()>>
                                                    =
                                                    ::std::thread::local_impl::EagerStorage::new(__RUST_STD_INTERNAL_INIT);
                                                __RUST_STD_INTERNAL_VAL.get()
                                            }
                                    } else {
                                        |_|
                                            {
                                                #[thread_local]
                                                static __RUST_STD_INTERNAL_VAL: ::std::cell::Cell<*const ()>
                                                    =
                                                    __RUST_STD_INTERNAL_INIT;
                                                &__RUST_STD_INTERNAL_VAL
                                            }
                                    }
                                })
                    }
                };
            &FOO
        },
        _marker: ::std::marker::PhantomData,
    };scoped_tls::scoped_thread_local!(static DERIVE_EXPAND_CTX: QueryDeriveExpandCtx);
53
54pub(crate) fn expand_derive_macro_cached(
55    invoc_id: LocalExpnId,
56    input: TokenStream,
57    ecx: &mut ExtCtxt<'_>,
58    client: DeriveClient,
59) -> Result<TokenStream, ()> {
60    tls::with(|tcx| {
61        let input = &*tcx.arena.alloc(input);
62        let key: (LocalExpnId, &TokenStream) = (invoc_id, input);
63
64        QueryDeriveExpandCtx::enter(ecx, client, move || tcx.derive_macro_expansion(key).cloned())
65    })
66}
67
68/// Provide a query for computing the output of a derive macro.
69pub(crate) fn derive_macro_expansion<'tcx>(
70    tcx: TyCtxt<'tcx>,
71    key: (LocalExpnId, &'tcx TokenStream),
72) -> Result<&'tcx TokenStream, ()> {
73    let (invoc_id, input) = key;
74
75    // Make sure that we invalidate the query when the crate defining the proc macro changes
76    let _ = tcx.crate_hash(invoc_id.expn_data().macro_def_id.unwrap().krate);
77
78    QueryDeriveExpandCtx::with(|ecx, client| {
79        rustc_expand::proc_macro::expand_derive_macro(invoc_id, input.clone(), ecx, client)
80            .map(|ts| &*tcx.arena.alloc(ts))
81    })
82}
rustc_expand_queries/derive.rs

rustc_expand_queries/
derive.rs
