1use proc_macro2::{Span, TokenStream};
2use quote::{quote, quote_spanned, ToTokens, TokenStreamExt};
3use std::collections::HashMap;
4use syn::{
5 parse::{Parse, ParseStream},
6 *,
7};
8
9#[derive(Debug, PartialEq, Eq, Clone, Hash)]
15pub struct IfmtInput {
16 pub source: LitStr,
17 pub segments: Vec<Segment>,
18}
19
20impl IfmtInput {
21 pub fn new(span: Span) -> Self {
22 Self {
23 source: LitStr::new("", span),
24 segments: Vec::new(),
25 }
26 }
27
28 pub fn new_litstr(source: LitStr) -> Result<Self> {
29 let segments = IfmtInput::from_raw(&source.value()).map_err(|e| {
30 let span = source.span();
32 syn::Error::new(span, e)
33 })?;
34 Ok(Self { segments, source })
35 }
36
37 pub fn span(&self) -> Span {
38 self.source.span()
39 }
40
41 pub fn push_raw_str(&mut self, other: String) {
42 self.segments.push(Segment::Literal(other.to_string()))
43 }
44
45 pub fn push_ifmt(&mut self, other: IfmtInput) {
46 self.segments.extend(other.segments);
47 }
48
49 pub fn push_expr(&mut self, expr: Expr) {
50 self.segments.push(Segment::Formatted(FormattedSegment {
51 format_args: String::new(),
52 segment: FormattedSegmentType::Expr(Box::new(expr)),
53 }));
54 }
55
56 pub fn is_static(&self) -> bool {
57 self.segments
58 .iter()
59 .all(|seg| matches!(seg, Segment::Literal(_)))
60 }
61
62 pub fn to_static(&self) -> Option<String> {
63 self.segments
64 .iter()
65 .try_fold(String::new(), |acc, segment| {
66 if let Segment::Literal(seg) = segment {
67 Some(acc + seg)
68 } else {
69 None
70 }
71 })
72 }
73
74 pub fn dynamic_segments(&self) -> Vec<&FormattedSegment> {
75 self.segments
76 .iter()
77 .filter_map(|seg| match seg {
78 Segment::Formatted(seg) => Some(seg),
79 _ => None,
80 })
81 .collect::<Vec<_>>()
82 }
83
84 pub fn dynamic_seg_frequency_map(&self) -> HashMap<&FormattedSegment, usize> {
85 let mut map = HashMap::new();
86 for seg in self.dynamic_segments() {
87 *map.entry(seg).or_insert(0) += 1;
88 }
89 map
90 }
91
92 fn is_simple_expr(&self) -> bool {
93 if !self.segments.is_empty() && self.source.span().byte_range().is_empty() {
95 return false;
96 }
97
98 self.segments.iter().all(|seg| match seg {
99 Segment::Literal(_) => true,
100 Segment::Formatted(FormattedSegment { segment, .. }) => {
101 matches!(segment, FormattedSegmentType::Ident(_))
102 }
103 })
104 }
105
106 fn try_to_string(&self) -> Option<TokenStream> {
110 let mut single_dynamic = None;
111 for segment in &self.segments {
112 match segment {
113 Segment::Literal(literal) => {
114 if !literal.is_empty() {
115 return None;
116 }
117 }
118 Segment::Formatted(FormattedSegment {
119 segment,
120 format_args,
121 }) => {
122 if format_args.is_empty() {
123 match single_dynamic {
124 Some(current_string) => {
125 single_dynamic =
126 Some(quote!(#current_string + &(#segment).to_string()));
127 }
128 None => {
129 single_dynamic = Some(quote!((#segment).to_string()));
130 }
131 }
132 } else {
133 return None;
134 }
135 }
136 }
137 }
138 single_dynamic
139 }
140
141 pub fn to_string_with_quotes(&self) -> String {
143 self.source.to_token_stream().to_string()
144 }
145
146 fn from_raw(input: &str) -> Result<Vec<Segment>> {
148 let mut chars = input.chars().peekable();
149 let mut segments = Vec::new();
150 let mut current_literal = String::new();
151 while let Some(c) = chars.next() {
152 if c == '{' {
153 if let Some(c) = chars.next_if(|c| *c == '{') {
154 current_literal.push(c);
155 continue;
156 }
157 if !current_literal.is_empty() {
158 segments.push(Segment::Literal(current_literal));
159 }
160 current_literal = String::new();
161 let mut current_captured = String::new();
162 while let Some(c) = chars.next() {
163 if c == ':' {
164 if chars.next_if(|c| *c == ':').is_some() {
166 current_captured.push_str("::");
167 continue;
168 }
169 let mut current_format_args = String::new();
170 for c in chars.by_ref() {
171 if c == '}' {
172 segments.push(Segment::Formatted(FormattedSegment {
173 format_args: current_format_args,
174 segment: FormattedSegmentType::parse(¤t_captured)?,
175 }));
176 break;
177 }
178 current_format_args.push(c);
179 }
180 break;
181 }
182 if c == '}' {
183 segments.push(Segment::Formatted(FormattedSegment {
184 format_args: String::new(),
185 segment: FormattedSegmentType::parse(¤t_captured)?,
186 }));
187 break;
188 }
189 current_captured.push(c);
190 }
191 } else {
192 if '}' == c {
193 if let Some(c) = chars.next_if(|c| *c == '}') {
194 current_literal.push(c);
195 continue;
196 } else {
197 return Err(Error::new(
198 Span::call_site(),
199 "unmatched closing '}' in format string",
200 ));
201 }
202 }
203 current_literal.push(c);
204 }
205 }
206
207 if !current_literal.is_empty() {
208 segments.push(Segment::Literal(current_literal));
209 }
210
211 Ok(segments)
212 }
213}
214
215impl ToTokens for IfmtInput {
216 fn to_tokens(&self, tokens: &mut TokenStream) {
217 if let Some(static_str) = self.to_static() {
219 return quote_spanned! { self.span() => #static_str }.to_tokens(tokens);
220 }
221
222 if !cfg!(debug_assertions) {
224 if let Some(single_dynamic) = self.try_to_string() {
225 tokens.extend(single_dynamic);
226 return;
227 }
228 }
229
230 if self.is_simple_expr() {
232 let raw = &self.source;
233 return quote_spanned! { raw.span() => ::std::format!(#raw) }.to_tokens(tokens);
234 }
235
236 let mut format_literal = String::new();
238 let mut expr_counter = 0;
239 for segment in self.segments.iter() {
240 match segment {
241 Segment::Literal(s) => format_literal += &s.replace('{', "{{").replace('}', "}}"),
242 Segment::Formatted(FormattedSegment { format_args, .. }) => {
243 format_literal += "{";
244 format_literal += &expr_counter.to_string();
245 expr_counter += 1;
246 format_literal += ":";
247 format_literal += format_args;
248 format_literal += "}";
249 }
250 }
251 }
252
253 let span = self.span();
254
255 let positional_args = self.segments.iter().filter_map(|seg| {
256 if let Segment::Formatted(FormattedSegment { segment, .. }) = seg {
257 let mut segment = segment.clone();
258 if let FormattedSegmentType::Ident(ident) = &mut segment {
260 ident.set_span(span);
261 }
262 Some(segment)
263 } else {
264 None
265 }
266 });
267
268 quote_spanned! {
269 span =>
270 ::std::format!(
271 #format_literal
272 #(, #positional_args)*
273 )
274 }
275 .to_tokens(tokens)
276 }
277}
278
279#[derive(Debug, PartialEq, Eq, Clone, Hash)]
280pub enum Segment {
281 Literal(String),
282 Formatted(FormattedSegment),
283}
284
285impl Segment {
286 pub fn is_literal(&self) -> bool {
287 matches!(self, Segment::Literal(_))
288 }
289
290 pub fn is_formatted(&self) -> bool {
291 matches!(self, Segment::Formatted(_))
292 }
293}
294
295#[derive(Debug, PartialEq, Eq, Clone, Hash)]
296pub struct FormattedSegment {
297 pub format_args: String,
298 pub segment: FormattedSegmentType,
299}
300
301impl ToTokens for FormattedSegment {
302 fn to_tokens(&self, tokens: &mut TokenStream) {
303 let (fmt, seg) = (&self.format_args, &self.segment);
304 let fmt = format!("{{0:{fmt}}}");
305 tokens.append_all(quote! {
306 format!(#fmt, #seg)
307 });
308 }
309}
310
311#[derive(Debug, PartialEq, Eq, Clone, Hash)]
312pub enum FormattedSegmentType {
313 Expr(Box<Expr>),
314 Ident(Ident),
315}
316
317impl FormattedSegmentType {
318 fn parse(input: &str) -> Result<Self> {
319 if let Ok(ident) = parse_str::<Ident>(input) {
320 if ident == input {
321 return Ok(Self::Ident(ident));
322 }
323 }
324 if let Ok(expr) = parse_str(input) {
325 Ok(Self::Expr(Box::new(expr)))
326 } else {
327 Err(Error::new(
328 Span::call_site(),
329 "Failed to parse formatted segment: Expected Ident or Expression",
330 ))
331 }
332 }
333}
334
335impl ToTokens for FormattedSegmentType {
336 fn to_tokens(&self, tokens: &mut TokenStream) {
337 match self {
338 Self::Expr(expr) => expr.to_tokens(tokens),
339 Self::Ident(ident) => ident.to_tokens(tokens),
340 }
341 }
342}
343
344impl Parse for IfmtInput {
345 fn parse(input: ParseStream) -> Result<Self> {
346 let source: LitStr = input.parse()?;
347 Self::new_litstr(source)
348 }
349}
350
351#[cfg(test)]
352mod tests {
353 use super::*;
354 use prettier_please::PrettyUnparse;
355
356 #[test]
357 fn raw_tokens() {
358 let input = syn::parse2::<IfmtInput>(quote! { r#"hello world"# }).unwrap();
359 println!("{}", input.to_token_stream().pretty_unparse());
360 assert_eq!(input.source.value(), "hello world");
361 assert_eq!(input.to_string_with_quotes(), "r#\"hello world\"#");
362 }
363
364 #[test]
365 fn segments_parse() {
366 let input: IfmtInput = parse_quote! { "blah {abc} {def}" };
367 assert_eq!(
368 input.segments,
369 vec![
370 Segment::Literal("blah ".to_string()),
371 Segment::Formatted(FormattedSegment {
372 format_args: String::new(),
373 segment: FormattedSegmentType::Ident(Ident::new("abc", Span::call_site()))
374 }),
375 Segment::Literal(" ".to_string()),
376 Segment::Formatted(FormattedSegment {
377 format_args: String::new(),
378 segment: FormattedSegmentType::Ident(Ident::new("def", Span::call_site()))
379 }),
380 ]
381 );
382 }
383
384 #[test]
385 fn printing_raw() {
386 let input = syn::parse2::<IfmtInput>(quote! { "hello {world}" }).unwrap();
387 println!("{}", input.to_string_with_quotes());
388
389 let input = syn::parse2::<IfmtInput>(quote! { "hello {world} {world} {world}" }).unwrap();
390 println!("{}", input.to_string_with_quotes());
391
392 let input = syn::parse2::<IfmtInput>(quote! { "hello {world} {world} {world()}" }).unwrap();
393 println!("{}", input.to_string_with_quotes());
394
395 let input =
396 syn::parse2::<IfmtInput>(quote! { r#"hello {world} {world} {world()}"# }).unwrap();
397 println!("{}", input.to_string_with_quotes());
398 assert!(!input.is_static());
399
400 let input = syn::parse2::<IfmtInput>(quote! { r#"hello"# }).unwrap();
401 println!("{}", input.to_string_with_quotes());
402 assert!(input.is_static());
403 }
404
405 #[test]
406 fn to_static() {
407 let input = syn::parse2::<IfmtInput>(quote! { "body {{ background: red; }}" }).unwrap();
408 assert_eq!(
409 input.to_static(),
410 Some("body { background: red; }".to_string())
411 );
412 }
413
414 #[test]
415 fn error_spans() {
416 let input = syn::parse2::<IfmtInput>(quote! { "body {{ background: red; }" }).unwrap_err();
417 assert_eq!(input.span().byte_range(), 0..28);
418 }
419}