Skip to main content

unsynn/
token_iter.rs

1//! Iterator type for parsing token streams.
2
3#[cfg(feature = "proc_macro2")]
4use proc_macro2::{TokenStream, TokenTree};
5
6#[cfg(not(feature = "proc_macro2"))]
7use proc_macro::{TokenStream, TokenTree};
8
9use crate::EndOfStream;
10use crate::Parser;
11
12/// The base iterator
13#[derive(Clone)]
14#[cfg_attr(feature = "proc_macro2", derive(Debug))]
15struct TokenIterBase {
16    /// The actual token iterator
17    iter: <TokenStream as IntoIterator>::IntoIter,
18    /// Counter for this level
19    counter: usize,
20
21    #[cfg(feature = "extra_asserts")]
22    /// Tracks if a nested iterator has been committed at the current position
23    /// Set to true on commit, cleared when parent advances via `.next()`
24    committed: bool,
25}
26
27impl TokenIterBase {
28    fn new(
29        iter: impl IntoIterator<IntoIter = <TokenStream as IntoIterator>::IntoIter>,
30        start_at: usize,
31    ) -> Self {
32        Self {
33            iter: iter.into_iter(),
34            counter: start_at,
35            #[cfg(feature = "extra_asserts")]
36            committed: false,
37        }
38    }
39}
40
41/// Iterator type for parsing token streams.
42///
43/// # Examples
44///
45/// ```rust
46/// # use unsynn::*;
47/// // Simple parsing
48/// let tokens = "foo bar".to_token_iter();
49/// ```
50#[derive(Clone)]
51#[cfg_attr(feature = "proc_macro2", derive(Debug))]
52pub struct TokenIter {
53    base: TokenIterBase,
54}
55
56impl TokenIter {
57    /// Create a new owned `TokenIter`
58    pub fn new(
59        stream: impl IntoIterator<IntoIter = <TokenStream as IntoIterator>::IntoIter>,
60    ) -> Self {
61        TokenIter {
62            base: TokenIterBase::new(stream, 0),
63        }
64    }
65}
66
67impl TokenIter {
68    /// Get the current counter value
69    #[must_use]
70    pub const fn counter(&self) -> usize {
71        self.base.counter
72    }
73
74    /// Adjust the counter by an offset
75    pub(crate) fn add(&mut self, add: usize) {
76        self.base.counter += add;
77    }
78
79    /// Scoped parsing of a nested group.
80    ///
81    /// Creates a new iterator with the stream to the inner of the group,
82    /// parses type `T`, and ensures all tokens are consumed.
83    pub(crate) fn parse_group<T>(
84        &mut self,
85        stream: impl IntoIterator<IntoIter = <TokenStream as IntoIterator>::IntoIter>,
86    ) -> crate::Result<T>
87    where
88        T: crate::Parser,
89    {
90        #[cfg(feature = "extra_asserts")]
91        assert!(!self.base.committed);
92        let mut tokens = TokenIter {
93            base: TokenIterBase::new(stream, self.counter()),
94        };
95        let parsed = T::parser(&mut tokens)?;
96        // Must parse exhaustively, otherwise we would miscount tokens
97        EndOfStream::parser(&mut tokens)?;
98        self.base.counter = tokens.counter();
99        #[cfg(feature = "extra_asserts")]
100        {
101            self.base.committed = true;
102        };
103        Ok(parsed)
104    }
105}
106
107// Iterator implementation
108impl Iterator for TokenIter {
109    type Item = TokenTree;
110
111    fn next(&mut self) -> Option<Self::Item> {
112        let item = self.base.iter.next()?;
113        self.base.counter += 1;
114        #[cfg(feature = "extra_asserts")]
115        {
116            self.base.committed = false;
117        };
118        Some(item)
119    }
120}
121
122/// Extension trait to convert iterators into [`TokenIter`].
123pub trait IntoTokenIter: Sized {
124    /// Convert this iterator into a [`TokenIter`].
125    fn into_token_iter(self) -> TokenIter;
126}
127
128impl IntoTokenIter for <TokenStream as IntoIterator>::IntoIter {
129    fn into_token_iter(self) -> TokenIter {
130        TokenIter {
131            base: TokenIterBase::new(self, 0),
132        }
133    }
134}
135
136/// Extension trait to convert `TokenStreams` into [`TokenIter`].
137pub trait ToTokenIter: Sized {
138    /// Convert this into a [`TokenIter`].
139    fn to_token_iter(self) -> TokenIter;
140}
141
142impl ToTokenIter for TokenStream {
143    fn to_token_iter(self) -> TokenIter {
144        TokenIter::new(self)
145    }
146}