tt/
lib.rs

1//! `tt` crate defines a `TokenTree` data structure: this is the interface (both
2//! input and output) of macros.
3//!
4//! The `TokenTree` is semantically a tree, but for performance reasons it is stored as a flat structure.
5
6#![cfg_attr(feature = "in-rust-tree", feature(rustc_private))]
7
8#[cfg(feature = "in-rust-tree")]
9extern crate rustc_driver as _;
10
11#[cfg(not(feature = "in-rust-tree"))]
12extern crate ra_ap_rustc_lexer as rustc_lexer;
13#[cfg(feature = "in-rust-tree")]
14extern crate rustc_lexer;
15
16pub mod buffer;
17pub mod iter;
18
19use std::fmt;
20
21use buffer::Cursor;
22use intern::Symbol;
23use iter::{TtElement, TtIter};
24use stdx::{impl_from, itertools::Itertools as _};
25
26pub use text_size::{TextRange, TextSize};
27
28pub const MAX_GLUED_PUNCT_LEN: usize = 3;
29
30#[derive(Clone, PartialEq, Debug)]
31pub struct Lit {
32    pub kind: LitKind,
33    pub symbol: Symbol,
34    pub suffix: Option<Symbol>,
35}
36
37#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
38pub enum IdentIsRaw {
39    No,
40    Yes,
41}
42impl IdentIsRaw {
43    pub fn yes(self) -> bool {
44        matches!(self, IdentIsRaw::Yes)
45    }
46    pub fn no(&self) -> bool {
47        matches!(self, IdentIsRaw::No)
48    }
49    pub fn as_str(self) -> &'static str {
50        match self {
51            IdentIsRaw::No => "",
52            IdentIsRaw::Yes => "r#",
53        }
54    }
55    pub fn split_from_symbol(sym: &str) -> (Self, &str) {
56        if let Some(sym) = sym.strip_prefix("r#") {
57            (IdentIsRaw::Yes, sym)
58        } else {
59            (IdentIsRaw::No, sym)
60        }
61    }
62}
63
64#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
65pub enum LitKind {
66    Byte,
67    Char,
68    Integer, // e.g. `1`, `1u8`, `1f32`
69    Float,   // e.g. `1.`, `1.0`, `1e3f32`
70    Str,
71    StrRaw(u8), // raw string delimited by `n` hash symbols
72    ByteStr,
73    ByteStrRaw(u8), // raw byte string delimited by `n` hash symbols
74    CStr,
75    CStrRaw(u8),
76    Err(()),
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Hash)]
80pub enum TokenTree<S = u32> {
81    Leaf(Leaf<S>),
82    Subtree(Subtree<S>),
83}
84impl_from!(Leaf<S>, Subtree<S> for TokenTree);
85impl<S: Copy> TokenTree<S> {
86    pub fn first_span(&self) -> S {
87        match self {
88            TokenTree::Leaf(l) => *l.span(),
89            TokenTree::Subtree(s) => s.delimiter.open,
90        }
91    }
92}
93
94#[derive(Debug, Clone, PartialEq, Eq, Hash)]
95pub enum Leaf<S> {
96    Literal(Literal<S>),
97    Punct(Punct<S>),
98    Ident(Ident<S>),
99}
100
101impl<S> Leaf<S> {
102    pub fn span(&self) -> &S {
103        match self {
104            Leaf::Literal(it) => &it.span,
105            Leaf::Punct(it) => &it.span,
106            Leaf::Ident(it) => &it.span,
107        }
108    }
109}
110impl_from!(Literal<S>, Punct<S>, Ident<S> for Leaf);
111
112#[derive(Debug, Clone, PartialEq, Eq, Hash)]
113pub struct Subtree<S> {
114    pub delimiter: Delimiter<S>,
115    /// Number of following token trees that belong to this subtree, excluding this subtree.
116    pub len: u32,
117}
118
119impl<S> Subtree<S> {
120    pub fn usize_len(&self) -> usize {
121        self.len as usize
122    }
123}
124
125#[derive(Clone, PartialEq, Eq, Hash)]
126pub struct TopSubtree<S>(pub Box<[TokenTree<S>]>);
127
128impl<S: Copy> TopSubtree<S> {
129    pub fn empty(span: DelimSpan<S>) -> Self {
130        Self(Box::new([TokenTree::Subtree(Subtree {
131            delimiter: Delimiter::invisible_delim_spanned(span),
132            len: 0,
133        })]))
134    }
135
136    pub fn invisible_from_leaves<const N: usize>(delim_span: S, leaves: [Leaf<S>; N]) -> Self {
137        let mut builder = TopSubtreeBuilder::new(Delimiter::invisible_spanned(delim_span));
138        builder.extend(leaves);
139        builder.build()
140    }
141
142    pub fn from_token_trees(delimiter: Delimiter<S>, token_trees: TokenTreesView<'_, S>) -> Self {
143        let mut builder = TopSubtreeBuilder::new(delimiter);
144        builder.extend_with_tt(token_trees);
145        builder.build()
146    }
147
148    pub fn from_subtree(subtree: SubtreeView<'_, S>) -> Self {
149        Self(subtree.0.into())
150    }
151
152    pub fn view(&self) -> SubtreeView<'_, S> {
153        SubtreeView::new(&self.0)
154    }
155
156    pub fn iter(&self) -> TtIter<'_, S> {
157        self.view().iter()
158    }
159
160    pub fn top_subtree(&self) -> &Subtree<S> {
161        self.view().top_subtree()
162    }
163
164    pub fn top_subtree_delimiter_mut(&mut self) -> &mut Delimiter<S> {
165        let TokenTree::Subtree(subtree) = &mut self.0[0] else {
166            unreachable!("the first token tree is always the top subtree");
167        };
168        &mut subtree.delimiter
169    }
170
171    pub fn token_trees(&self) -> TokenTreesView<'_, S> {
172        self.view().token_trees()
173    }
174}
175
176#[derive(Debug, Clone, PartialEq, Eq, Hash)]
177pub struct TopSubtreeBuilder<S> {
178    unclosed_subtree_indices: Vec<usize>,
179    token_trees: Vec<TokenTree<S>>,
180    last_closed_subtree: Option<usize>,
181}
182
183impl<S: Copy> TopSubtreeBuilder<S> {
184    pub fn new(top_delimiter: Delimiter<S>) -> Self {
185        let mut result = Self {
186            unclosed_subtree_indices: Vec::new(),
187            token_trees: Vec::new(),
188            last_closed_subtree: None,
189        };
190        let top_subtree = TokenTree::Subtree(Subtree { delimiter: top_delimiter, len: 0 });
191        result.token_trees.push(top_subtree);
192        result
193    }
194
195    pub fn open(&mut self, delimiter_kind: DelimiterKind, open_span: S) {
196        self.unclosed_subtree_indices.push(self.token_trees.len());
197        self.token_trees.push(TokenTree::Subtree(Subtree {
198            delimiter: Delimiter {
199                open: open_span,
200                close: open_span, // Will be overwritten on close.
201                kind: delimiter_kind,
202            },
203            len: 0,
204        }));
205    }
206
207    pub fn close(&mut self, close_span: S) {
208        let last_unclosed_index = self
209            .unclosed_subtree_indices
210            .pop()
211            .expect("attempt to close a `tt::Subtree` when none is open");
212        let subtree_len = (self.token_trees.len() - last_unclosed_index - 1) as u32;
213        let TokenTree::Subtree(subtree) = &mut self.token_trees[last_unclosed_index] else {
214            unreachable!("unclosed token tree is always a subtree");
215        };
216        subtree.len = subtree_len;
217        subtree.delimiter.close = close_span;
218        self.last_closed_subtree = Some(last_unclosed_index);
219    }
220
221    /// You cannot call this consecutively, it will only work once after close.
222    pub fn remove_last_subtree_if_invisible(&mut self) {
223        let Some(last_subtree_idx) = self.last_closed_subtree else { return };
224        if let TokenTree::Subtree(Subtree {
225            delimiter: Delimiter { kind: DelimiterKind::Invisible, .. },
226            ..
227        }) = self.token_trees[last_subtree_idx]
228        {
229            self.token_trees.remove(last_subtree_idx);
230            self.last_closed_subtree = None;
231        }
232    }
233
234    pub fn push(&mut self, leaf: Leaf<S>) {
235        self.token_trees.push(TokenTree::Leaf(leaf));
236    }
237
238    pub fn extend(&mut self, leaves: impl IntoIterator<Item = Leaf<S>>) {
239        self.token_trees.extend(leaves.into_iter().map(TokenTree::Leaf));
240    }
241
242    /// This does not check the token trees are valid, beware!
243    pub fn extend_tt_dangerous(&mut self, tt: impl IntoIterator<Item = TokenTree<S>>) {
244        self.token_trees.extend(tt);
245    }
246
247    pub fn extend_with_tt(&mut self, tt: TokenTreesView<'_, S>) {
248        self.token_trees.extend(tt.0.iter().cloned());
249    }
250
251    /// Like [`Self::extend_with_tt()`], but makes sure the new tokens will never be
252    /// joint with whatever comes after them.
253    pub fn extend_with_tt_alone(&mut self, tt: TokenTreesView<'_, S>) {
254        if let Some((last, before_last)) = tt.0.split_last() {
255            self.token_trees.reserve(tt.0.len());
256            self.token_trees.extend(before_last.iter().cloned());
257            let last = if let TokenTree::Leaf(Leaf::Punct(last)) = last {
258                let mut last = *last;
259                last.spacing = Spacing::Alone;
260                TokenTree::Leaf(Leaf::Punct(last))
261            } else {
262                last.clone()
263            };
264            self.token_trees.push(last);
265        }
266    }
267
268    pub fn expected_delimiters(&self) -> impl Iterator<Item = &Delimiter<S>> {
269        self.unclosed_subtree_indices.iter().rev().map(|&subtree_idx| {
270            let TokenTree::Subtree(subtree) = &self.token_trees[subtree_idx] else {
271                unreachable!("unclosed token tree is always a subtree")
272            };
273            &subtree.delimiter
274        })
275    }
276
277    /// Builds, and remove the top subtree if it has only one subtree child.
278    pub fn build_skip_top_subtree(mut self) -> TopSubtree<S> {
279        let top_tts = TokenTreesView::new(&self.token_trees[1..]);
280        match top_tts.try_into_subtree() {
281            Some(_) => {
282                assert!(
283                    self.unclosed_subtree_indices.is_empty(),
284                    "attempt to build an unbalanced `TopSubtreeBuilder`"
285                );
286                TopSubtree(self.token_trees.drain(1..).collect())
287            }
288            None => self.build(),
289        }
290    }
291
292    pub fn build(mut self) -> TopSubtree<S> {
293        assert!(
294            self.unclosed_subtree_indices.is_empty(),
295            "attempt to build an unbalanced `TopSubtreeBuilder`"
296        );
297        let total_len = self.token_trees.len() as u32;
298        let TokenTree::Subtree(top_subtree) = &mut self.token_trees[0] else {
299            unreachable!("first token tree is always a subtree");
300        };
301        top_subtree.len = total_len - 1;
302        TopSubtree(self.token_trees.into_boxed_slice())
303    }
304
305    pub fn restore_point(&self) -> SubtreeBuilderRestorePoint {
306        SubtreeBuilderRestorePoint {
307            unclosed_subtree_indices_len: self.unclosed_subtree_indices.len(),
308            token_trees_len: self.token_trees.len(),
309            last_closed_subtree: self.last_closed_subtree,
310        }
311    }
312
313    pub fn restore(&mut self, restore_point: SubtreeBuilderRestorePoint) {
314        self.unclosed_subtree_indices.truncate(restore_point.unclosed_subtree_indices_len);
315        self.token_trees.truncate(restore_point.token_trees_len);
316        self.last_closed_subtree = restore_point.last_closed_subtree;
317    }
318}
319
320#[derive(Clone, Copy)]
321pub struct SubtreeBuilderRestorePoint {
322    unclosed_subtree_indices_len: usize,
323    token_trees_len: usize,
324    last_closed_subtree: Option<usize>,
325}
326
327#[derive(Clone, Copy)]
328pub struct TokenTreesView<'a, S>(&'a [TokenTree<S>]);
329
330impl<'a, S: Copy> TokenTreesView<'a, S> {
331    pub fn new(tts: &'a [TokenTree<S>]) -> Self {
332        if cfg!(debug_assertions) {
333            tts.iter().enumerate().for_each(|(idx, tt)| {
334                if let TokenTree::Subtree(tt) = &tt {
335                    // `<` and not `<=` because `Subtree.len` does not include the subtree node itself.
336                    debug_assert!(
337                        idx + tt.usize_len() < tts.len(),
338                        "`TokenTreeView::new()` was given a cut-in-half list"
339                    );
340                }
341            });
342        }
343        Self(tts)
344    }
345
346    pub fn iter(&self) -> TtIter<'a, S> {
347        TtIter::new(self.0)
348    }
349
350    pub fn cursor(&self) -> Cursor<'a, S> {
351        Cursor::new(self.0)
352    }
353
354    pub fn len(&self) -> usize {
355        self.0.len()
356    }
357
358    pub fn is_empty(&self) -> bool {
359        self.0.is_empty()
360    }
361
362    pub fn try_into_subtree(self) -> Option<SubtreeView<'a, S>> {
363        if let Some(TokenTree::Subtree(subtree)) = self.0.first()
364            && subtree.usize_len() == (self.0.len() - 1)
365        {
366            return Some(SubtreeView::new(self.0));
367        }
368        None
369    }
370
371    pub fn strip_invisible(self) -> TokenTreesView<'a, S> {
372        self.try_into_subtree().map(|subtree| subtree.strip_invisible()).unwrap_or(self)
373    }
374
375    /// This returns a **flat** structure of tokens (subtrees will be represented by a single node
376    /// preceding their children), so it isn't suited for most use cases, only for matching leaves
377    /// at the beginning/end with no subtrees before them. If you need a structured pass, use [`TtIter`].
378    pub fn flat_tokens(&self) -> &'a [TokenTree<S>] {
379        self.0
380    }
381
382    pub fn split(
383        self,
384        mut split_fn: impl FnMut(TtElement<'a, S>) -> bool,
385    ) -> impl Iterator<Item = TokenTreesView<'a, S>> {
386        let mut subtree_iter = self.iter();
387        let mut need_to_yield_even_if_empty = true;
388
389        std::iter::from_fn(move || {
390            if subtree_iter.is_empty() && !need_to_yield_even_if_empty {
391                return None;
392            };
393
394            need_to_yield_even_if_empty = false;
395            let savepoint = subtree_iter.savepoint();
396            let mut result = subtree_iter.from_savepoint(savepoint);
397            while let Some(tt) = subtree_iter.next() {
398                if split_fn(tt) {
399                    need_to_yield_even_if_empty = true;
400                    break;
401                }
402                result = subtree_iter.from_savepoint(savepoint);
403            }
404            Some(result)
405        })
406    }
407}
408
409impl<S: fmt::Debug + Copy> fmt::Debug for TokenTreesView<'_, S> {
410    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
411        let mut iter = self.iter();
412        while let Some(tt) = iter.next() {
413            print_debug_token(f, 0, tt)?;
414            if !iter.is_empty() {
415                writeln!(f)?;
416            }
417        }
418        Ok(())
419    }
420}
421
422impl<S: Copy> fmt::Display for TokenTreesView<'_, S> {
423    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
424        return token_trees_display(f, self.iter());
425
426        fn subtree_display<S>(
427            subtree: &Subtree<S>,
428            f: &mut fmt::Formatter<'_>,
429            iter: TtIter<'_, S>,
430        ) -> fmt::Result {
431            let (l, r) = match subtree.delimiter.kind {
432                DelimiterKind::Parenthesis => ("(", ")"),
433                DelimiterKind::Brace => ("{", "}"),
434                DelimiterKind::Bracket => ("[", "]"),
435                DelimiterKind::Invisible => ("", ""),
436            };
437            f.write_str(l)?;
438            token_trees_display(f, iter)?;
439            f.write_str(r)?;
440            Ok(())
441        }
442
443        fn token_trees_display<S>(f: &mut fmt::Formatter<'_>, iter: TtIter<'_, S>) -> fmt::Result {
444            let mut needs_space = false;
445            for child in iter {
446                if needs_space {
447                    f.write_str(" ")?;
448                }
449                needs_space = true;
450
451                match child {
452                    TtElement::Leaf(Leaf::Punct(p)) => {
453                        needs_space = p.spacing == Spacing::Alone;
454                        fmt::Display::fmt(p, f)?;
455                    }
456                    TtElement::Leaf(leaf) => fmt::Display::fmt(leaf, f)?,
457                    TtElement::Subtree(subtree, subtree_iter) => {
458                        subtree_display(subtree, f, subtree_iter)?
459                    }
460                }
461            }
462            Ok(())
463        }
464    }
465}
466
467#[derive(Clone, Copy)]
468// Invariant: always starts with `Subtree` that covers the entire thing.
469pub struct SubtreeView<'a, S>(&'a [TokenTree<S>]);
470
471impl<'a, S: Copy> SubtreeView<'a, S> {
472    pub fn new(tts: &'a [TokenTree<S>]) -> Self {
473        if cfg!(debug_assertions) {
474            let TokenTree::Subtree(subtree) = &tts[0] else {
475                panic!("first token tree must be a subtree in `SubtreeView`");
476            };
477            assert_eq!(
478                subtree.usize_len(),
479                tts.len() - 1,
480                "subtree must cover the entire `SubtreeView`"
481            );
482        }
483        Self(tts)
484    }
485
486    pub fn as_token_trees(self) -> TokenTreesView<'a, S> {
487        TokenTreesView::new(self.0)
488    }
489
490    pub fn iter(&self) -> TtIter<'a, S> {
491        TtIter::new(&self.0[1..])
492    }
493
494    pub fn top_subtree(&self) -> &'a Subtree<S> {
495        let TokenTree::Subtree(subtree) = &self.0[0] else {
496            unreachable!("the first token tree is always the top subtree");
497        };
498        subtree
499    }
500
501    pub fn strip_invisible(&self) -> TokenTreesView<'a, S> {
502        if self.top_subtree().delimiter.kind == DelimiterKind::Invisible {
503            TokenTreesView::new(&self.0[1..])
504        } else {
505            TokenTreesView::new(self.0)
506        }
507    }
508
509    pub fn token_trees(&self) -> TokenTreesView<'a, S> {
510        TokenTreesView::new(&self.0[1..])
511    }
512}
513
514impl<S: fmt::Debug + Copy> fmt::Debug for SubtreeView<'_, S> {
515    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
516        fmt::Debug::fmt(&TokenTreesView(self.0), f)
517    }
518}
519
520impl<S: Copy> fmt::Display for SubtreeView<'_, S> {
521    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
522        fmt::Display::fmt(&TokenTreesView(self.0), f)
523    }
524}
525
526#[derive(Debug, Copy, Clone, PartialEq)]
527pub struct DelimSpan<S> {
528    pub open: S,
529    pub close: S,
530}
531
532impl<Span: Copy> DelimSpan<Span> {
533    pub fn from_single(sp: Span) -> Self {
534        DelimSpan { open: sp, close: sp }
535    }
536
537    pub fn from_pair(open: Span, close: Span) -> Self {
538        DelimSpan { open, close }
539    }
540}
541#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
542pub struct Delimiter<S> {
543    pub open: S,
544    pub close: S,
545    pub kind: DelimiterKind,
546}
547
548impl<S: Copy> Delimiter<S> {
549    pub const fn invisible_spanned(span: S) -> Self {
550        Delimiter { open: span, close: span, kind: DelimiterKind::Invisible }
551    }
552
553    pub const fn invisible_delim_spanned(span: DelimSpan<S>) -> Self {
554        Delimiter { open: span.open, close: span.close, kind: DelimiterKind::Invisible }
555    }
556
557    pub fn delim_span(&self) -> DelimSpan<S> {
558        DelimSpan { open: self.open, close: self.close }
559    }
560}
561
562#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
563pub enum DelimiterKind {
564    Parenthesis,
565    Brace,
566    Bracket,
567    Invisible,
568}
569
570#[derive(Debug, Clone, PartialEq, Eq, Hash)]
571pub struct Literal<S> {
572    // escaped
573    pub symbol: Symbol,
574    pub span: S,
575    pub kind: LitKind,
576    pub suffix: Option<Symbol>,
577}
578
579pub fn token_to_literal<S>(text: &str, span: S) -> Literal<S>
580where
581    S: Copy,
582{
583    use rustc_lexer::LiteralKind;
584
585    let token = rustc_lexer::tokenize(text, rustc_lexer::FrontmatterAllowed::No).next_tuple();
586    let Some((rustc_lexer::Token {
587        kind: rustc_lexer::TokenKind::Literal { kind, suffix_start },
588        ..
589    },)) = token
590    else {
591        return Literal {
592            span,
593            symbol: Symbol::intern(text),
594            kind: LitKind::Err(()),
595            suffix: None,
596        };
597    };
598
599    let (kind, start_offset, end_offset) = match kind {
600        LiteralKind::Int { .. } => (LitKind::Integer, 0, 0),
601        LiteralKind::Float { .. } => (LitKind::Float, 0, 0),
602        LiteralKind::Char { terminated } => (LitKind::Char, 1, terminated as usize),
603        LiteralKind::Byte { terminated } => (LitKind::Byte, 2, terminated as usize),
604        LiteralKind::Str { terminated } => (LitKind::Str, 1, terminated as usize),
605        LiteralKind::ByteStr { terminated } => (LitKind::ByteStr, 2, terminated as usize),
606        LiteralKind::CStr { terminated } => (LitKind::CStr, 2, terminated as usize),
607        LiteralKind::RawStr { n_hashes } => (
608            LitKind::StrRaw(n_hashes.unwrap_or_default()),
609            2 + n_hashes.unwrap_or_default() as usize,
610            1 + n_hashes.unwrap_or_default() as usize,
611        ),
612        LiteralKind::RawByteStr { n_hashes } => (
613            LitKind::ByteStrRaw(n_hashes.unwrap_or_default()),
614            3 + n_hashes.unwrap_or_default() as usize,
615            1 + n_hashes.unwrap_or_default() as usize,
616        ),
617        LiteralKind::RawCStr { n_hashes } => (
618            LitKind::CStrRaw(n_hashes.unwrap_or_default()),
619            3 + n_hashes.unwrap_or_default() as usize,
620            1 + n_hashes.unwrap_or_default() as usize,
621        ),
622    };
623
624    let (lit, suffix) = text.split_at(suffix_start as usize);
625    let lit = &lit[start_offset..lit.len() - end_offset];
626    let suffix = match suffix {
627        "" | "_" => None,
628        // ill-suffixed literals
629        _ if !matches!(kind, LitKind::Integer | LitKind::Float | LitKind::Err(_)) => {
630            return Literal {
631                span,
632                symbol: Symbol::intern(text),
633                kind: LitKind::Err(()),
634                suffix: None,
635            };
636        }
637        suffix => Some(Symbol::intern(suffix)),
638    };
639
640    Literal { span, symbol: Symbol::intern(lit), kind, suffix }
641}
642
643#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
644pub struct Punct<S> {
645    pub char: char,
646    pub spacing: Spacing,
647    pub span: S,
648}
649
650/// Indicates whether a token can join with the following token to form a
651/// compound token. Used for conversions to `proc_macro::Spacing`. Also used to
652/// guide pretty-printing, which is where the `JointHidden` value (which isn't
653/// part of `proc_macro::Spacing`) comes in useful.
654#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
655pub enum Spacing {
656    /// The token cannot join with the following token to form a compound
657    /// token.
658    ///
659    /// In token streams parsed from source code, the compiler will use `Alone`
660    /// for any token immediately followed by whitespace, a non-doc comment, or
661    /// EOF.
662    ///
663    /// When constructing token streams within the compiler, use this for each
664    /// token that (a) should be pretty-printed with a space after it, or (b)
665    /// is the last token in the stream. (In the latter case the choice of
666    /// spacing doesn't matter because it is never used for the last token. We
667    /// arbitrarily use `Alone`.)
668    ///
669    /// Converts to `proc_macro::Spacing::Alone`, and
670    /// `proc_macro::Spacing::Alone` converts back to this.
671    Alone,
672
673    /// The token can join with the following token to form a compound token.
674    ///
675    /// In token streams parsed from source code, the compiler will use `Joint`
676    /// for any token immediately followed by punctuation (as determined by
677    /// `Token::is_punct`).
678    ///
679    /// When constructing token streams within the compiler, use this for each
680    /// token that (a) should be pretty-printed without a space after it, and
681    /// (b) is followed by a punctuation token.
682    ///
683    /// Converts to `proc_macro::Spacing::Joint`, and
684    /// `proc_macro::Spacing::Joint` converts back to this.
685    Joint,
686
687    /// The token can join with the following token to form a compound token,
688    /// but this will not be visible at the proc macro level. (This is what the
689    /// `Hidden` means; see below.)
690    ///
691    /// In token streams parsed from source code, the compiler will use
692    /// `JointHidden` for any token immediately followed by anything not
693    /// covered by the `Alone` and `Joint` cases: an identifier, lifetime,
694    /// literal, delimiter, doc comment.
695    ///
696    /// When constructing token streams, use this for each token that (a)
697    /// should be pretty-printed without a space after it, and (b) is followed
698    /// by a non-punctuation token.
699    ///
700    /// Converts to `proc_macro::Spacing::Alone`, but
701    /// `proc_macro::Spacing::Alone` converts back to `token::Spacing::Alone`.
702    /// Because of that, pretty-printing of `TokenStream`s produced by proc
703    /// macros is unavoidably uglier (with more whitespace between tokens) than
704    /// pretty-printing of `TokenStream`'s produced by other means (i.e. parsed
705    /// source code, internally constructed token streams, and token streams
706    /// produced by declarative macros).
707    JointHidden,
708}
709
710/// Identifier or keyword.
711#[derive(Debug, Clone, PartialEq, Eq, Hash)]
712pub struct Ident<S> {
713    pub sym: Symbol,
714    pub span: S,
715    pub is_raw: IdentIsRaw,
716}
717
718impl<S> Ident<S> {
719    pub fn new(text: &str, span: S) -> Self {
720        // let raw_stripped = IdentIsRaw::split_from_symbol(text.as_ref());
721        let (is_raw, text) = IdentIsRaw::split_from_symbol(text);
722        Ident { sym: Symbol::intern(text), span, is_raw }
723    }
724}
725
726fn print_debug_subtree<S: fmt::Debug>(
727    f: &mut fmt::Formatter<'_>,
728    subtree: &Subtree<S>,
729    level: usize,
730    iter: TtIter<'_, S>,
731) -> fmt::Result {
732    let align = "  ".repeat(level);
733
734    let Delimiter { kind, open, close } = &subtree.delimiter;
735    let delim = match kind {
736        DelimiterKind::Invisible => "$$",
737        DelimiterKind::Parenthesis => "()",
738        DelimiterKind::Brace => "{}",
739        DelimiterKind::Bracket => "[]",
740    };
741
742    write!(f, "{align}SUBTREE {delim} ",)?;
743    write!(f, "{open:#?}")?;
744    write!(f, " ")?;
745    write!(f, "{close:#?}")?;
746    for child in iter {
747        writeln!(f)?;
748        print_debug_token(f, level + 1, child)?;
749    }
750
751    Ok(())
752}
753
754fn print_debug_token<S: fmt::Debug>(
755    f: &mut fmt::Formatter<'_>,
756    level: usize,
757    tt: TtElement<'_, S>,
758) -> fmt::Result {
759    let align = "  ".repeat(level);
760
761    match tt {
762        TtElement::Leaf(leaf) => match leaf {
763            Leaf::Literal(lit) => {
764                write!(
765                    f,
766                    "{}LITERAL {:?} {}{} {:#?}",
767                    align,
768                    lit.kind,
769                    lit.symbol,
770                    lit.suffix.as_ref().map(|it| it.as_str()).unwrap_or(""),
771                    lit.span
772                )?;
773            }
774            Leaf::Punct(punct) => {
775                write!(
776                    f,
777                    "{}PUNCH   {} [{}] {:#?}",
778                    align,
779                    punct.char,
780                    if punct.spacing == Spacing::Alone { "alone" } else { "joint" },
781                    punct.span
782                )?;
783            }
784            Leaf::Ident(ident) => {
785                write!(
786                    f,
787                    "{}IDENT   {}{} {:#?}",
788                    align,
789                    ident.is_raw.as_str(),
790                    ident.sym,
791                    ident.span
792                )?;
793            }
794        },
795        TtElement::Subtree(subtree, subtree_iter) => {
796            print_debug_subtree(f, subtree, level, subtree_iter)?;
797        }
798    }
799
800    Ok(())
801}
802
803impl<S: fmt::Debug + Copy> fmt::Debug for TopSubtree<S> {
804    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
805        fmt::Debug::fmt(&self.view(), f)
806    }
807}
808
809impl<S: fmt::Display + Copy> fmt::Display for TopSubtree<S> {
810    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
811        fmt::Display::fmt(&self.view(), f)
812    }
813}
814
815impl<S> fmt::Display for Leaf<S> {
816    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
817        match self {
818            Leaf::Ident(it) => fmt::Display::fmt(it, f),
819            Leaf::Literal(it) => fmt::Display::fmt(it, f),
820            Leaf::Punct(it) => fmt::Display::fmt(it, f),
821        }
822    }
823}
824
825impl<S> fmt::Display for Ident<S> {
826    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
827        fmt::Display::fmt(&self.is_raw.as_str(), f)?;
828        fmt::Display::fmt(&self.sym, f)
829    }
830}
831
832impl<S> fmt::Display for Literal<S> {
833    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
834        match self.kind {
835            LitKind::Byte => write!(f, "b'{}'", self.symbol),
836            LitKind::Char => write!(f, "'{}'", self.symbol),
837            LitKind::Integer | LitKind::Float | LitKind::Err(_) => write!(f, "{}", self.symbol),
838            LitKind::Str => write!(f, "\"{}\"", self.symbol),
839            LitKind::ByteStr => write!(f, "b\"{}\"", self.symbol),
840            LitKind::CStr => write!(f, "c\"{}\"", self.symbol),
841            LitKind::StrRaw(num_of_hashes) => {
842                let num_of_hashes = num_of_hashes as usize;
843                write!(
844                    f,
845                    r#"r{0:#<num_of_hashes$}"{text}"{0:#<num_of_hashes$}"#,
846                    "",
847                    text = self.symbol
848                )
849            }
850            LitKind::ByteStrRaw(num_of_hashes) => {
851                let num_of_hashes = num_of_hashes as usize;
852                write!(
853                    f,
854                    r#"br{0:#<num_of_hashes$}"{text}"{0:#<num_of_hashes$}"#,
855                    "",
856                    text = self.symbol
857                )
858            }
859            LitKind::CStrRaw(num_of_hashes) => {
860                let num_of_hashes = num_of_hashes as usize;
861                write!(
862                    f,
863                    r#"cr{0:#<num_of_hashes$}"{text}"{0:#<num_of_hashes$}"#,
864                    "",
865                    text = self.symbol
866                )
867            }
868        }?;
869        if let Some(suffix) = &self.suffix {
870            write!(f, "{suffix}")?;
871        }
872        Ok(())
873    }
874}
875
876impl<S> fmt::Display for Punct<S> {
877    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
878        fmt::Display::fmt(&self.char, f)
879    }
880}
881
882impl<S> Subtree<S> {
883    /// Count the number of tokens recursively
884    pub fn count(&self) -> usize {
885        self.usize_len()
886    }
887}
888
889impl<S> TopSubtree<S> {
890    /// A simple line string used for debugging
891    pub fn subtree_as_debug_string(&self, subtree_idx: usize) -> String {
892        fn debug_subtree<S>(
893            output: &mut String,
894            subtree: &Subtree<S>,
895            iter: &mut std::slice::Iter<'_, TokenTree<S>>,
896        ) {
897            let delim = match subtree.delimiter.kind {
898                DelimiterKind::Brace => ("{", "}"),
899                DelimiterKind::Bracket => ("[", "]"),
900                DelimiterKind::Parenthesis => ("(", ")"),
901                DelimiterKind::Invisible => ("$", "$"),
902            };
903
904            output.push_str(delim.0);
905            let mut last = None;
906            let mut idx = 0;
907            while idx < subtree.len {
908                let child = iter.next().unwrap();
909                debug_token_tree(output, child, last, iter);
910                last = Some(child);
911                idx += 1;
912            }
913
914            output.push_str(delim.1);
915        }
916
917        fn debug_token_tree<S>(
918            output: &mut String,
919            tt: &TokenTree<S>,
920            last: Option<&TokenTree<S>>,
921            iter: &mut std::slice::Iter<'_, TokenTree<S>>,
922        ) {
923            match tt {
924                TokenTree::Leaf(it) => {
925                    let s = match it {
926                        Leaf::Literal(it) => it.symbol.to_string(),
927                        Leaf::Punct(it) => it.char.to_string(),
928                        Leaf::Ident(it) => format!("{}{}", it.is_raw.as_str(), it.sym),
929                    };
930                    match (it, last) {
931                        (Leaf::Ident(_), Some(&TokenTree::Leaf(Leaf::Ident(_)))) => {
932                            output.push(' ');
933                            output.push_str(&s);
934                        }
935                        (Leaf::Punct(_), Some(TokenTree::Leaf(Leaf::Punct(punct)))) => {
936                            if punct.spacing == Spacing::Alone {
937                                output.push(' ');
938                                output.push_str(&s);
939                            } else {
940                                output.push_str(&s);
941                            }
942                        }
943                        _ => output.push_str(&s),
944                    }
945                }
946                TokenTree::Subtree(it) => debug_subtree(output, it, iter),
947            }
948        }
949
950        let mut res = String::new();
951        debug_token_tree(
952            &mut res,
953            &self.0[subtree_idx],
954            None,
955            &mut self.0[subtree_idx + 1..].iter(),
956        );
957        res
958    }
959}
960
961pub fn pretty<S>(mut tkns: &[TokenTree<S>]) -> String {
962    fn tokentree_to_text<S>(tkn: &TokenTree<S>, tkns: &mut &[TokenTree<S>]) -> String {
963        match tkn {
964            TokenTree::Leaf(Leaf::Ident(ident)) => {
965                format!("{}{}", ident.is_raw.as_str(), ident.sym)
966            }
967            TokenTree::Leaf(Leaf::Literal(literal)) => format!("{literal}"),
968            TokenTree::Leaf(Leaf::Punct(punct)) => format!("{}", punct.char),
969            TokenTree::Subtree(subtree) => {
970                let (subtree_content, rest) = tkns.split_at(subtree.usize_len());
971                let content = pretty(subtree_content);
972                *tkns = rest;
973                let (open, close) = match subtree.delimiter.kind {
974                    DelimiterKind::Brace => ("{", "}"),
975                    DelimiterKind::Bracket => ("[", "]"),
976                    DelimiterKind::Parenthesis => ("(", ")"),
977                    DelimiterKind::Invisible => ("", ""),
978                };
979                format!("{open}{content}{close}")
980            }
981        }
982    }
983
984    let mut last = String::new();
985    let mut last_to_joint = true;
986
987    while let Some((tkn, rest)) = tkns.split_first() {
988        tkns = rest;
989        last = [last, tokentree_to_text(tkn, &mut tkns)].join(if last_to_joint { "" } else { " " });
990        last_to_joint = false;
991        if let TokenTree::Leaf(Leaf::Punct(punct)) = tkn
992            && punct.spacing == Spacing::Joint
993        {
994            last_to_joint = true;
995        }
996    }
997    last
998}