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    /// Returns the source span of the first node, or `Span::call_site()` if empty.
28    pub fn span(&self) -> Span {
29        self.nodes
30            .first()
31            .map(|n| n.span())
32            .unwrap_or_else(Span::call_site)
33    }
34
35    /// Expands the template into a `TokenStream` without an `Input` binding.
36    pub fn to_token_stream(&self) -> TokenStream {
37        let mut idents = ident::Iter::new();
38        let output = idents.next().unwrap();
39        let expanded = self.expand(&output, &mut idents);
40
41        quote! {
42            {
43                let mut #output = ::zyn::proc_macro2::TokenStream::new();
44                #expanded
45                #output
46            }
47        }
48    }
49
50    /// Expands the template with the given `Input` bound as `input` in the generated code.
51    pub fn render(&self, input: &Input) -> TokenStream {
52        let expanded = self.to_token_stream();
53        quote! {
54            {
55                let input: ::zyn::Input = ::zyn::parse!(#input).unwrap();
56                #expanded
57            }
58        }
59    }
60
61    fn flush(pending: &mut TokenStream, nodes: &mut Vec<Node>) {
62        if pending.is_empty() {
63            return;
64        }
65
66        let span = pending
67            .clone()
68            .into_iter()
69            .next()
70            .map(|tt| tt.span())
71            .unwrap_or_else(Span::call_site);
72
73        nodes.push(
74            TokensNode {
75                span,
76                stream: pending.clone(),
77            }
78            .into(),
79        );
80
81        *pending = TokenStream::new();
82    }
83}
84
85impl Expand for Template {
86    fn expand(&self, output: &Ident, idents: &mut ident::Iter) -> TokenStream {
87        let mut result = TokenStream::new();
88
89        for node in &self.nodes {
90            result.extend(node.expand(output, idents));
91        }
92
93        result
94    }
95}
96
97impl Parse for Template {
98    fn parse(input: ParseStream) -> syn::Result<Self> {
99        let mut nodes = Vec::new();
100        let mut pending = TokenStream::new();
101
102        while !input.is_empty() {
103            if input.peek(Token![@]) {
104                Self::flush(&mut pending, &mut nodes);
105                nodes.push(input.parse::<AtNode>()?.into());
106            } else if input.peek(syn::token::Brace) {
107                Self::flush(&mut pending, &mut nodes);
108                nodes.push(Node::parse_brace(input)?);
109            } else if input.peek(syn::token::Paren) || input.peek(syn::token::Bracket) {
110                Self::flush(&mut pending, &mut nodes);
111                nodes.push(input.parse::<GroupNode>()?.into());
112            } else {
113                let tt: TokenTree = input.parse()?;
114                tt.to_tokens(&mut pending);
115            }
116        }
117
118        Self::flush(&mut pending, &mut nodes);
119        Ok(Self { nodes })
120    }
121}