rune/macros/
format_args.rs

1use core::str;
2
3use crate as rune;
4use crate::alloc::prelude::*;
5use crate::alloc::{self, BTreeMap, BTreeSet, Box, HashMap, String, Vec};
6use crate::ast::{self, Span, Spanned};
7use crate::compile::{self, WithSpan};
8use crate::macros::{quote, MacroContext, Quote, ToTokens, TokenStream};
9use crate::parse::{Parse, Parser, Peek, Peeker};
10use crate::runtime::{format, Inline};
11
12/// A format specification: A format string followed by arguments to be
13/// formatted in accordance with that string.
14///
15/// This type can only be parsed inside of a macro context since it performs
16/// constant evaluation.
17pub struct FormatArgs {
18    /// Format argument.
19    format: ast::Expr,
20    /// Format arguments.
21    args: Vec<FormatArg>,
22}
23
24impl FormatArgs {
25    /// Expand the format specification.
26    pub fn expand(&self, cx: &mut MacroContext<'_, '_, '_>) -> compile::Result<Quote<'_>> {
27        let format = cx.eval(&self.format)?;
28
29        let mut pos = Vec::new();
30        let mut named = HashMap::<Box<str>, _>::new();
31
32        for a in &self.args {
33            match a {
34                FormatArg::Positional(expr) => {
35                    if !named.is_empty() {
36                        return Err(compile::Error::msg(
37                            expr.span(),
38                            "unnamed positional arguments must come before named ones",
39                        ));
40                    }
41
42                    pos.try_push(expr)?;
43                }
44                FormatArg::Named(n) => {
45                    let name = cx.resolve(n.key)?;
46                    named.try_insert(name.try_into()?, n)?;
47                }
48            }
49        }
50
51        let format = format.downcast::<String>().with_span(&self.format)?;
52
53        let mut unused_pos = (0..pos.len()).try_collect::<BTreeSet<_>>()?;
54        let mut unused_named = named
55            .iter()
56            .map(|(key, n)| Ok::<_, alloc::Error>((key.try_clone()?, n.span())))
57            .try_collect::<alloc::Result<BTreeMap<_, _>>>()??;
58
59        let result = expand_format_spec(
60            cx,
61            self.format.span(),
62            &format,
63            &pos,
64            &mut unused_pos,
65            &named,
66            &mut unused_named,
67        );
68
69        let expanded = match result {
70            Ok(expanded) => expanded,
71            Err(message) => return Err(compile::Error::msg(self.format.span(), message)),
72        };
73
74        if let Some(expr) = unused_pos.into_iter().flat_map(|n| pos.get(n)).next() {
75            return Err(compile::Error::msg(
76                expr.span(),
77                "unused positional argument",
78            ));
79        }
80
81        if let Some((key, span)) = unused_named.into_iter().next() {
82            return Err(compile::Error::msg(
83                span,
84                format!("unused named argument `{key}`"),
85            ));
86        }
87
88        Ok(expanded)
89    }
90}
91
92impl Parse for FormatArgs {
93    /// Parse format arguments inside of a macro.
94    fn parse(p: &mut Parser<'_>) -> compile::Result<Self> {
95        if p.is_eof()? {
96            return Err(compile::Error::msg(
97                p.last_span(),
98                "expected format specifier",
99            ));
100        }
101
102        let format = p.parse::<ast::Expr>()?;
103
104        let mut args = Vec::new();
105
106        while p.parse::<Option<T![,]>>()?.is_some() {
107            if p.is_eof()? {
108                break;
109            }
110
111            args.try_push(p.parse()?)?;
112        }
113
114        Ok(Self { format, args })
115    }
116}
117
118impl Peek for FormatArgs {
119    fn peek(p: &mut Peeker<'_>) -> bool {
120        !p.is_eof()
121    }
122}
123
124/// A named format argument.
125#[derive(Debug, TryClone, Parse, Spanned)]
126pub struct NamedFormatArg {
127    /// The key of the named argument.
128    pub key: ast::Ident,
129    /// The `=` token.
130    pub eq_token: T![=],
131    /// The value expression.
132    pub expr: ast::Expr,
133}
134
135/// A single format argument.
136#[derive(Debug, TryClone)]
137pub enum FormatArg {
138    /// A positional argument.
139    Positional(ast::Expr),
140    /// A named argument.
141    Named(NamedFormatArg),
142}
143
144impl Parse for FormatArg {
145    fn parse(p: &mut Parser<'_>) -> compile::Result<Self> {
146        Ok(if let (K![ident], K![=]) = (p.nth(0)?, p.nth(1)?) {
147            FormatArg::Named(p.parse()?)
148        } else {
149            FormatArg::Positional(p.parse()?)
150        })
151    }
152}
153
154fn expand_format_spec<'a>(
155    cx: &mut MacroContext<'_, '_, '_>,
156    span: Span,
157    input: &str,
158    pos: &[&'a ast::Expr],
159    unused_pos: &mut BTreeSet<usize>,
160    named: &HashMap<Box<str>, &'a NamedFormatArg>,
161    unused_named: &mut BTreeMap<Box<str>, Span>,
162) -> compile::Result<Quote<'a>> {
163    let mut iter = Iter::new(input);
164
165    let mut name = String::new();
166    let mut width = String::new();
167    let mut precision = String::new();
168
169    let mut buf = String::new();
170    let mut components = Vec::new();
171    let mut count = 0;
172    let mut start = Some(0);
173
174    while let Some((at, a, b)) = iter.next() {
175        match (a, b) {
176            ('}', '}') => {
177                if let Some(start) = start.take() {
178                    buf.try_push_str(&input[start..at])?;
179                }
180
181                buf.try_push('}')?;
182                iter.next();
183            }
184            ('{', '{') => {
185                if let Some(start) = start.take() {
186                    buf.try_push_str(&input[start..at])?;
187                }
188
189                buf.try_push('{')?;
190                iter.next();
191            }
192            ('}', _) => {
193                return Err(compile::Error::msg(
194                    span,
195                    "unsupported close `}`, if you meant to escape this use `}}`",
196                ));
197            }
198            ('{', _) => {
199                if let Some(start) = start.take() {
200                    buf.try_push_str(&input[start..at])?;
201                }
202
203                if !buf.is_empty() {
204                    components.try_push(C::Literal(Box::try_from(&buf[..])?))?;
205                    buf.clear();
206                }
207
208                components.try_push(parse_group(
209                    cx,
210                    span,
211                    &mut iter,
212                    &mut count,
213                    &mut name,
214                    &mut width,
215                    &mut precision,
216                    pos,
217                    unused_pos,
218                    named,
219                    unused_named,
220                )?)?;
221            }
222            _ => {
223                if start.is_none() {
224                    start = Some(at);
225                }
226            }
227        }
228    }
229
230    if let Some(start) = start.take() {
231        buf.try_push_str(&input[start..])?;
232    }
233
234    if !buf.is_empty() {
235        components.try_push(C::Literal(Box::try_from(&buf[..])?))?;
236        buf.clear();
237    }
238
239    if components.is_empty() {
240        return Ok(quote!(""));
241    }
242
243    let mut args = Vec::<Quote<'static>>::new();
244
245    for c in components {
246        match c {
247            C::Literal(literal) => {
248                let lit = cx.lit(literal.as_ref())?;
249                args.try_push(quote!(#lit))?;
250            }
251            C::Format {
252                expr,
253                fill,
254                align,
255                width,
256                precision,
257                flags,
258                format_type,
259            } => {
260                let mut specs = Vec::new();
261
262                let fill = fill
263                    .map(|fill| {
264                        let fill = cx.lit(fill)?;
265                        Ok::<_, alloc::Error>(quote!(fill = #fill))
266                    })
267                    .transpose()?;
268
269                let width = width
270                    .map(|width| {
271                        let width = cx.lit(width)?;
272                        Ok::<_, alloc::Error>(quote!(width = #width))
273                    })
274                    .transpose()?;
275
276                let precision = precision
277                    .map(|precision| {
278                        let precision = cx.lit(precision)?;
279                        Ok::<_, alloc::Error>(quote!(precision = #precision))
280                    })
281                    .transpose()?;
282
283                let align = align
284                    .map(|align| {
285                        let align = align.try_to_string()?;
286                        let align = cx.ident(&align)?;
287                        Ok::<_, alloc::Error>(quote!(align = #align))
288                    })
289                    .transpose()?;
290
291                specs.try_extend(fill)?;
292                specs.try_extend(width)?;
293                specs.try_extend(precision)?;
294                specs.try_extend(align)?;
295
296                if !flags.is_empty() {
297                    let flags = cx.lit(flags.into_u32())?;
298                    specs.try_push(quote!(flags = #flags))?;
299                }
300
301                let format_type = format_type
302                    .map(|format_type| {
303                        let format_type = format_type.try_to_string()?;
304                        let format_type = cx.ident(&format_type)?;
305                        Ok::<_, alloc::Error>(quote!(type = #format_type))
306                    })
307                    .transpose()?;
308
309                specs.try_extend(format_type)?;
310
311                if specs.is_empty() {
312                    args.try_push(quote!(#expr))?;
313                } else {
314                    args.try_push(quote!(
315                        #[builtin]
316                        format!(#expr, #(specs),*)
317                    ))?;
318                }
319            }
320        }
321    }
322
323    return Ok(quote! {
324        #[builtin] template!(#(args),*)
325    });
326
327    enum ExprOrIdent<'a> {
328        Expr(&'a ast::Expr),
329        Ident(ast::Ident),
330    }
331
332    impl ToTokens for ExprOrIdent<'_> {
333        fn to_tokens(
334            &self,
335            cx: &mut MacroContext<'_, '_, '_>,
336            stream: &mut TokenStream,
337        ) -> alloc::Result<()> {
338            match self {
339                Self::Expr(expr) => expr.to_tokens(cx, stream),
340                Self::Ident(ident) => ident.to_tokens(cx, stream),
341            }
342        }
343    }
344
345    enum C<'a> {
346        Literal(Box<str>),
347        Format {
348            expr: ExprOrIdent<'a>,
349            fill: Option<char>,
350            align: Option<format::Alignment>,
351            width: Option<usize>,
352            precision: Option<usize>,
353            flags: format::Flags,
354            format_type: Option<format::Type>,
355        },
356    }
357
358    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
359    enum Mode {
360        /// Start of parser.
361        Start,
362        // Parse alignment.
363        FillAllign,
364        // '+' or '-' encountered.
365        Sign,
366        // Alternate '#' encountered.
367        Alternate,
368        // Sign aware zero pad `0` encountered.
369        SignAwareZeroPad,
370        // Parse width.
371        Width,
372        /// We've parsed precision fully already.
373        Precision,
374        // Type e.g. `?` encountered.
375        Type,
376        // Final mode.
377        End,
378    }
379
380    /// Parse a single expansion group.
381    fn parse_group<'a>(
382        cx: &mut MacroContext<'_, '_, '_>,
383        span: Span,
384        iter: &mut Iter<'_>,
385        count: &mut usize,
386        name: &mut String,
387        width: &mut String,
388        precision: &mut String,
389        pos: &[&'a ast::Expr],
390        unused_pos: &mut BTreeSet<usize>,
391        named: &HashMap<Box<str>, &'a NamedFormatArg>,
392        unused_named: &mut BTreeMap<Box<str>, Span>,
393    ) -> compile::Result<C<'a>> {
394        use num::ToPrimitive as _;
395
396        // Parsed flags.
397        let mut flags = format::Flags::default();
398        // Parsed fill character.
399        let mut fill = None;
400        // Parsed alignment.
401        let mut align = None;
402        // We are expecting to receive precision as a positional parameter.
403        let mut input_precision = false;
404        // Parsed formatting type.
405        let mut format_type = None;
406
407        // Clear re-used temporary buffers.
408        name.clear();
409        width.clear();
410        precision.clear();
411
412        let mut mode = Mode::Start;
413
414        loop {
415            let Some((_, a, b)) = iter.current() else {
416                return Err(compile::Error::msg(span, "unexpected end of format string"));
417            };
418
419            match mode {
420                Mode::Start => match a {
421                    ':' => {
422                        mode = Mode::FillAllign;
423                        iter.next();
424                    }
425                    '}' => {
426                        mode = Mode::End;
427                    }
428                    c => {
429                        name.try_push(c)?;
430                        iter.next();
431                    }
432                },
433                Mode::FillAllign => {
434                    // NB: parse alignment, if present.
435                    if matches!(a, '<' | '^' | '>') {
436                        align = Some(parse_align(a));
437                        iter.next();
438                    } else if matches!(b, '<' | '^' | '>') {
439                        fill = Some(a);
440                        align = Some(parse_align(b));
441
442                        iter.next();
443                        iter.next();
444                    }
445
446                    mode = Mode::Sign;
447                }
448                Mode::Sign => {
449                    match a {
450                        '-' => {
451                            flags.set(format::Flag::SignMinus);
452                            iter.next();
453                        }
454                        '+' => {
455                            flags.set(format::Flag::SignPlus);
456                            iter.next();
457                        }
458                        _ => (),
459                    }
460
461                    mode = Mode::Alternate;
462                }
463                Mode::Alternate => {
464                    if a == '#' {
465                        flags.set(format::Flag::Alternate);
466                        iter.next();
467                    }
468
469                    mode = Mode::SignAwareZeroPad;
470                }
471                Mode::SignAwareZeroPad => {
472                    if a == '0' {
473                        flags.set(format::Flag::SignAwareZeroPad);
474                        iter.next();
475                    }
476
477                    mode = Mode::Width;
478                }
479                Mode::Width => {
480                    match a {
481                        '0'..='9' => {
482                            width.try_push(a)?;
483                            iter.next();
484                            continue;
485                        }
486                        '.' => {
487                            mode = Mode::Precision;
488                            iter.next();
489                            continue;
490                        }
491                        _ => (),
492                    }
493
494                    mode = Mode::Type;
495                }
496                Mode::Precision => {
497                    match a {
498                        '*' if precision.is_empty() => {
499                            input_precision = true;
500                            iter.next();
501                        }
502                        '0'..='9' => {
503                            precision.try_push(a)?;
504                            iter.next();
505                            continue;
506                        }
507                        _ => (),
508                    }
509
510                    mode = Mode::Type;
511                }
512                Mode::Type => {
513                    match a {
514                        '?' => {
515                            format_type = Some(format::Type::Debug);
516                            iter.next();
517                        }
518                        'x' => {
519                            format_type = Some(format::Type::LowerHex);
520                            iter.next();
521                        }
522                        'X' => {
523                            format_type = Some(format::Type::UpperHex);
524                            iter.next();
525                        }
526                        'b' => {
527                            format_type = Some(format::Type::Binary);
528                            iter.next();
529                        }
530                        'p' => {
531                            format_type = Some(format::Type::Pointer);
532                            iter.next();
533                        }
534                        _ => (),
535                    }
536
537                    mode = Mode::End;
538                }
539                Mode::End => {
540                    match a {
541                        '}' => (),
542                        c => {
543                            return Err(compile::Error::msg(
544                                span,
545                                format!("unsupported char `{c}` in spec"),
546                            ));
547                        }
548                    }
549
550                    iter.next();
551                    break;
552                }
553            }
554        }
555
556        let precision = if input_precision {
557            let &expr = match pos.get(*count) {
558                Some(expr) => expr,
559                None => {
560                    return Err(compile::Error::msg(
561                        span,
562                        format!(
563                            "missing positional argument #{count} \
564                            which is required for position parameter",
565                        ),
566                    ));
567                }
568            };
569
570            unused_pos.remove(count);
571
572            let value = cx.eval(expr)?;
573
574            let number = match value.as_inline() {
575                Some(Inline::Signed(n)) => n.to_usize(),
576                _ => None,
577            };
578
579            let precision = if let Some(number) = number {
580                number
581            } else {
582                let span = expr.span();
583
584                return Err(compile::Error::msg(
585                    span,
586                    format!(
587                        "expected position argument #{} \
588                        to be a positive number in use as precision, \
589                        but got `{}`",
590                        count,
591                        value.type_info()
592                    ),
593                ));
594            };
595
596            *count += 1;
597            Some(precision)
598        } else if !precision.is_empty() {
599            str::parse::<usize>(precision).ok()
600        } else {
601            None
602        };
603
604        let expr = 'expr: {
605            if name.is_empty() {
606                let Some(expr) = pos.get(*count) else {
607                    return Err(compile::Error::msg(
608                        span,
609                        format!("missing positional argument #{count}"),
610                    ));
611                };
612
613                unused_pos.remove(count);
614                *count += 1;
615                break 'expr ExprOrIdent::Expr(expr);
616            };
617
618            if let Ok(n) = str::parse::<usize>(name) {
619                let expr = match pos.get(n) {
620                    Some(expr) => *expr,
621                    None => {
622                        return Err(compile::Error::msg(
623                            span,
624                            format!("missing positional argument #{n}"),
625                        ));
626                    }
627                };
628
629                unused_pos.remove(&n);
630                break 'expr ExprOrIdent::Expr(expr);
631            }
632
633            if let Some(n) = named.get(name.as_str()) {
634                unused_named.remove(name.as_str());
635                break 'expr ExprOrIdent::Expr(&n.expr);
636            }
637
638            let mut ident = cx.ident(name.as_str())?;
639            ident.span = span;
640            ExprOrIdent::Ident(ident)
641        };
642
643        let width = if !width.is_empty() {
644            str::parse::<usize>(width).ok()
645        } else {
646            None
647        };
648
649        Ok(C::Format {
650            expr,
651            fill,
652            align,
653            width,
654            precision,
655            format_type,
656            flags,
657        })
658    }
659
660    fn parse_align(c: char) -> format::Alignment {
661        match c {
662            '<' => format::Alignment::Left,
663            '^' => format::Alignment::Center,
664            _ => format::Alignment::Right,
665        }
666    }
667}
668
669struct Iter<'a> {
670    iter: str::CharIndices<'a>,
671    a: Option<(usize, char)>,
672    b: Option<(usize, char)>,
673}
674
675impl<'a> Iter<'a> {
676    fn new(input: &'a str) -> Self {
677        let mut iter = input.char_indices();
678        let a = iter.next();
679        let b = iter.next();
680        Self { iter, a, b }
681    }
682
683    fn current(&self) -> Option<(usize, char, char)> {
684        let (pos, a) = self.a?;
685        let (_, b) = self.b.unwrap_or_default();
686        Some((pos, a, b))
687    }
688}
689
690impl Iterator for Iter<'_> {
691    type Item = (usize, char, char);
692
693    fn next(&mut self) -> Option<Self::Item> {
694        let value = self.current()?;
695
696        self.a = self.b;
697        self.b = self.iter.next();
698
699        Some(value)
700    }
701}
rune/macros/format_args.rs

rune/macros/
format_args.rs
