Skip to main content

zyn_core/ast/at/
for_node.rs

1use proc_macro2::Ident;
2use proc_macro2::Span;
3use proc_macro2::TokenStream;
4
5use quote::quote;
6
7use syn::ext::IdentExt;
8use syn::parse::Parse;
9use syn::parse::ParseStream;
10
11use crate::template::Template;
12
13use crate::Expand;
14
15pub struct ForNode {
16    pub span: Span,
17    pub binding: syn::Ident,
18    pub iter: TokenStream,
19    pub body: Box<Template>,
20}
21
22impl ForNode {
23    pub fn span(&self) -> Span {
24        self.span
25    }
26}
27
28impl Parse for ForNode {
29    fn parse(input: ParseStream) -> syn::Result<Self> {
30        let params;
31        syn::parenthesized!(params in input);
32
33        let fork = params.fork();
34        let is_classic = if fork.call(syn::Ident::parse_any).is_ok() {
35            let Ok(kw) = fork.call(syn::Ident::parse_any) else {
36                return Self::parse_classic(&params, input);
37            };
38            kw != "in"
39        } else {
40            true
41        };
42
43        if is_classic {
44            return Self::parse_classic(&params, input);
45        }
46
47        let binding: syn::Ident = params.call(syn::Ident::parse_any)?;
48        let _in_kw: syn::Ident = params.call(syn::Ident::parse_any)?;
49        let iter: TokenStream = params.parse()?;
50
51        let body_content;
52        syn::braced!(body_content in input);
53        let body = body_content.parse::<Template>()?;
54
55        Ok(Self {
56            span: Span::call_site(),
57            binding,
58            iter,
59            body: Box::new(body),
60        })
61    }
62}
63
64impl ForNode {
65    fn parse_classic(params: ParseStream, input: ParseStream) -> syn::Result<Self> {
66        let count: TokenStream = params.parse()?;
67        let binding = syn::Ident::new("_", Span::call_site());
68        let iter: TokenStream = quote! { 0..#count };
69
70        let body_content;
71        syn::braced!(body_content in input);
72        let body = body_content.parse::<Template>()?;
73
74        Ok(Self {
75            span: Span::call_site(),
76            binding,
77            iter,
78            body: Box::new(body),
79        })
80    }
81}
82
83impl Expand for ForNode {
84    fn expand(&self, output: &Ident, idents: &mut crate::ident::Iter) -> TokenStream {
85        let binding = &self.binding;
86        let iter = &self.iter;
87        let body_expanded = self.body.expand(output, idents);
88
89        quote! {
90            for #binding in #iter {
91                #body_expanded
92            }
93        }
94    }
95}