Skip to content

Commit 4f49fff

Browse files
ZoxcZoxc
committed
Clean up parsing code and split out codegen for the QueryDescription impl
1 parent 9e9d03f commit 4f49fff

File tree

3 files changed

+142
-185
lines changed

3 files changed

+142
-185
lines changed

‎src/librustc_macros/src/lib.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ use proc_macro::TokenStream;
99

1010
mod hash_stable;
1111
mod query;
12-
mod tt;
1312

1413
#[proc_macro]
1514
pub fn rustc_queries(input: TokenStream) -> TokenStream {

‎src/librustc_macros/src/query.rs

Lines changed: 142 additions & 119 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,21 @@
11
use proc_macro::TokenStream;
22
use proc_macro2::Span;
33
use syn::{
4-
Token, Ident, Type, Attribute, ReturnType, Expr, Block,
4+
Token, Ident, Type, Attribute, ReturnType, Expr, Block, Error,
55
braced, parenthesized, parse_macro_input,
66
};
7+
use syn::spanned::Spanned;
78
use syn::parse::{Result, Parse, ParseStream};
89
use syn::punctuated::Punctuated;
10+
use syn;
911
use quote::quote;
10-
use crate::tt::TS;
1112

13+
#[allow(non_camel_case_types)]
14+
mod kw {
15+
syn::custom_keyword!(query);
16+
}
17+
18+
/// Ident or a wildcard `_`.
1219
struct IdentOrWild(Ident);
1320

1421
impl Parse for IdentOrWild {
@@ -22,17 +29,27 @@ impl Parse for IdentOrWild {
2229
}
2330
}
2431

25-
enum QueryAttribute {
32+
/// A modifier for a query
33+
enum QueryModifier {
34+
/// The description of the query
2635
Desc(Option<Ident>, Punctuated<Expr, Token![,]>),
36+
37+
/// Cache the query to disk if the `Expr` returns true.
2738
Cache(Option<Ident>, Expr),
39+
40+
/// Custom code to load the query from disk.
2841
LoadCached(Ident, Ident, Block),
42+
43+
/// A cycle error for this query aborting the compilation with a fatal error.
2944
FatalCycle,
3045
}
3146

32-
impl Parse for QueryAttribute {
47+
impl Parse for QueryModifier {
3348
fn parse(input: ParseStream<'_>) -> Result<Self> {
34-
let attr: Ident = input.parse()?;
35-
if attr == "desc" {
49+
let modifier: Ident = input.parse()?;
50+
if modifier == "desc" {
51+
// Parse a description modifier like:
52+
// `desc { |tcx| "foo {}", tcx.item_path(key) }`
3653
let attr_content;
3754
braced!(attr_content in input);
3855
let tcx = if attr_content.peek(Token![|]) {
@@ -44,11 +61,10 @@ impl Parse for QueryAttribute {
4461
None
4562
};
4663
let desc = attr_content.parse_terminated(Expr::parse)?;
47-
if !attr_content.is_empty() {
48-
panic!("unexpected tokens in block");
49-
};
50-
Ok(QueryAttribute::Desc(tcx, desc))
51-
} else if attr == "cache" {
64+
Ok(QueryModifier::Desc(tcx, desc))
65+
} else if modifier == "cache" {
66+
// Parse a cache modifier like:
67+
// `cache { |tcx| key.is_local() }`
5268
let attr_content;
5369
braced!(attr_content in input);
5470
let tcx = if attr_content.peek(Token![|]) {
@@ -60,68 +76,59 @@ impl Parse for QueryAttribute {
6076
None
6177
};
6278
let expr = attr_content.parse()?;
63-
if !attr_content.is_empty() {
64-
panic!("unexpected tokens in block");
65-
};
66-
Ok(QueryAttribute::Cache(tcx, expr))
67-
} else if attr == "load_cached" {
79+
Ok(QueryModifier::Cache(tcx, expr))
80+
} else if modifier == "load_cached" {
81+
// Parse a load_cached modifier like:
82+
// `load_cached(tcx, id) { tcx.queries.on_disk_cache.try_load_query_result(tcx, id) }`
6883
let args;
6984
parenthesized!(args in input);
7085
let tcx = args.parse()?;
7186
args.parse::<Token![,]>()?;
7287
let id = args.parse()?;
73-
if !args.is_empty() {
74-
panic!("unexpected tokens in arguments");
75-
};
7688
let block = input.parse()?;
77-
Ok(QueryAttribute::LoadCached(tcx, id, block))
78-
} else if attr == "fatal_cycle" {
79-
Ok(QueryAttribute::FatalCycle)
89+
Ok(QueryModifier::LoadCached(tcx, id, block))
90+
} else if modifier == "fatal_cycle" {
91+
Ok(QueryModifier::FatalCycle)
8092
} else {
81-
panic!("unknown query modifier {}", attr)
93+
Err(Error::new(modifier.span(), "unknown query modifier"))
8294
}
8395
}
8496
}
8597

98+
/// Ensures only doc comment attributes are used
99+
fn check_attributes(attrs: Vec<Attribute>) -> Result<()> {
100+
for attr in attrs {
101+
if !attr.path.is_ident("doc") {
102+
return Err(Error::new(attr.span(), "attributes not supported on queries"));
103+
}
104+
}
105+
Ok(())
106+
}
107+
108+
/// A compiler query. `query ... { ... }`
86109
struct Query {
87-
attrs: List<QueryAttribute>,
110+
attrs: List<QueryModifier>,
88111
name: Ident,
89112
key: IdentOrWild,
90113
arg: Type,
91114
result: ReturnType,
92115
}
93116

94-
fn check_attributes(attrs: Vec<Attribute>) {
95-
for attr in attrs {
96-
let path = attr.path;
97-
let path = quote! { #path };
98-
let path = TS(&path);
99-
100-
if path != TS(&quote! { doc }) {
101-
panic!("attribute `{}` not supported on queries", path.0)
102-
}
103-
}
104-
}
105-
106117
impl Parse for Query {
107118
fn parse(input: ParseStream<'_>) -> Result<Self> {
108-
check_attributes(input.call(Attribute::parse_outer)?);
119+
check_attributes(input.call(Attribute::parse_outer)?)?;
109120

110-
let query: Ident = input.parse()?;
111-
if query != "query" {
112-
panic!("expected `query`");
113-
}
121+
// Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
122+
input.parse::<kw::query>()?;
114123
let name: Ident = input.parse()?;
115124
let arg_content;
116125
parenthesized!(arg_content in input);
117126
let key = arg_content.parse()?;
118127
arg_content.parse::<Token![:]>()?;
119128
let arg = arg_content.parse()?;
120-
if !arg_content.is_empty() {
121-
panic!("expected only one query argument");
122-
};
123129
let result = input.parse()?;
124130

131+
// Parse the query modifiers
125132
let content;
126133
braced!(content in input);
127134
let attrs = content.parse()?;
@@ -136,6 +143,7 @@ impl Parse for Query {
136143
}
137144
}
138145

146+
/// A type used to greedily parse another type until the input is empty.
139147
struct List<T>(Vec<T>);
140148

141149
impl<T: Parse> Parse for List<T> {
@@ -148,6 +156,7 @@ impl<T: Parse> Parse for List<T> {
148156
}
149157
}
150158

159+
/// A named group containing queries.
151160
struct Group {
152161
name: Ident,
153162
queries: List<Query>,
@@ -165,6 +174,88 @@ impl Parse for Group {
165174
}
166175
}
167176

177+
/// Add the impl of QueryDescription for the query to `impls` if one is requested
178+
fn add_query_description_impl(query: &Query, impls: &mut proc_macro2::TokenStream) {
179+
let name = &query.name;
180+
let arg = &query.arg;
181+
let key = &query.key.0;
182+
183+
// Find custom code to load the query from disk
184+
let load_cached = query.attrs.0.iter().find_map(|attr| match attr {
185+
QueryModifier::LoadCached(tcx, id, block) => Some((tcx, id, block)),
186+
_ => None,
187+
});
188+
189+
// Find out if we should cache the query on disk
190+
let cache = query.attrs.0.iter().find_map(|attr| match attr {
191+
QueryModifier::Cache(tcx, expr) => Some((tcx, expr)),
192+
_ => None,
193+
}).map(|(tcx, expr)| {
194+
let try_load_from_disk = if let Some((tcx, id, block)) = load_cached {
195+
quote! {
196+
#[inline]
197+
fn try_load_from_disk(
198+
#tcx: TyCtxt<'_, 'tcx, 'tcx>,
199+
#id: SerializedDepNodeIndex
200+
) -> Option<Self::Value> {
201+
#block
202+
}
203+
}
204+
} else {
205+
quote! {
206+
#[inline]
207+
fn try_load_from_disk(
208+
tcx: TyCtxt<'_, 'tcx, 'tcx>,
209+
id: SerializedDepNodeIndex
210+
) -> Option<Self::Value> {
211+
tcx.queries.on_disk_cache.try_load_query_result(tcx, id)
212+
}
213+
}
214+
};
215+
216+
let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
217+
quote! {
218+
#[inline]
219+
fn cache_on_disk(#tcx: TyCtxt<'_, 'tcx, 'tcx>, #key: Self::Key) -> bool {
220+
#expr
221+
}
222+
223+
#try_load_from_disk
224+
}
225+
});
226+
227+
if cache.is_none() && load_cached.is_some() {
228+
panic!("load_cached modifier on query `{}` without a cache modifier", name);
229+
}
230+
231+
let desc = query.attrs.0.iter().find_map(|attr| match attr {
232+
QueryModifier::Desc(tcx, desc) => Some((tcx, desc)),
233+
_ => None,
234+
}).map(|(tcx, desc)| {
235+
let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
236+
quote! {
237+
fn describe(
238+
#tcx: TyCtxt<'_, '_, '_>,
239+
#key: #arg,
240+
) -> Cow<'static, str> {
241+
format!(#desc).into()
242+
}
243+
}
244+
});
245+
246+
if desc.is_some() || cache.is_some() {
247+
let cache = cache.unwrap_or(quote! {});
248+
let desc = desc.unwrap_or(quote! {});
249+
250+
impls.extend(quote! {
251+
impl<'tcx> QueryDescription<'tcx> for queries::#name<'tcx> {
252+
#desc
253+
#cache
254+
}
255+
});
256+
}
257+
}
258+
168259
pub fn rustc_queries(input: TokenStream) -> TokenStream {
169260
let groups = parse_macro_input!(input as List<Group>);
170261

@@ -178,99 +269,31 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
178269
for query in &group.queries.0 {
179270
let name = &query.name;
180271
let arg = &query.arg;
181-
let key = &query.key.0;
182272
let result_full = &query.result;
183273
let result = match query.result {
184274
ReturnType::Default => quote! { -> () },
185275
_ => quote! { #result_full },
186276
};
187277

188-
let load_cached = query.attrs.0.iter().find_map(|attr| match attr {
189-
QueryAttribute::LoadCached(tcx, id, block) => Some((tcx, id, block)),
190-
_ => None,
191-
});
192-
193-
// Find out if we should cache the query on disk
194-
let cache = query.attrs.0.iter().find_map(|attr| match attr {
195-
QueryAttribute::Cache(tcx, expr) => Some((tcx, expr)),
196-
_ => None,
197-
}).map(|(tcx, expr)| {
198-
let try_load_from_disk = if let Some((tcx, id, block)) = load_cached {
199-
quote! {
200-
#[inline]
201-
fn try_load_from_disk(
202-
#tcx: TyCtxt<'_, 'tcx, 'tcx>,
203-
#id: SerializedDepNodeIndex
204-
) -> Option<Self::Value> {
205-
#block
206-
}
207-
}
208-
} else {
209-
quote! {
210-
#[inline]
211-
fn try_load_from_disk(
212-
tcx: TyCtxt<'_, 'tcx, 'tcx>,
213-
id: SerializedDepNodeIndex
214-
) -> Option<Self::Value> {
215-
tcx.queries.on_disk_cache.try_load_query_result(tcx, id)
216-
}
217-
}
218-
};
219-
220-
let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
221-
quote! {
222-
#[inline]
223-
fn cache_on_disk(#tcx: TyCtxt<'_, 'tcx, 'tcx>, #key: Self::Key) -> bool {
224-
#expr
225-
}
226-
227-
#try_load_from_disk
228-
}
229-
});
230-
231-
if cache.is_none() && load_cached.is_some() {
232-
panic!("load_cached modifier on query `{}` without a cache modifier", name);
233-
}
234-
278+
// Look for a fatal_cycle modifier to pass on
235279
let fatal_cycle = query.attrs.0.iter().find_map(|attr| match attr {
236-
QueryAttribute::FatalCycle => Some(()),
280+
QueryModifier::FatalCycle => Some(()),
237281
_ => None,
238282
}).map(|_| quote! { fatal_cycle }).unwrap_or(quote! {});
239283

284+
// Add the query to the group
240285
group_stream.extend(quote! {
241286
[#fatal_cycle] fn #name: #name(#arg) #result,
242287
});
243288

244-
let desc = query.attrs.0.iter().find_map(|attr| match attr {
245-
QueryAttribute::Desc(tcx, desc) => Some((tcx, desc)),
246-
_ => None,
247-
}).map(|(tcx, desc)| {
248-
let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
249-
quote! {
250-
fn describe(
251-
#tcx: TyCtxt<'_, '_, '_>,
252-
#key: #arg,
253-
) -> Cow<'static, str> {
254-
format!(#desc).into()
255-
}
256-
}
257-
});
258-
259-
if desc.is_some() || cache.is_some() {
260-
let cache = cache.unwrap_or(quote! {});
261-
let desc = desc.unwrap_or(quote! {});
262-
263-
query_description_stream.extend(quote! {
264-
impl<'tcx> QueryDescription<'tcx> for queries::#name<'tcx> {
265-
#desc
266-
#cache
267-
}
268-
});
269-
}
289+
add_query_description_impl(query, &mut query_description_stream);
270290

291+
// Create a dep node for the query
271292
dep_node_def_stream.extend(quote! {
272293
[] #name(#arg),
273294
});
295+
296+
// Add a match arm to force the query given the dep node
274297
dep_node_force_stream.extend(quote! {
275298
DepKind::#name => {
276299
if let Some(key) = RecoverKey::recover($tcx, $dep_node) {

0 commit comments

Comments
 (0)