1use std::iter;
4
5use proc_macro2::{Delimiter, Group, TokenStream, TokenTree};
6use syn::parse::{Nothing, Parse, ParseStream};
7use tokel_engine::prelude::{Registry, Transformer};
8
9#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
11pub struct Flatten;
12
13impl Flatten {
14 fn flatten_recursive(stream: TokenStream) -> TokenStream {
15 let mut output = TokenStream::new();
16 for tree in stream {
17 match tree {
18 TokenTree::Group(g) => output.extend(Self::flatten_recursive(g.stream())),
19 other => output.extend(iter::once(other)),
20 }
21 }
22 output
23 }
24}
25
26impl Transformer for Flatten {
27 fn transform(
28 &mut self,
29 input: TokenStream,
30 argument: TokenStream,
31 ) -> Result<TokenStream, syn::Error> {
32 let _: Nothing = syn::parse2(argument)?;
33
34 Ok(Self::flatten_recursive(input))
35 }
36}
37
38#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
40pub struct Encapsulate;
41
42impl Transformer for Encapsulate {
43 fn transform(
44 &mut self,
45 input: TokenStream,
46 argument: TokenStream,
47 ) -> Result<TokenStream, syn::Error> {
48 struct Argument(Delimiter);
49
50 impl Parse for Argument {
51 fn parse(input: ParseStream) -> syn::Result<Self> {
52 let group = input.parse::<Group>()?;
53
54 let _: Nothing = syn::parse2(group.stream())?;
55
56 let _: Nothing = input.parse()?;
57
58 Ok(Self(group.delimiter()))
59 }
60 }
61
62 let Argument(delimiter) = syn::parse2(argument)?;
63
64 let group = Group::new(delimiter, input);
65
66 Ok(iter::once(TokenTree::Group(group)).collect())
67 }
68}
69
70#[inline]
78pub fn register(registry: &mut Registry) -> Result<(), Box<dyn Transformer>> {
79 registry
80 .try_insert("flatten", Flatten)
81 .map_err(Box::new)
82 .map_err(|target_value| target_value as Box<dyn Transformer>)?;
83
84 registry
85 .try_insert("encapsulate", Encapsulate)
86 .map_err(Box::new)
87 .map_err(|target_value| target_value as Box<dyn Transformer>)?;
88
89 Ok(())
90}