1use crate::{
2 used_proc_macro::{
3 token_stream::IntoIter,
4 Delimiter, Ident, Group, Literal, Punct, Span, Spacing, TokenStream, TokenTree,
5 },
6 macro_utils_shared::{
7 cmp_ts::{self, ComparableTT, Found},
8 list_generation::{ListIter, parse_bounded, parse_unbounded},
9 expect_no_tokens,
10 out_braced_tt,
11 parse_count_param, parse_ident, parse_int_or_range_param,
12 parse_keyword, parse_check_punct,
13 parse_parentheses, parse_bounded_range_param,
14 macro_span, out_parenthesized_tt,
15 match_token,
16 },
17 parsing_shared::{out_parenthesized, parse_macro_invocation},
18 mmatches,
19 try_,
20};
21
22use core::{
23 iter::{Peekable, once},
24 ops::Range,
25 mem,
26};
27
28use alloc::{
29 collections::VecDeque,
30 string::{String, ToString},
31 vec::Vec,
32 format,
33};
34
35
36#[cfg(test)]
37mod mu_tests;
38
39
40pub fn rewrap_macro_parameters(tokens: TokenStream) -> TokenStream {
41 let mut prev_tilde;
42 let mut curr_tilde = false;
43 let mut out = TokenStream::new();
44
45 for tt in tokens {
46 prev_tilde = mem::replace(&mut curr_tilde, false);
47
48 let tt_out = match tt {
49 TokenTree::Group(group) => {
50 let out = rewrap_macro_parameters(group.stream());
51 let span = group.span();
52
53 let delim = if prev_tilde && group.delimiter() == Delimiter::None {
54 Delimiter::Parenthesis
55 } else {
56 group.delimiter()
57 };
58
59 let mut group = Group::new(delim, out);
60 group.set_span(span);
61 TokenTree::Group(group)
62 }
63 TokenTree::Punct(punct) => {
64 curr_tilde = punct.as_char() == '~';
65 if !prev_tilde && curr_tilde {
66 continue;
67 } else {
68 curr_tilde = false;
69 TokenTree::Punct(punct)
70 }
71 },
72 tt @ TokenTree::Ident(_) if prev_tilde => {
73 let span = tt.span();
74 let mut group = Group::new(Delimiter::Parenthesis, TokenStream::from(tt));
75 group.set_span(span);
76 TokenTree::Group(group)
77 }
78 tt => tt,
79 };
80
81 out.extend(once(tt_out));
82 }
83 out
84}
85
86enum ExpandedInto{
87 Macro,
88 Expr,
89}
90
91pub(crate) fn count_tts(tokens: TokenStream) -> crate::Result<TokenStream> {
92 let mut iter = tokens.into_iter().peekable();
93
94 fn output_counted(counted: Group, ei: ExpandedInto, out: &mut TokenStream) {
95 let count = counted.stream().into_iter().count();
96 let mut lit = match ei {
97 ExpandedInto::Macro => Literal::usize_unsuffixed(count),
98 ExpandedInto::Expr => Literal::usize_suffixed(count),
99 };
100 lit.set_span(counted.span());
101 out.extend(once(TokenTree::Literal(lit)));
102 }
103
104 if mmatches!{
106 iter.peek(), Some(TokenTree::Group(group))
107 if mmatches!(group.delimiter(), Delimiter::Parenthesis)
108 } {
109 let mut out = TokenStream::new();
110
111 output_counted(parse_parentheses(&mut iter)?, ExpandedInto::Expr, &mut out);
112
113 Ok(out)
114 } else {
115 let mut macro_ = parse_macro_invocation(&mut iter)?;
116
117 output_counted(parse_parentheses(&mut iter)?, ExpandedInto::Macro, &mut macro_.args);
118
119 Ok(macro_.into_token_stream())
120 }
121}
122
123
124pub(crate) fn gen_ident_range(tokens: TokenStream) -> crate::Result<TokenStream> {
125 let mut iter = tokens.into_iter().peekable();
126 let mut macro_ = parse_macro_invocation(&mut iter)?;
127
128 let idents = try_!(gen_ident_range_just_idents(&mut iter, |x| parse_bounded_range_param(x)));
129
130 let paren = Group::new(Delimiter::Parenthesis, idents.collect());
131
132 macro_.args.extend(once(TokenTree::Group(paren)));
133
134 Ok(macro_.into_token_stream())
135}
136
137pub(crate) fn gen_ident_range_just_idents<F>(
138 iter: &mut Peekable<IntoIter>,
139 parse_range: F,
140) -> crate::Result<GenIdentRange>
141where
142 F: FnOnce(&mut Peekable<IntoIter>) -> crate::Result<Range<usize>>
143{
144 try_!(parse_keyword(&mut *iter, "for"));
145
146 let prefix = try_!(parse_ident(&mut *iter));
147 let sprefix = prefix.to_string();
148 let span = prefix.span();
149
150 try_!(parse_check_punct(&mut *iter, '*'));
151
152 try_!(parse_keyword(&mut *iter, "in"));
153
154 let range = try_!(parse_range(&mut *iter));
155
156 try_!(expect_no_tokens(iter));
157
158 Ok(GenIdentRange{sprefix, range, span})
159}
160
161pub(crate) struct GenIdentRange {
162 sprefix: String,
163 range: Range<usize>,
164 span: Span,
165}
166
167impl GenIdentRange{
168 pub(crate) fn span(&self) -> Span {
169 self.span
170 }
171 pub(crate) fn is_unbounded(&self) -> bool {
172 const M: usize = !0;
173 self.range.end == M
174 }
175}
176
177impl Iterator for GenIdentRange {
178 type Item = TokenTree;
179
180 fn next(&mut self) -> Option<TokenTree> {
181 self.range
182 .next()
183 .map(|n| {
184 let ident = Ident::new(&format!("{}{}", self.sprefix, n), self.span);
185 TokenTree::Ident(ident)
186 })
187 }
188}
189
190
191
192pub(crate) fn macro_attr(attr: TokenStream, item: TokenStream) -> crate::Result<TokenStream> {
193 let mut attr = attr.into_iter();
194
195 let mut macro_ = crate::macro_utils_shared::parse_path_and_span(&mut attr)?;
196
197 let (bang, more_tokens) = match macro_.terminator {
198 Some(TokenTree::Punct(punct)) if punct.as_char() == '!' =>
199 (punct, true),
200 Some(tt) =>
201 return Err(crate::Error::one_tt(tt.span(), "expected a `!`")),
202 None => {
203 let mut bang = Punct::new('!', Spacing::Alone);
204 bang.set_span(macro_.start_span);
205 (bang, false)
206 }
207 };
208
209 macro_.path.extend(once(TokenTree::Punct(bang)));
210
211 let (args, bspan) = if more_tokens {
212 let group = crate::macro_utils_shared::parse_group(&mut attr)?;
213 let mut args = group.stream();
214 args.extend(item);
215
216 (args, group.span())
217 } else {
218 (item, macro_.end_span)
219 };
220
221 let mut args = Group::new(Delimiter::Brace, args);
222 args.set_span(bspan);
223 macro_.path.extend(once(TokenTree::Group(args)));
224
225 Ok(macro_.path)
226}
227
228
229pub(crate) fn tokens_method(tokens: TokenStream) -> crate::Result<TokenStream> {
230 let mut iter = tokens.into_iter();
231
232 let mut macro_ = parse_macro_invocation(&mut iter)?;
233 let args = &mut macro_.args;
234
235 macro_rules! declare_methods {
236 (
237 $fname:literal => $fblock:block
238 $( $name:literal => $block:block )*
239 ) => {
240 const ERR_MSG: &str = concat!(
241 "expected one of ",
242 "`", $fname, "`",
243 $(", `", $name, "`",)*
244 "."
245 );
246
247 match parse_ident(&mut iter) {
248 Ok(ident) => {
249 let keyword = ident.to_string();
250
251 match &keyword[..] {
252 $fname => $fblock
253 $($name => $block)*
254 other => {
255 let err = format!("{}\nFound {}", ERR_MSG, other);
256 return Err(crate::Error::one_tt(ident.span(), &err));
257 }
258 }
259 }
260 Err(e) => {
261 return Err(crate::Error::one_tt(e.start_span(), &ERR_MSG));
262 }
263 }
264 };
265 }
266
267 declare_methods!{
268 "first" => {
269 parse_no_params(&mut iter)?;
270 let xx = parse_unbounded(&mut iter)?;
271 let span = xx.spans().start;
272 let last_token: TokenStream = xx.into_iter().take(1).collect();
273
274 out_parenthesized(last_token, span, args);
275 }
276 "last" => {
277 parse_no_params(&mut iter)?;
278 let group = parse_bounded(&mut iter)?;
279
280 let last_token: TokenStream =
281 group.stream().into_iter().last().into_iter().collect();
282
283 out_parenthesized(last_token, group.span(), args);
284 }
285 "split_first" => {
286 parse_no_params(&mut iter)?;
287 let group = parse_bounded(&mut iter)?;
288
289 let mut iter = group.stream().into_iter();
290 let first: TokenStream = (&mut iter).take(1).collect();
291 let rest: TokenStream = iter.collect();
292
293 out_parenthesized(first, group.span(), args);
294 out_parenthesized(rest, group.span(), args);
295 }
296 "split_last" => {
297 parse_no_params(&mut iter)?;
298 let group = parse_bounded(&mut iter)?;
299
300 let mut iter = group.stream().into_iter();
301
302 let mut first = TokenStream::new();
303 let mut last = iter.next();
304 for tt in iter {
305 first.extend(last);
306 last = Some(tt);
307 }
308 let last = last.into_iter().collect::<TokenStream>();
309
310 out_parenthesized(first, group.span(), args);
311 out_parenthesized(last, group.span(), args);
312 }
313 "split_last_n" => {
314 let mut params = parse_params(&mut iter)?.stream().into_iter();
315 let (last_count, _) = parse_count_param(&mut params)?;
316 crate::macro_utils_shared::expect_no_tokens(params)?;
317
318 let group = parse_bounded(&mut iter)?;
319
320 let elems = group.stream().into_iter().collect::<Vec<TokenTree>>();
321
322 let taken = elems.len().saturating_sub(last_count);
323 let mut iter = elems.into_iter();
324 let first = (&mut iter).take(taken).collect::<TokenStream>();
325 let last = iter.collect::<TokenStream>();
326
327 out_parenthesized(first, group.span(), args);
328 out_parenthesized(last, group.span(), args);
329 }
330 "split_at" => {
331 let mut params = parse_params(&mut iter)?.stream().into_iter();
332 let (split_at, _) = parse_count_param(&mut params)?;
333 crate::macro_utils_shared::expect_no_tokens(params)?;
334
335 let group = parse_bounded(&mut iter)?;
336
337 let mut iter = group.stream().into_iter();
338
339 let start: TokenStream = (&mut iter).take(split_at).collect();
340 let rest: TokenStream = iter.collect();
341
342 out_parenthesized(start, group.span(), args);
343 out_parenthesized(rest, group.span(), args);
344 }
345 "get" => {
346 let mut params = parse_params(&mut iter)?.stream().into_iter().peekable();
347 let range = parse_int_or_range_param(&mut params)?;
348 crate::macro_utils_shared::expect_no_tokens(params)?;
349
350 let span: Span;
351 let middle: TokenStream = if let Some(end) = range.end {
352 let xx = parse_unbounded(&mut iter)?;
353 span = xx.spans().start;
354 xx.into_iter().take(end).skip(range.start).collect()
355 } else {
356 let xx = parse_bounded(&mut iter)?;
357 span = xx.span();
358 xx.stream().into_iter().skip(range.start).collect()
359 };
360
361 out_parenthesized(middle, span, args);
362 }
363 "split" => {
364 let (needle, group, mut iter) = split_shared(&mut iter)?;
365 loop {
366 let (tokens, found) = cmp_ts::skip_until_match(&mut iter, &needle);
367 out_parenthesized(tokens, group.span(), args);
368 if let Found::No = found { break }
369 }
370 }
371 "split_terminator" => {
372 let (needle, group, mut iter) = split_shared(&mut iter)?;
373 loop {
374 let (tokens, found) = cmp_ts::skip_until_match(&mut iter, &needle);
375 if mmatches!(found, Found::Yes) || !tokens.is_empty() {
376 out_parenthesized(tokens, group.span(), args);
377 }
378 if let Found::No = found { break }
379 }
380 }
381 "split_starter" => {
382 let (needle, group, mut iter) = split_shared(&mut iter)?;
383
384 let mut start = true;
385 loop {
386 let (tokens, found) = cmp_ts::skip_until_match(&mut iter, &needle);
387 if !start || ( start && (!tokens.is_empty() || mmatches!(found, Found::No))) {
388 out_parenthesized(tokens, group.span(), args);
389 }
390 if let Found::No = found { break }
391 start = false;
392 }
393 }
394 "zip_shortest" => {
395 parse_no_params(&mut iter)?;
396 let ZipArgs{mut iters, ..} = parse_for_zip(iter)?;
397 let outer_span = macro_span();
398
399 'outer: loop {
400 let mut zipped = TokenStream::new();
401 for tt_iter in &mut iters {
402 if let Some(tt) = tt_iter.next() {
403 out_parenthesized_tt(tt, &mut zipped);
404 } else {
405 break 'outer;
406 }
407 }
408 out_parenthesized(zipped, outer_span, args)
409 }
410 }
411 "zip_longest" => {
412 parse_no_params(&mut iter)?;
413 let ZipArgs{mut iters, finite_arg_count} = parse_for_zip(iter)?;
414 let outer_span = macro_span();
415
416 loop {
417 let mut zipped = TokenStream::new();
418
419 let mut none_count = 0;
420 for tt_iter in &mut iters {
421 if let Some(tt) = tt_iter.next() {
422 out_parenthesized_tt(tt, &mut zipped);
423 } else {
424 none_count+=1;
425 out_parenthesized(TokenStream::new(), outer_span, &mut zipped)
426 }
427 }
428 if none_count == finite_arg_count { break }
429
430 out_parenthesized(zipped, outer_span, args)
431 }
432 }
433 "iterate" => {
434 parse_no_params(&mut iter)?;
435 let mut ingroups = parse_bounded_args(iter)?;
436
437 let mut outgroups = VecDeque::<Group>::new();
438 outgroups.push_front(ingroups.pop().unwrap());
439
440 for ingroup in ingroups.iter().rev() {
441 let nested = outgroups.front().unwrap();
442
443 let mut out_elem = TokenStream::new();
444
445 for tt in ingroup.stream() {
446 out_braced_tt(tt, &mut out_elem);
447 out_elem.extend(once(TokenTree::Group(nested.clone())));
448 }
449
450 let mut group = Group::new(Delimiter::Parenthesis, out_elem);
451 group.set_span(ingroup.span());
452 outgroups.push_front(group);
453 }
454
455 args.extend(once(TokenTree::Group(outgroups.pop_front().unwrap())));
456 }
457 }
458
459 Ok(macro_.into_token_stream())
460}
461
462fn parse_params(iter: &mut IntoIter) -> crate::Result<Group> {
463 match_token!{"expected parentheses followed by colon", iter.next() =>
464 Some(TokenTree::Group(group)) if mmatches!(group.delimiter(), Delimiter::Parenthesis) => {
465 parse_no_params(iter)?;
466 Ok(group)
467 }
468 }
469}
470
471fn parse_no_params(iter: &mut IntoIter) -> crate::Result<()> {
472 match_token!{"expected colon", iter.next() =>
473 Some(TokenTree::Punct(p)) if p.as_char() == ':' => {
474 Ok(())
475 }
476 }
477}
478
479fn split_shared(iter: &mut IntoIter) -> crate::Result<(Vec<ComparableTT>, Group, IntoIter)> {
480 let params = parse_params(iter)?;
481 let needle = ComparableTT::many(params.stream().into_iter());
482
483 let group = parse_bounded(&mut *iter)?;
484 let iter = group.stream().into_iter();
485
486 Ok((needle, group, iter))
487}
488
489
490struct ZipArgs {
491 iters: Vec<ListIter>,
492 finite_arg_count: usize,
493}
494
495fn parse_for_zip(iter: IntoIter) -> crate::Result<ZipArgs> {
496 let mut groups = Vec::new();
497 let mut iter = iter.peekable();
498
499 let mut finite_arg_count = 0;
500
501 loop {
502 let elem = try_!(parse_unbounded(&mut iter));
503 if elem.is_finite() {
504 finite_arg_count += 1;
505 }
506 groups.push(elem.into_iter());
507
508 if let None = iter.peek() { break }
509 }
510
511 if finite_arg_count == 0 {
512 Err(crate::Error::one_tt(macro_span(), "Expected at least one finite list"))
513 } else {
514 Ok(ZipArgs{iters: groups, finite_arg_count})
515 }
516
517}
518
519fn parse_bounded_args(iter: IntoIter) -> crate::Result<Vec<Group>> {
520 let mut groups = Vec::new();
521 let mut iter = iter.peekable();
522
523 loop {
524 groups.push(try_!(parse_bounded(&mut iter)));
525 if let None = iter.peek() { break }
526 }
527
528 Ok(groups)
529}
530
531
532