Skip to main content

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;
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    /// The largest width or precision a value can be written with.
381    ///
382    /// A precision is handed to the formatter the platform provides, which
383    /// takes one that fits in a `u16` and panics on anything larger. A width is
384    /// written out here instead, but it is held to the same bound so that a
385    /// specification is not accepted under one and rejected under the other.
386    const MAX_FORMAT_ARGUMENT: usize = u16::MAX as usize;
387
388    /// Bound a width or a precision which was given as a number.
389    fn bound_format_argument(span: Span, what: &str, n: usize) -> compile::Result<usize> {
390        if n > MAX_FORMAT_ARGUMENT {
391            return Err(compile::Error::msg(
392                span,
393                format!("{what} {n} is larger than the maximum of {MAX_FORMAT_ARGUMENT}"),
394            ));
395        }
396
397        Ok(n)
398    }
399
400    /// Parse the digits a width or a precision was written as.
401    ///
402    /// A number too large to use was dropped rather than reported, so a
403    /// specification carrying one was honoured as if it had not been written.
404    fn parse_format_argument(span: Span, what: &str, digits: &str) -> compile::Result<usize> {
405        // Digits which do not parse are ones too large to hold, which is past
406        // the bound either way.
407        let Ok(n) = str::parse::<usize>(digits) else {
408            return Err(compile::Error::msg(
409                span,
410                format!("{what} {digits} is larger than the maximum of {MAX_FORMAT_ARGUMENT}"),
411            ));
412        };
413
414        bound_format_argument(span, what, n)
415    }
416
417    /// Parse a single expansion group.
418    fn parse_group<'a>(
419        cx: &mut MacroContext<'_, '_, '_>,
420        span: Span,
421        iter: &mut Iter<'_>,
422        count: &mut usize,
423        name: &mut String,
424        width: &mut String,
425        precision: &mut String,
426        pos: &[&'a ast::Expr],
427        unused_pos: &mut BTreeSet<usize>,
428        named: &HashMap<Box<str>, &'a NamedFormatArg>,
429        unused_named: &mut BTreeMap<Box<str>, Span>,
430    ) -> compile::Result<C<'a>> {
431        // Parsed flags.
432        let mut flags = format::Flags::default();
433        // Parsed fill character.
434        let mut fill = None;
435        // Parsed alignment.
436        let mut align = None;
437        // We are expecting to receive precision as a positional parameter.
438        let mut input_precision = false;
439        // Parsed formatting type.
440        let mut format_type = None;
441
442        // Clear re-used temporary buffers.
443        name.clear();
444        width.clear();
445        precision.clear();
446
447        let mut mode = Mode::Start;
448
449        loop {
450            let Some((_, a, b)) = iter.current() else {
451                return Err(compile::Error::msg(span, "unexpected end of format string"));
452            };
453
454            match mode {
455                Mode::Start => match a {
456                    ':' => {
457                        mode = Mode::FillAllign;
458                        iter.next();
459                    }
460                    '}' => {
461                        mode = Mode::End;
462                    }
463                    c => {
464                        name.try_push(c)?;
465                        iter.next();
466                    }
467                },
468                Mode::FillAllign => {
469                    // NB: parse alignment, if present.
470                    if matches!(a, '<' | '^' | '>') {
471                        align = Some(parse_align(a));
472                        iter.next();
473                    } else if matches!(b, '<' | '^' | '>') {
474                        fill = Some(a);
475                        align = Some(parse_align(b));
476
477                        iter.next();
478                        iter.next();
479                    }
480
481                    mode = Mode::Sign;
482                }
483                Mode::Sign => {
484                    match a {
485                        '-' => {
486                            flags.set(format::Flag::SignMinus);
487                            iter.next();
488                        }
489                        '+' => {
490                            flags.set(format::Flag::SignPlus);
491                            iter.next();
492                        }
493                        _ => (),
494                    }
495
496                    mode = Mode::Alternate;
497                }
498                Mode::Alternate => {
499                    if a == '#' {
500                        flags.set(format::Flag::Alternate);
501                        iter.next();
502                    }
503
504                    mode = Mode::SignAwareZeroPad;
505                }
506                Mode::SignAwareZeroPad => {
507                    if a == '0' {
508                        flags.set(format::Flag::SignAwareZeroPad);
509                        iter.next();
510                    }
511
512                    mode = Mode::Width;
513                }
514                Mode::Width => {
515                    match a {
516                        '0'..='9' => {
517                            width.try_push(a)?;
518                            iter.next();
519                            continue;
520                        }
521                        '.' => {
522                            mode = Mode::Precision;
523                            iter.next();
524                            continue;
525                        }
526                        _ => (),
527                    }
528
529                    mode = Mode::Type;
530                }
531                Mode::Precision => {
532                    match a {
533                        '*' if precision.is_empty() => {
534                            input_precision = true;
535                            iter.next();
536                        }
537                        '0'..='9' => {
538                            precision.try_push(a)?;
539                            iter.next();
540                            continue;
541                        }
542                        _ => (),
543                    }
544
545                    mode = Mode::Type;
546                }
547                Mode::Type => {
548                    match a {
549                        '?' => {
550                            format_type = Some(format::Type::Debug);
551                            iter.next();
552                        }
553                        'x' => {
554                            format_type = Some(format::Type::LowerHex);
555                            iter.next();
556                        }
557                        'X' => {
558                            format_type = Some(format::Type::UpperHex);
559                            iter.next();
560                        }
561                        'b' => {
562                            format_type = Some(format::Type::Binary);
563                            iter.next();
564                        }
565                        'p' => {
566                            format_type = Some(format::Type::Pointer);
567                            iter.next();
568                        }
569                        _ => (),
570                    }
571
572                    mode = Mode::End;
573                }
574                Mode::End => {
575                    match a {
576                        '}' => (),
577                        c => {
578                            return Err(compile::Error::msg(
579                                span,
580                                format!("unsupported char `{c}` in spec"),
581                            ));
582                        }
583                    }
584
585                    iter.next();
586                    break;
587                }
588            }
589        }
590
591        let precision = if input_precision {
592            let &expr = match pos.get(*count) {
593                Some(expr) => expr,
594                None => {
595                    return Err(compile::Error::msg(
596                        span,
597                        format!(
598                            "missing positional argument #{count} \
599                            which is required for position parameter",
600                        ),
601                    ));
602                }
603            };
604
605            unused_pos.remove(count);
606
607            let value = cx.eval(expr)?;
608            let precision = value.as_usize().with_span(span)?;
609
610            *count += 1;
611            Some(bound_format_argument(span, "precision", precision)?)
612        } else if !precision.is_empty() {
613            Some(parse_format_argument(span, "precision", precision)?)
614        } else {
615            None
616        };
617
618        let expr = 'expr: {
619            if name.is_empty() {
620                let Some(expr) = pos.get(*count) else {
621                    return Err(compile::Error::msg(
622                        span,
623                        format!("missing positional argument #{count}"),
624                    ));
625                };
626
627                unused_pos.remove(count);
628                *count += 1;
629                break 'expr ExprOrIdent::Expr(expr);
630            };
631
632            if let Ok(n) = str::parse::<usize>(name) {
633                let expr = match pos.get(n) {
634                    Some(expr) => *expr,
635                    None => {
636                        return Err(compile::Error::msg(
637                            span,
638                            format!("missing positional argument #{n}"),
639                        ));
640                    }
641                };
642
643                unused_pos.remove(&n);
644                break 'expr ExprOrIdent::Expr(expr);
645            }
646
647            if let Some(n) = named.get(name.as_str()) {
648                unused_named.remove(name.as_str());
649                break 'expr ExprOrIdent::Expr(&n.expr);
650            }
651
652            let mut ident = cx.ident(name.as_str())?;
653            ident.span = span;
654            ExprOrIdent::Ident(ident)
655        };
656
657        let width = if !width.is_empty() {
658            Some(parse_format_argument(span, "width", width)?)
659        } else {
660            None
661        };
662
663        Ok(C::Format {
664            expr,
665            fill,
666            align,
667            width,
668            precision,
669            format_type,
670            flags,
671        })
672    }
673
674    fn parse_align(c: char) -> format::Alignment {
675        match c {
676            '<' => format::Alignment::Left,
677            '^' => format::Alignment::Center,
678            _ => format::Alignment::Right,
679        }
680    }
681}
682
683struct Iter<'a> {
684    iter: str::CharIndices<'a>,
685    a: Option<(usize, char)>,
686    b: Option<(usize, char)>,
687}
688
689impl<'a> Iter<'a> {
690    fn new(input: &'a str) -> Self {
691        let mut iter = input.char_indices();
692        let a = iter.next();
693        let b = iter.next();
694        Self { iter, a, b }
695    }
696
697    fn current(&self) -> Option<(usize, char, char)> {
698        let (pos, a) = self.a?;
699        let (_, b) = self.b.unwrap_or_default();
700        Some((pos, a, b))
701    }
702}
703
704impl Iterator for Iter<'_> {
705    type Item = (usize, char, char);
706
707    fn next(&mut self) -> Option<Self::Item> {
708        let value = self.current()?;
709
710        self.a = self.b;
711        self.b = self.iter.next();
712
713        Some(value)
714    }
715}