proc_macro_tool/
parse_iter.rs

1use crate::TokenStreamExt as _;
2use proc_macro::{Spacing::*, TokenStream, TokenTree};
3use std::{collections::VecDeque, iter::FusedIterator};
4
5/// Create [`ParseIter`]
6pub trait ParseIterExt: IntoIterator<Item = TokenTree> + Sized {
7    /// Create [`ParseIter`]
8    fn parse_iter(self) -> ParseIter<Self::IntoIter> {
9        ParseIter { iter: self.into_iter(), buf: VecDeque::new() }
10    }
11}
12impl<I: IntoIterator<Item = TokenTree>> ParseIterExt for I { }
13
14/// Peek `n` iterator adapter
15#[derive(Debug, Clone)]
16pub struct ParseIter<I: Iterator<Item = TokenTree>> {
17    iter: I,
18    buf: VecDeque<TokenTree>,
19}
20
21impl<I: Iterator<Item = TokenTree>> ParseIter<I> {
22    pub fn peek(&mut self) -> Option<&TokenTree> {
23        self.peek_i(0)
24    }
25
26    pub fn next_if<F>(&mut self, f: F) -> Option<TokenTree>
27    where F: FnOnce(&TokenTree) -> bool,
28    {
29        let peek = self.peek()?;
30
31        if f(peek) {
32            self.next()
33        } else {
34            None
35        }
36    }
37
38    pub fn peek_i(&mut self, i: usize) -> Option<&TokenTree> {
39        for _ in self.buf.len()..=i {
40            self.buf.push_back(self.iter.next()?);
41        }
42        Some(&self.buf[i])
43    }
44
45    pub fn peek_is<F>(&mut self, f: F) -> bool
46    where F: FnOnce(&TokenTree) -> bool,
47    {
48        self.peek().is_some_and(f)
49    }
50
51    pub fn peek_i_is<F>(&mut self, i: usize, f: F) -> bool
52    where F: FnOnce(&TokenTree) -> bool,
53    {
54        self.peek_i(i).is_some_and(f)
55    }
56
57    /// Peek jointed puncts
58    pub fn peek_puncts(&mut self, puncts: impl AsRef<[u8]>) -> Option<
59        impl Iterator<Item = &TokenTree>
60    > {
61        let mut prev = None;
62
63        for (i, ch) in puncts.as_ref().iter()
64            .copied().map(char::from).enumerate()
65        {
66            if let Some(prev) = prev {
67                if prev == Alone { return None }
68            }
69
70            let tt = self.peek_i(i)?;
71            let TokenTree::Punct(p) = tt else { return None };
72            if p.as_char() != ch { return None }
73
74            prev = Some(p.spacing());
75        }
76        Some(self.buf.iter())
77    }
78
79    /// Next jointed puncts
80    pub fn next_puncts(&mut self, puncts: impl AsRef<[u8]>) -> Option<
81        impl Iterator<Item = TokenTree> + '_
82    > {
83        let _ = self.peek_puncts(puncts.as_ref())?;
84        Some(self.buf.drain(..puncts.as_ref().len()))
85    }
86
87    /// Split [`TokenStream`] to `predicate` false and true
88    ///
89    /// Like `"+-,-+".split_puncts(",")` -> `("+-", "-+")`
90    pub fn split_puncts(&mut self, puncts: impl AsRef<[u8]>) -> Option<TokenStream> {
91        let mut left = TokenStream::new();
92        let puncts = puncts.as_ref();
93
94        loop {
95            if self.next_puncts(puncts).is_some() {
96                break Some(left);
97            }
98            let Some(next) = self.next() else {
99                break None;
100            };
101            left.push(next);
102        }
103    }
104}
105impl<I: Iterator<Item = TokenTree>> Iterator for ParseIter<I> {
106    type Item = TokenTree;
107
108    fn next(&mut self) -> Option<Self::Item> {
109        self.buf.pop_front()
110            .or_else(|| self.iter.next())
111    }
112
113    fn count(self) -> usize
114    where Self: Sized,
115    {
116        self.buf.len() + self.iter.count()
117    }
118
119    fn fold<B, F>(self, init: B, f: F) -> B
120    where Self: Sized,
121          F: FnMut(B, Self::Item) -> B,
122    {
123        self.buf.into_iter().chain(self.iter).fold(init, f)
124    }
125
126    fn size_hint(&self) -> (usize, Option<usize>) {
127        let (lo, hi) = self.iter.size_hint();
128        let lo = lo.saturating_add(self.buf.len());
129        let hi = hi.and_then(|hi| hi.checked_add(self.buf.len()));
130        (lo, hi)
131    }
132}
133impl<I: DoubleEndedIterator<Item = TokenTree>> DoubleEndedIterator for ParseIter<I> {
134    fn next_back(&mut self) -> Option<Self::Item> {
135        self.iter.next_back()
136            .or_else(|| self.buf.pop_back())
137    }
138
139    fn rfold<B, F>(self, init: B, f: F) -> B
140    where Self: Sized,
141          F: FnMut(B, Self::Item) -> B,
142    {
143        self.buf.into_iter().chain(self.iter).rfold(init, f)
144    }
145}
146impl<I: ExactSizeIterator<Item = TokenTree>> ExactSizeIterator for ParseIter<I> { }
147impl<I: FusedIterator<Item = TokenTree>> FusedIterator for ParseIter<I> { }