Skip to main content

zyn_core/
template.rs

1use proc_macro2::Ident;
2use proc_macro2::Span;
3use proc_macro2::TokenStream;
4use proc_macro2::TokenTree;
5
6use quote::ToTokens;
7use quote::quote;
8
9use syn::Token;
10use syn::parse::Parse;
11use syn::parse::ParseStream;
12
13use crate::Expand;
14use crate::ast::AtNode;
15use crate::ast::GroupNode;
16use crate::ast::Node;
17use crate::ast::TokensNode;
18use crate::ident;
19use crate::types::Input;
20
21/// A parsed template AST. Created by parsing template syntax via `syn::parse2::<Template>(tokens)`.
22pub struct Template {
23    pub nodes: Vec<Node>,
24}
25
26impl Template {
27    pub fn span(&self) -> Span {
28        self.nodes
29            .first()
30            .map(|n| n.span())
31            .unwrap_or_else(Span::call_site)
32    }
33
34    /// Expands the template into a `TokenStream` without an `Input` binding.
35    pub fn to_token_stream(&self) -> TokenStream {
36        let mut idents = ident::Iter::new();
37        let output = idents.next().unwrap();
38        let expanded = self.expand(&output, &mut idents);
39
40        quote! {
41            {
42                let mut #output = ::zyn::proc_macro2::TokenStream::new();
43                #expanded
44                #output
45            }
46        }
47    }
48
49    /// Expands the template with the given `Input` bound as `input` in the generated code.
50    pub fn render(&self, input: &Input) -> TokenStream {
51        let expanded = self.to_token_stream();
52        quote! {
53            {
54                let input: ::zyn::Input = ::zyn::parse!(#input).unwrap();
55                #expanded
56            }
57        }
58    }
59
60    fn flush(pending: &mut TokenStream, nodes: &mut Vec<Node>) {
61        if pending.is_empty() {
62            return;
63        }
64
65        let span = pending
66            .clone()
67            .into_iter()
68            .next()
69            .map(|tt| tt.span())
70            .unwrap_or_else(Span::call_site);
71
72        nodes.push(
73            TokensNode {
74                span,
75                stream: pending.clone(),
76            }
77            .into(),
78        );
79
80        *pending = TokenStream::new();
81    }
82}
83
84impl Expand for Template {
85    fn expand(&self, output: &Ident, idents: &mut ident::Iter) -> TokenStream {
86        let mut result = TokenStream::new();
87
88        for node in &self.nodes {
89            result.extend(node.expand(output, idents));
90        }
91
92        result
93    }
94}
95
96impl Parse for Template {
97    fn parse(input: ParseStream) -> syn::Result<Self> {
98        let mut nodes = Vec::new();
99        let mut pending = TokenStream::new();
100
101        while !input.is_empty() {
102            if input.peek(Token![@]) {
103                Self::flush(&mut pending, &mut nodes);
104                nodes.push(input.parse::<AtNode>()?.into());
105            } else if input.peek(syn::token::Brace) {
106                Self::flush(&mut pending, &mut nodes);
107                nodes.push(Node::parse_brace(input)?);
108            } else if input.peek(syn::token::Paren) || input.peek(syn::token::Bracket) {
109                Self::flush(&mut pending, &mut nodes);
110                nodes.push(input.parse::<GroupNode>()?.into());
111            } else {
112                let tt: TokenTree = input.parse()?;
113                tt.to_tokens(&mut pending);
114            }
115        }
116
117        Self::flush(&mut pending, &mut nodes);
118        Ok(Self { nodes })
119    }
120}