dfir_lang/
process_singletons.rs

1//! Utility methods for processing singleton references: `#my_var`, `#mut my_var`, `#{N} my_var`, `#{N} mut my_var`.
2
3use proc_macro2::{Group, TokenStream, TokenTree};
4use syn::parse::Parser;
5use syn::punctuated::Punctuated;
6use syn::{Expr, Token};
7
8use crate::parse::SingletonRef;
9
10/// Finds all the singleton references and appends them to `found`. Returns the
11/// `TokenStream` but with the `#`, `{N}`, and `mut` removed from the varnames.
12///
13/// Syntax: `#var`, `#mut var`, `#{N} var`, `#{N} mut var`
14///
15/// The returned tokens are used for "preflight" parsing, to check that the rest of the syntax is
16/// OK. However the returned tokens are not used in the codegen as we need to use [`postprocess_singletons`]
17/// later to substitute-in the context referencing code for each singleton
18pub fn preprocess_singletons(tokens: TokenStream, found: &mut Vec<SingletonRef>) -> TokenStream {
19    process_singletons(tokens, &mut |ref_token| {
20        let ident = ref_token.ident.clone();
21        found.push(ref_token);
22        TokenTree::Ident(ident)
23    })
24}
25
26/// Replaces singleton references with the code needed to actually get the value inside.
27///
28/// * `tokens` - The tokens to update singleton references within.
29/// * `resolved_exprs` - Token streams that correspond 1:1 and in the same
30///   order as the singleton references within `tokens` (found in-order via [`preprocess_singletons`]).
31///
32/// For shared refs: splices the resolved expression directly (e.g. `&T` or `&Option<T>`).
33/// For mutable refs: splices the resolved expression directly (e.g. `&mut T` or `&mut Option<T>`).
34pub fn postprocess_singletons(
35    tokens: TokenStream,
36    resolved_exprs: impl IntoIterator<Item = TokenStream>,
37) -> Punctuated<Expr, Token![,]> {
38    let mut resolved_exprs_iter = resolved_exprs.into_iter();
39    let processed = process_singletons(tokens, &mut |_ref_token| {
40        let span = _ref_token.ident.span();
41        let expr_tokens = resolved_exprs_iter.next().unwrap();
42        // Wrap in parentheses to preserve precedence.
43        let mut group = Group::new(proc_macro2::Delimiter::Parenthesis, expr_tokens);
44        group.set_span(span);
45        TokenTree::Group(group)
46    });
47    Punctuated::parse_terminated.parse2(processed).unwrap()
48}
49
50/// Traverse the token stream, applying the `map_singleton_fn` whenever a singleton is found,
51/// returning the transformed token stream.
52///
53/// Parses: `#ident`, `#mut ident`, `#{N} ident`, `#{N} mut ident`
54fn process_singletons(
55    tokens: TokenStream,
56    map_singleton_fn: &mut impl FnMut(SingletonRef) -> TokenTree,
57) -> TokenStream {
58    let mut iter = tokens.into_iter().peekable();
59    std::iter::from_fn(|| {
60        let out = match iter.peek()? {
61            TokenTree::Group(group) => {
62                let mut new_group = Group::new(
63                    group.delimiter(),
64                    process_singletons(group.stream(), map_singleton_fn),
65                );
66                new_group.set_span(group.span());
67
68                let _ = iter.next().unwrap(); // Advance past the `peek`ed group.
69                TokenTree::Group(new_group)
70            }
71            TokenTree::Punct(punct) if '#' == punct.as_char() => {
72                let tokens = iter.by_ref().collect::<TokenStream>();
73                let (opt_singleton, tokens_rest) = SingletonRef::try_parse
74                    .parse2(tokens)
75                    .expect("bug: should be infallible");
76                iter = tokens_rest.into_iter().peekable();
77                if let Some(singleton) = opt_singleton {
78                    (map_singleton_fn)(singleton)
79                } else {
80                    iter.next().unwrap()
81                }
82            }
83            _ => iter.next().unwrap(),
84        };
85        Some(out)
86    })
87    .collect()
88}
dfir_lang/process_singletons.rs

dfir_lang/
process_singletons.rs
