1#![forbid(unsafe_code)]
6
7use std::borrow::Cow;
8
9use ast::Markup;
10use itertools::Itertools;
11use proc_macro_error2::proc_macro_error;
12use proc_macro2::TokenStream;
13use quote::{ToTokens, TokenStreamExt, quote};
14use syn::{
15 Expr, Macro, Pat, Stmt,
16 parse::{ParseStream, Parser},
17};
18
19mod ast;
20
21#[proc_macro]
22#[proc_macro_error]
23pub fn markup(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
24 expand(input.into()).into()
25}
26
27#[proc_macro]
28#[proc_macro_error]
29pub fn markup_part(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
30 expand_part(input.into()).into()
31}
32
33fn expand(input: TokenStream) -> TokenStream {
34 let reserve = input.to_string().len();
35 match Parser::parse2(|input: ParseStream| input.parse::<Markup>(), input) {
36 Ok(markup) => {
37 let mut parts = Parts::new();
38 markup.append(&mut parts);
39 let mut tokens = TokenStream::new();
40 tokens.append_all(parts.0);
41 quote! {{
42 extern crate alloc;
43 extern crate collage;
44 &|__collage_buffer: &mut alloc::string::String| {
45 __collage_buffer.reserve(#reserve);
46 #tokens
47 }
48 }}
49 }
50 Err(err) => err.into_compile_error(),
51 }
52}
53
54fn expand_part(input: TokenStream) -> TokenStream {
55 match Parser::parse2(|input: ParseStream| input.parse::<Markup>(), input) {
56 Ok(markup) => {
57 let mut parts = Parts::new();
58 markup.append(&mut parts);
59 let mut tokens = TokenStream::new();
60 tokens.append_all(parts.0);
61 tokens
62 }
63 Err(err) => err.into_compile_error(),
64 }
65}
66
67fn is_void<T: PartialEq<str>>(tag: &T) -> bool {
68 [
69 "area", "base", "br", "col", "embed", "hr", "img", "input", "link", "meta", "source", "track", "wbr",
70 ]
71 .iter()
72 .any(|void| tag == *void)
73}
74
75trait PartBuilder {
76 fn append(&self, parts: &mut Parts);
77}
78
79impl<T: PartBuilder> PartBuilder for Option<T> {
80 fn append(&self, parts: &mut Parts) {
81 if let Some(t) = self {
82 t.append(parts);
83 }
84 }
85}
86
87struct Parts(Vec<Part>);
88
89impl Parts {
90 fn new() -> Self {
91 Self(Vec::new())
92 }
93
94 fn push(&mut self, part: Part) {
95 if let Some(prev) = self.0.pop() {
96 match (prev, part) {
97 (Part::Static(prev), Part::Static(part)) => {
98 self.0.push(format!("{prev}{part}").into());
99 }
100 (prev, part) => {
101 self.0.push(prev);
102 self.0.push(part);
103 }
104 }
105 } else {
106 self.0.push(part);
107 }
108 }
109}
110
111enum Part {
112 Static(String),
113 Dynamic((Expr, bool)),
114}
115
116impl From<String> for Part {
117 fn from(value: String) -> Self {
118 Self::Static(value)
119 }
120}
121
122impl From<&str> for Part {
123 fn from(value: &str) -> Self {
124 Self::Static(value.into())
125 }
126}
127
128impl From<Cow<'_, str>> for Part {
129 fn from(value: Cow<'_, str>) -> Self {
130 match value {
131 Cow::Borrowed(value) => Self::Static(value.into()),
132 Cow::Owned(value) => Self::Static(value),
133 }
134 }
135}
136
137impl From<Expr> for Part {
138 fn from(value: Expr) -> Self {
139 Self::Dynamic((value, false))
140 }
141}
142
143impl ToTokens for Part {
144 fn to_tokens(&self, tokens: &mut TokenStream) {
145 match self {
146 Part::Static(s) => {
147 quote! { __collage_buffer.push_str(#s); }.to_tokens(tokens);
148 }
149 Part::Dynamic((Expr::If(expr), attr)) if expr.else_branch.is_none() => {
150 tokens.append_all(&expr.attrs);
151 expr.if_token.to_tokens(tokens);
152 expr.cond.to_tokens(tokens);
153 expr.then_branch.brace_token.surround(tokens, |tokens| {
154 let stmts = stmts_replace_mac(&expr.then_branch.stmts);
155 if *attr {
156 quote! { __collage_buffer.push(' '); }.to_tokens(tokens);
157 }
158 quote! { collage::Render::render_to({ #(#stmts)* }, __collage_buffer); }.to_tokens(tokens);
159 });
160 }
161 Part::Dynamic((Expr::Match(expr), _)) => {
162 expr.match_token.to_tokens(tokens);
163 expr.expr.to_tokens(tokens);
164 expr.brace_token.surround(tokens, |tokens| {
165 for arm in &expr.arms {
166 tokens.append_all(&arm.attrs);
167 arm.pat.to_tokens(tokens);
168 if let Some((if_token, expr)) = &arm.guard {
169 if_token.to_tokens(tokens);
170 expr.to_tokens(tokens);
171 }
172 arm.fat_arrow_token.to_tokens(tokens);
173 match &*arm.body {
174 Expr::Macro(expr) if is_markup(&expr.mac) => {
175 tokens.append_all(&expr.attrs);
176 let mac = &expr.mac.tokens;
177 quote! {{ collage::markup_part! { #mac } }}.to_tokens(tokens);
178 }
179 expr => expr.to_tokens(tokens),
180 }
181 arm.comma.to_tokens(tokens);
182 }
183 if !expr.arms.iter().any(|arm| matches!(arm.pat, Pat::Wild(_))) {
184 quote! { _ => {} }.to_tokens(tokens);
185 }
186 });
187 }
188 Part::Dynamic((Expr::ForLoop(expr), _)) => {
189 tokens.append_all(&expr.attrs);
190 expr.for_token.to_tokens(tokens);
191 expr.pat.to_tokens(tokens);
192 expr.in_token.to_tokens(tokens);
193 expr.expr.to_tokens(tokens);
194 let stmts = stmts_replace_mac(&expr.body.stmts);
195 quote! {{ #(#stmts)* }}.to_tokens(tokens);
196 }
197 Part::Dynamic((Expr::Macro(expr), _)) if is_markup(&expr.mac) => {
198 tokens.append_all(&expr.attrs);
199 let mac = &expr.mac.tokens;
200 quote! { collage::markup_part! { #mac } }.to_tokens(tokens);
201 }
202 Part::Dynamic((Expr::Macro(expr), _)) if is_markup_part(&expr.mac) => {
203 tokens.append_all(&expr.attrs);
204 expr.mac.to_tokens(tokens);
205 }
206 Part::Dynamic((expr, attr)) => {
207 if *attr {
208 quote! { __collage_buffer.push(' '); }.to_tokens(tokens);
209 }
210 quote! { collage::Render::render_to(&#expr, __collage_buffer); }.to_tokens(tokens);
211 }
212 }
213 }
214}
215
216fn stmts_replace_mac(stmts: &[Stmt]) -> Vec<TokenStream> {
217 stmts
218 .iter()
219 .map(|stmt| match stmt {
220 Stmt::Macro(stmt) if is_markup(&stmt.mac) => {
221 let mac = &stmt.mac.tokens;
222 quote! { collage::markup_part! { #mac } }
223 }
224 stmt => quote! { #stmt },
225 })
226 .collect_vec()
227}
228
229fn is_markup(mac: &Macro) -> bool {
230 mac.path.is_ident("markup")
231 || mac
232 .path
233 .segments
234 .iter()
235 .rev()
236 .next_tuple()
237 .is_some_and(|(b, a)| a.ident == "collage" && b.ident == "markup")
238}
239
240fn is_markup_part(mac: &Macro) -> bool {
241 mac.path.is_ident("markup_part")
242 || mac
243 .path
244 .segments
245 .iter()
246 .rev()
247 .next_tuple()
248 .is_some_and(|(b, a)| a.ident == "collage" && b.ident == "markup_part")
249}