cgp_component_macro_lib/
for_each_replace.rs

1use alloc::vec::Vec;
2
3use proc_macro2::{Group, TokenStream, TokenTree};
4use quote::{quote, ToTokens};
5use syn::__private::parse_brackets;
6use syn::parse::discouraged::Speculative;
7use syn::parse::{Parse, ParseStream};
8use syn::punctuated::Punctuated;
9use syn::token::{Comma, Or};
10use syn::{braced, Ident, Type};
11
12use crate::delegate_components::ast::ComponentAst;
13
14pub struct ReplaceSpecs {
15    pub target_ident: Ident,
16    pub replacements: Vec<TokenStream>,
17    pub body: TokenStream,
18}
19
20impl Parse for ReplaceSpecs {
21    fn parse(input: ParseStream) -> syn::Result<Self> {
22        let raw_replacements: Vec<ComponentAst> = {
23            let content = parse_brackets(input)?.content;
24            let types = <Punctuated<ComponentAst, Comma>>::parse_terminated(&content)?;
25            types.into_iter().collect()
26        };
27
28        Comma::parse(input)?;
29
30        let exclude: Vec<Type> = {
31            let fork = input.fork();
32
33            match parse_brackets(&fork) {
34                Ok(bracket) => {
35                    let types = <Punctuated<Type, Comma>>::parse_terminated(&bracket.content)?;
36
37                    input.advance_to(&fork);
38                    Comma::parse(input)?;
39
40                    types.into_iter().collect()
41                }
42                _ => Vec::new(),
43            }
44        };
45
46        Or::parse(input)?;
47
48        let target_ident = Ident::parse(input)?;
49
50        Or::parse(input)?;
51
52        let body = {
53            let content;
54            braced!(content in input);
55            TokenStream::parse(&content)?
56        };
57
58        let replacements = raw_replacements
59            .into_iter()
60            .filter(|replacement| {
61                !exclude
62                    .iter()
63                    .any(|exclude| exclude == &replacement.component_type)
64            })
65            .map(|ast| ast.to_token_stream())
66            .collect();
67
68        Ok(ReplaceSpecs {
69            target_ident,
70            replacements,
71            body,
72        })
73    }
74}
75
76pub fn handle_for_each_replace(tokens: TokenStream) -> syn::Result<TokenStream> {
77    let specs: ReplaceSpecs = syn::parse2(tokens)?;
78
79    Ok(for_each_replace(
80        &specs.target_ident,
81        &specs.replacements,
82        &specs.body,
83    ))
84}
85
86pub fn handle_replace(tokens: TokenStream) -> syn::Result<TokenStream> {
87    let specs: ReplaceSpecs = syn::parse2(tokens)?;
88
89    let items: Punctuated<TokenStream, Comma> = specs.replacements.into_iter().collect();
90
91    let tokens = quote! { [ #items ] };
92
93    Ok(replace_stream(&specs.target_ident, &tokens, specs.body))
94}
95
96pub fn for_each_replace(
97    target_ident: &Ident,
98    replacements: &[TokenStream],
99    body: &TokenStream,
100) -> TokenStream {
101    replacements
102        .iter()
103        .map(|replacement| replace_stream(target_ident, replacement, body.clone()))
104        .collect()
105}
106
107pub fn replace_stream(
108    target_ident: &Ident,
109    replacement: &TokenStream,
110    body: TokenStream,
111) -> TokenStream {
112    body.into_iter()
113        .map(|tree| replace_tree(target_ident, replacement, tree))
114        .collect()
115}
116
117pub fn replace_tree(
118    target_ident: &Ident,
119    replacement: &TokenStream,
120    body: TokenTree,
121) -> TokenStream {
122    match body {
123        TokenTree::Group(group) => TokenTree::Group(Group::new(
124            group.delimiter(),
125            replace_stream(target_ident, replacement, group.stream()),
126        ))
127        .into(),
128        TokenTree::Ident(ident) => {
129            if &ident == target_ident {
130                replacement.to_token_stream()
131            } else {
132                TokenTree::Ident(ident).into()
133            }
134        }
135        tokens => tokens.into(),
136    }
137}