tokel_std/
iter.rs

1//! Iteration-related Tokel [`Transformer`]s.
2
3use std::{iter, ops::Range};
4
5use proc_macro2::{Literal, TokenStream, TokenTree};
6use syn::{
7    LitInt, Token,
8    parse::{Nothing, Parse, ParseStream},
9};
10use tokel_engine::prelude::{Registry, Transformer};
11
12/// A transformer that reverses the sequence of the token trees in the stream.
13///
14/// # Example
15/// `[< a b c >]:reverse` -> `c b a`
16#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
17pub struct Reverse;
18
19impl Transformer for Reverse {
20    fn transform(
21        &mut self,
22        input: TokenStream,
23        argument: TokenStream,
24    ) -> Result<TokenStream, syn::Error> {
25        let _: Nothing = syn::parse2(argument)?;
26
27        let mut tokens: Vec<_> = input.into_iter().collect();
28
29        tokens.reverse();
30
31        Ok(tokens.into_iter().collect())
32    }
33}
34
35/// Inserts the argument between every token tree in the input.
36#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
37pub struct Intersperse;
38
39impl Transformer for Intersperse {
40    fn transform(
41        &mut self,
42        input: TokenStream,
43        argument: TokenStream,
44    ) -> Result<TokenStream, syn::Error> {
45        let mut output = TokenStream::new();
46        let mut iter = input.into_iter().peekable();
47
48        while let Some(tree) = iter.next() {
49            output.extend(iter::once(tree));
50
51            if iter.peek().is_some() {
52                output.extend(argument.clone());
53            }
54        }
55
56        Ok(output)
57    }
58}
59
60/// Prepends the argument to the start of the stream.
61#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
62pub struct PushLeft;
63
64impl Transformer for PushLeft {
65    fn transform(
66        &mut self,
67        input: TokenStream,
68        argument: TokenStream,
69    ) -> Result<TokenStream, syn::Error> {
70        let mut output = argument;
71
72        output.extend(input);
73
74        Ok(output)
75    }
76}
77
78/// Removes the last token tree from the stream (useful for trailing commas).
79#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
80pub struct PopRight;
81
82impl Transformer for PopRight {
83    fn transform(
84        &mut self,
85        input: TokenStream,
86        argument: TokenStream,
87    ) -> Result<TokenStream, syn::Error> {
88        let _: Nothing = syn::parse2(argument)?;
89
90        let mut tokens: Vec<_> = input.into_iter().collect();
91
92        tokens.pop(); // Discard the last token
93
94        Ok(tokens.into_iter().collect())
95    }
96}
97
98/// Appends the argument to the end of the stream.
99#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
100pub struct PushRight;
101
102impl Transformer for PushRight {
103    fn transform(
104        &mut self,
105        mut input: TokenStream,
106        argument: TokenStream,
107    ) -> Result<TokenStream, syn::Error> {
108        input.extend(argument);
109        Ok(input)
110    }
111}
112
113/// Removes the first token tree from the stream.
114#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
115pub struct PopLeft;
116
117impl Transformer for PopLeft {
118    fn transform(
119        &mut self,
120        input: TokenStream,
121        argument: TokenStream,
122    ) -> Result<TokenStream, syn::Error> {
123        let _: Nothing = syn::parse2(argument)?;
124
125        let mut iter = input.into_iter();
126        iter.next(); // Discard the first token
127
128        Ok(iter.collect())
129    }
130}
131
132/// Keeps only the first `N` tokens from the stream.
133///
134/// The argument must be a valid integer literal, e.g., `[[ 3 ]]`.
135#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
136pub struct Take;
137
138impl Transformer for Take {
139    fn transform(
140        &mut self,
141        input: TokenStream,
142        argument: TokenStream,
143    ) -> Result<TokenStream, syn::Error> {
144        let lit: syn::LitInt = syn::parse2(argument)?;
145        let n: usize = lit.base10_parse()?;
146
147        Ok(input.into_iter().take(n).collect())
148    }
149}
150
151/// Discards the first `N` tokens from the stream.
152///
153/// The argument must be a valid integer literal, e.g., `[[ 2 ]]`.
154#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
155pub struct Skip;
156
157impl Transformer for Skip {
158    fn transform(
159        &mut self,
160        input: TokenStream,
161        argument: TokenStream,
162    ) -> Result<TokenStream, syn::Error> {
163        let lit: syn::LitInt = syn::parse2(argument)?;
164
165        Ok(input.into_iter().skip(lit.base10_parse()?).collect())
166    }
167}
168
169/// Duplicates the input stream `N` times.
170///
171/// The argument must be a valid integer literal, e.g., `[[ 5 ]]`.
172#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
173pub struct Repeat;
174
175impl Transformer for Repeat {
176    fn transform(
177        &mut self,
178        input: TokenStream,
179        argument: TokenStream,
180    ) -> Result<TokenStream, syn::Error> {
181        let lit: syn::LitInt = syn::parse2(argument)?;
182
183        let n: usize = lit.base10_parse()?;
184
185        let mut output = TokenStream::new();
186
187        for _ in 0..n {
188            output.extend(input.clone());
189        }
190
191        Ok(output)
192    }
193}
194
195/// Emits the number of token trees present in the stream.
196///
197/// E.g., `[< a b c >]:measure` -> `3`
198#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
199pub struct Count;
200
201impl Transformer for Count {
202    fn transform(
203        &mut self,
204        input: TokenStream,
205        argument: TokenStream,
206    ) -> Result<TokenStream, syn::Error> {
207        let _: Nothing = syn::parse2(argument)?;
208
209        let target_count = input.into_iter().count();
210
211        let target_literal = Literal::usize_unsuffixed(target_count);
212
213        Ok(quote::quote!(#target_literal))
214    }
215}
216
217/// Generates a sequence of integers based on a provided range.
218///
219/// The input stream is ignored. The argument must be a valid Rust range literal.
220/// Supported formats: `[[ 0..5 ]]` or `[[ 1..=10 ]]`.
221///
222/// # Example
223/// `[< >]:sequence[[ 1..=3 ]]` -> `1 2 3`
224#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
225pub struct Sequence;
226
227impl Transformer for Sequence {
228    fn transform(&mut self, input: TokenStream, argument: TokenStream) -> syn::Result<TokenStream> {
229        struct Argument(Range<i32>);
230
231        impl Parse for Argument {
232            fn parse(input: ParseStream) -> syn::Result<Self> {
233                let start: LitInt = input.parse()?;
234
235                let _: Token![..] = input.parse()?;
236
237                let end: LitInt = input.parse()?;
238
239                Ok(Self(start.base10_parse()?..end.base10_parse()?))
240            }
241        }
242
243        let _: Nothing = syn::parse2(input)?;
244
245        let Argument(target_range) = syn::parse2(argument)?;
246
247        let mut output = TokenStream::new();
248
249        for target_value in target_range {
250            output.extend(iter::once(TokenTree::Literal(Literal::i32_unsuffixed(
251                target_value,
252            ))));
253        }
254
255        Ok(output)
256    }
257}
258
259/// Inserts all `iter`-related [`Transformer`]s into the specified [`Registry`].
260///
261/// # Errors
262///
263/// This will fail if at least one standard `iter`-related [`Transformer`] is already present by-name in the [`Registry`].
264///
265/// On failure, there is no guarantee that other non-colliding transformers have not been registered.
266#[inline]
267pub fn register(registry: &mut Registry) -> Result<(), Box<dyn Transformer>> {
268    registry
269        .try_insert("reverse", Reverse)
270        .map_err(Box::new)
271        .map_err(|t| t as Box<dyn Transformer>)?;
272
273    registry
274        .try_insert("intersperse", Intersperse)
275        .map_err(Box::new)
276        .map_err(|t| t as Box<dyn Transformer>)?;
277
278    registry
279        .try_insert("push_left", PushLeft)
280        .map_err(Box::new)
281        .map_err(|t| t as Box<dyn Transformer>)?;
282
283    registry
284        .try_insert("push_right", PushRight)
285        .map_err(Box::new)
286        .map_err(|t| t as Box<dyn Transformer>)?;
287
288    registry
289        .try_insert("pop_left", PopLeft)
290        .map_err(Box::new)
291        .map_err(|t| t as Box<dyn Transformer>)?;
292
293    registry
294        .try_insert("pop_right", PopRight)
295        .map_err(Box::new)
296        .map_err(|t| t as Box<dyn Transformer>)?;
297
298    registry
299        .try_insert("take", Take)
300        .map_err(Box::new)
301        .map_err(|t| t as Box<dyn Transformer>)?;
302
303    registry
304        .try_insert("skip", Skip)
305        .map_err(Box::new)
306        .map_err(|t| t as Box<dyn Transformer>)?;
307
308    registry
309        .try_insert("repeat", Repeat)
310        .map_err(Box::new)
311        .map_err(|t| t as Box<dyn Transformer>)?;
312
313    registry
314        .try_insert("count", Count)
315        .map_err(Box::new)
316        .map_err(|t| t as Box<dyn Transformer>)?;
317
318    registry
319        .try_insert("sequence", Sequence)
320        .map_err(Box::new)
321        .map_err(|t| t as Box<dyn Transformer>)?;
322
323    Ok(())
324}
tokel_std/iter.rs

tokel_std/
iter.rs
