1use std::{mem, slice};
2
3use rustc_ast::ptr::P;
4use rustc_ast::visit::{self, Visitor};
5use rustc_ast::{self as ast, HasNodeId, NodeId, attr};
6use rustc_ast_pretty::pprust;
7use rustc_attr_parsing::AttributeParser;
8use rustc_errors::DiagCtxtHandle;
9use rustc_expand::base::{ExtCtxt, ResolverExpand};
10use rustc_expand::expand::{AstFragment, ExpansionConfig};
11use rustc_feature::Features;
12use rustc_hir::attrs::AttributeKind;
13use rustc_session::Session;
14use rustc_span::hygiene::AstPass;
15use rustc_span::source_map::SourceMap;
16use rustc_span::{DUMMY_SP, Ident, Span, Symbol, kw, sym};
17use smallvec::smallvec;
18use thin_vec::{ThinVec, thin_vec};
19
20use crate::errors;
21
22struct ProcMacroDerive {
23 id: NodeId,
24 trait_name: Symbol,
25 function_ident: Ident,
26 span: Span,
27 attrs: ThinVec<Symbol>,
28}
29
30struct ProcMacroDef {
31 id: NodeId,
32 function_ident: Ident,
33 span: Span,
34}
35
36enum ProcMacro {
37 Derive(ProcMacroDerive),
38 Attr(ProcMacroDef),
39 Bang(ProcMacroDef),
40}
41
42struct CollectProcMacros<'a> {
43 macros: Vec<ProcMacro>,
44 in_root: bool,
45 dcx: DiagCtxtHandle<'a>,
46 session: &'a Session,
47 source_map: &'a SourceMap,
48 is_proc_macro_crate: bool,
49 is_test_crate: bool,
50}
51
52pub fn inject(
53 krate: &mut ast::Crate,
54 sess: &Session,
55 features: &Features,
56 resolver: &mut dyn ResolverExpand,
57 is_proc_macro_crate: bool,
58 has_proc_macro_decls: bool,
59 is_test_crate: bool,
60 dcx: DiagCtxtHandle<'_>,
61) {
62 let ecfg = ExpansionConfig::default(sym::proc_macro, features);
63 let mut cx = ExtCtxt::new(sess, ecfg, resolver, None);
64
65 let mut collect = CollectProcMacros {
66 macros: Vec::new(),
67 in_root: true,
68 dcx,
69 session: sess,
70 source_map: sess.source_map(),
71 is_proc_macro_crate,
72 is_test_crate,
73 };
74
75 if has_proc_macro_decls || is_proc_macro_crate {
76 visit::walk_crate(&mut collect, krate);
77 }
78 let macros = collect.macros;
79
80 if !is_proc_macro_crate {
81 return;
82 }
83
84 if is_test_crate {
85 return;
86 }
87
88 let decls = mk_decls(&mut cx, ¯os);
89 krate.items.push(decls);
90}
91
92impl<'a> CollectProcMacros<'a> {
93 fn check_not_pub_in_root(&self, vis: &ast::Visibility, sp: Span) {
94 if self.is_proc_macro_crate && self.in_root && vis.kind.is_pub() {
95 self.dcx.emit_err(errors::ProcMacro { span: sp });
96 }
97 }
98
99 fn collect_custom_derive(
100 &mut self,
101 item: &'a ast::Item,
102 function_ident: Ident,
103 attr: &'a ast::Attribute,
104 ) {
105 let Some(rustc_hir::Attribute::Parsed(AttributeKind::ProcMacroDerive {
106 trait_name,
107 helper_attrs,
108 ..
109 })) = AttributeParser::parse_limited(
110 self.session,
111 slice::from_ref(attr),
112 sym::proc_macro_derive,
113 item.span,
114 item.node_id(),
115 None,
116 )
117 else {
118 return;
119 };
120
121 if self.in_root && item.vis.kind.is_pub() {
122 self.macros.push(ProcMacro::Derive(ProcMacroDerive {
123 id: item.id,
124 span: item.span,
125 trait_name,
126 function_ident,
127 attrs: helper_attrs,
128 }));
129 } else {
130 let msg = if !self.in_root {
131 "functions tagged with `#[proc_macro_derive]` must \
132 currently reside in the root of the crate"
133 } else {
134 "functions tagged with `#[proc_macro_derive]` must be `pub`"
135 };
136 self.dcx.span_err(self.source_map.guess_head_span(item.span), msg);
137 }
138 }
139
140 fn collect_attr_proc_macro(&mut self, item: &'a ast::Item, function_ident: Ident) {
141 if self.in_root && item.vis.kind.is_pub() {
142 self.macros.push(ProcMacro::Attr(ProcMacroDef {
143 id: item.id,
144 span: item.span,
145 function_ident,
146 }));
147 } else {
148 let msg = if !self.in_root {
149 "functions tagged with `#[proc_macro_attribute]` must \
150 currently reside in the root of the crate"
151 } else {
152 "functions tagged with `#[proc_macro_attribute]` must be `pub`"
153 };
154 self.dcx.span_err(self.source_map.guess_head_span(item.span), msg);
155 }
156 }
157
158 fn collect_bang_proc_macro(&mut self, item: &'a ast::Item, function_ident: Ident) {
159 if self.in_root && item.vis.kind.is_pub() {
160 self.macros.push(ProcMacro::Bang(ProcMacroDef {
161 id: item.id,
162 span: item.span,
163 function_ident,
164 }));
165 } else {
166 let msg = if !self.in_root {
167 "functions tagged with `#[proc_macro]` must \
168 currently reside in the root of the crate"
169 } else {
170 "functions tagged with `#[proc_macro]` must be `pub`"
171 };
172 self.dcx.span_err(self.source_map.guess_head_span(item.span), msg);
173 }
174 }
175}
176
177impl<'a> Visitor<'a> for CollectProcMacros<'a> {
178 fn visit_item(&mut self, item: &'a ast::Item) {
179 if let ast::ItemKind::MacroDef(..) = item.kind {
180 if self.is_proc_macro_crate && attr::contains_name(&item.attrs, sym::macro_export) {
181 self.dcx.emit_err(errors::ExportMacroRules {
182 span: self.source_map.guess_head_span(item.span),
183 });
184 }
185 }
186
187 let mut found_attr: Option<&'a ast::Attribute> = None;
188
189 for attr in &item.attrs {
190 if attr.is_proc_macro_attr() {
191 if let Some(prev_attr) = found_attr {
192 let prev_item = prev_attr.get_normal_item();
193 let item = attr.get_normal_item();
194 let path_str = pprust::path_to_string(&item.path);
195 let msg = if item.path.segments[0].ident.name
196 == prev_item.path.segments[0].ident.name
197 {
198 format!(
199 "only one `#[{path_str}]` attribute is allowed on any given function",
200 )
201 } else {
202 format!(
203 "`#[{}]` and `#[{}]` attributes cannot both be applied
204 to the same function",
205 path_str,
206 pprust::path_to_string(&prev_item.path),
207 )
208 };
209
210 self.dcx
211 .struct_span_err(attr.span, msg)
212 .with_span_label(prev_attr.span, "previous attribute here")
213 .emit();
214
215 return;
216 }
217
218 found_attr = Some(attr);
219 }
220 }
221
222 let Some(attr) = found_attr else {
223 self.check_not_pub_in_root(&item.vis, self.source_map.guess_head_span(item.span));
224 let prev_in_root = mem::replace(&mut self.in_root, false);
225 visit::walk_item(self, item);
226 self.in_root = prev_in_root;
227 return;
228 };
229
230 let fn_ident = if let ast::ItemKind::Fn(fn_) = &item.kind {
233 fn_.ident
234 } else {
235 self.dcx
236 .create_err(errors::AttributeOnlyBeUsedOnBareFunctions {
237 span: attr.span,
238 path: &pprust::path_to_string(&attr.get_normal_item().path),
239 })
240 .emit();
241 return;
242 };
243
244 if self.is_test_crate {
245 return;
246 }
247
248 if !self.is_proc_macro_crate {
249 self.dcx
250 .create_err(errors::AttributeOnlyUsableWithCrateType {
251 span: attr.span,
252 path: &pprust::path_to_string(&attr.get_normal_item().path),
253 })
254 .emit();
255 return;
256 }
257
258 if attr.has_name(sym::proc_macro_derive) {
260 self.collect_custom_derive(item, fn_ident, attr);
261 } else if attr.has_name(sym::proc_macro_attribute) {
262 self.collect_attr_proc_macro(item, fn_ident);
263 } else if attr.has_name(sym::proc_macro) {
264 self.collect_bang_proc_macro(item, fn_ident);
265 };
266
267 let prev_in_root = mem::replace(&mut self.in_root, false);
268 visit::walk_item(self, item);
269 self.in_root = prev_in_root;
270 }
271}
272
273fn mk_decls(cx: &mut ExtCtxt<'_>, macros: &[ProcMacro]) -> P<ast::Item> {
290 let expn_id = cx.resolver.expansion_for_ast_pass(
291 DUMMY_SP,
292 AstPass::ProcMacroHarness,
293 &[sym::rustc_attrs, sym::proc_macro_internals],
294 None,
295 );
296 let span = DUMMY_SP.with_def_site_ctxt(expn_id.to_expn_id());
297
298 let proc_macro = Ident::new(sym::proc_macro, span);
299 let krate = cx.item(span, ast::AttrVec::new(), ast::ItemKind::ExternCrate(None, proc_macro));
300
301 let bridge = Ident::new(sym::bridge, span);
302 let client = Ident::new(sym::client, span);
303 let proc_macro_ty = Ident::new(sym::ProcMacro, span);
304 let custom_derive = Ident::new(sym::custom_derive, span);
305 let attr = Ident::new(sym::attr, span);
306 let bang = Ident::new(sym::bang, span);
307
308 let decls = macros
313 .iter()
314 .map(|m| {
315 let harness_span = span;
316 let span = match m {
317 ProcMacro::Derive(m) => m.span,
318 ProcMacro::Attr(m) | ProcMacro::Bang(m) => m.span,
319 };
320 let local_path = |cx: &ExtCtxt<'_>, ident| cx.expr_path(cx.path(span, vec![ident]));
321 let proc_macro_ty_method_path = |cx: &ExtCtxt<'_>, method| {
322 cx.expr_path(cx.path(
323 span.with_ctxt(harness_span.ctxt()),
324 vec![proc_macro, bridge, client, proc_macro_ty, method],
325 ))
326 };
327 match m {
328 ProcMacro::Derive(cd) => {
329 cx.resolver.declare_proc_macro(cd.id);
330 cx.expr_call(
333 harness_span,
334 proc_macro_ty_method_path(cx, custom_derive),
335 thin_vec![
336 cx.expr_str(span, cd.trait_name),
337 cx.expr_array_ref(
338 span,
339 cd.attrs
340 .iter()
341 .map(|&s| cx.expr_str(span, s))
342 .collect::<ThinVec<_>>(),
343 ),
344 local_path(cx, cd.function_ident),
345 ],
346 )
347 }
348 ProcMacro::Attr(ca) | ProcMacro::Bang(ca) => {
349 cx.resolver.declare_proc_macro(ca.id);
350 let ident = match m {
351 ProcMacro::Attr(_) => attr,
352 ProcMacro::Bang(_) => bang,
353 ProcMacro::Derive(_) => unreachable!(),
354 };
355
356 cx.expr_call(
359 harness_span,
360 proc_macro_ty_method_path(cx, ident),
361 thin_vec![
362 cx.expr_str(span, ca.function_ident.name),
363 local_path(cx, ca.function_ident),
364 ],
365 )
366 }
367 }
368 })
369 .collect();
370
371 let mut decls_static = cx.item_static(
372 span,
373 Ident::new(sym::_DECLS, span),
374 cx.ty_ref(
375 span,
376 cx.ty(
377 span,
378 ast::TyKind::Slice(
379 cx.ty_path(cx.path(span, vec![proc_macro, bridge, client, proc_macro_ty])),
380 ),
381 ),
382 None,
383 ast::Mutability::Not,
384 ),
385 ast::Mutability::Not,
386 cx.expr_array_ref(span, decls),
387 );
388 decls_static.attrs.extend([
389 cx.attr_word(sym::rustc_proc_macro_decls, span),
390 cx.attr_word(sym::used, span),
391 cx.attr_nested_word(sym::allow, sym::deprecated, span),
392 ]);
393
394 let block = cx.expr_block(
395 cx.block(span, thin_vec![cx.stmt_item(span, krate), cx.stmt_item(span, decls_static)]),
396 );
397
398 let anon_constant = cx.item_const(
399 span,
400 Ident::new(kw::Underscore, span),
401 cx.ty(span, ast::TyKind::Tup(ThinVec::new())),
402 block,
403 );
404
405 let items = AstFragment::Items(smallvec![anon_constant]);
407 cx.monotonic_expander().fully_expand_fragment(items).make_items().pop().unwrap()
408}