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

rune/macros/
format_args.rs
