dfir_lang/
process_singletons.rs

1//! Utility methods for processing singleton references: `#my_var`.
2
3use itertools::Itertools;
4use proc_macro2::{Group, Ident, TokenStream, TokenTree};
5use quote::quote_spanned;
6use syn::punctuated::Punctuated;
7use syn::{Expr, Token};
8
9use crate::parse::parse_terminated;
10
11/// Finds all the singleton references `#my_var` and appends them to `found_idents`. Returns the
12/// `TokenStream` but with the hashes removed from the varnames.
13///
14/// The returned tokens are used for "preflight" parsing, to check that the rest of the syntax is
15/// OK. However the returned tokens are not used in the codegen as we need to use [`postprocess_singletons`]
16/// later to substitute-in the context referencing code for each singleton
17pub fn preprocess_singletons(tokens: TokenStream, found_idents: &mut Vec<Ident>) -> TokenStream {
18    process_singletons(tokens, &mut |singleton_ident| {
19        found_idents.push(singleton_ident.clone());
20        TokenTree::Ident(singleton_ident)
21    })
22}
23
24/// Replaces singleton references `#my_var` with the code needed to actually get the value inside.
25///
26/// * `tokens` - The tokens to update singleton references within.
27/// * `resolved_idents` - The `RefCell<_>` local variable idents that correspond 1:1 and in the same
28///   order as the singleton references within `tokens` (found in-order via [`preprocess_singletons`]).
29///
30/// Generates borrowing code ([`std::cell::RefCell::borrow`]). Use
31/// [`postprocess_singletons_handles`] for just the raw idents.
32pub fn postprocess_singletons(
33    tokens: TokenStream,
34    resolved_idents: impl IntoIterator<Item = Ident>,
35    _context: &Ident,
36) -> Punctuated<Expr, Token![,]> {
37    let mut resolved_idents_iter = resolved_idents.into_iter();
38    let processed = process_singletons(tokens, &mut |singleton_ident| {
39        let span = singleton_ident.span();
40        let mut resolved_ident = resolved_idents_iter.next().unwrap();
41        resolved_ident.set_span(span);
42        let mut group = Group::new(
43            proc_macro2::Delimiter::Parenthesis,
44            quote_spanned! {span=>
45                *#resolved_ident.borrow()
46            },
47        );
48        group.set_span(singleton_ident.span());
49        TokenTree::Group(group)
50    });
51    parse_terminated(processed).unwrap()
52}
53
54/// Same as [`postprocess_singletons`] but generates just the raw ident rather than
55/// `RefCell` borrowing code.
56pub fn postprocess_singletons_handles(
57    tokens: TokenStream,
58    resolved_idents: impl IntoIterator<Item = Ident>,
59) -> Punctuated<Expr, Token![,]> {
60    let mut resolved_idents_iter = resolved_idents.into_iter();
61    let processed = process_singletons(tokens, &mut |singleton_ident| {
62        let mut resolved_ident = resolved_idents_iter.next().unwrap();
63        resolved_ident.set_span(singleton_ident.span().resolved_at(resolved_ident.span()));
64        TokenTree::Ident(resolved_ident)
65    });
66    parse_terminated(processed).unwrap()
67}
68
69/// Traverse the token stream, applying the `map_singleton_fn` whenever a singleton is found,
70/// returning the transformed token stream.
71fn process_singletons(
72    tokens: TokenStream,
73    map_singleton_fn: &mut impl FnMut(Ident) -> TokenTree,
74) -> TokenStream {
75    tokens
76        .into_iter()
77        .peekable()
78        .batching(|iter| {
79            let out = match iter.next()? {
80                TokenTree::Group(group) => {
81                    let mut new_group = Group::new(
82                        group.delimiter(),
83                        process_singletons(group.stream(), map_singleton_fn),
84                    );
85                    new_group.set_span(group.span());
86                    TokenTree::Group(new_group)
87                }
88                TokenTree::Ident(ident) => TokenTree::Ident(ident),
89                TokenTree::Punct(punct) => {
90                    if '#' == punct.as_char() && matches!(iter.peek(), Some(TokenTree::Ident(_))) {
91                        // Found a singleton.
92                        let Some(TokenTree::Ident(mut singleton_ident)) = iter.next() else {
93                            unreachable!()
94                        };
95                        {
96                            // Include the `#` in the span.
97                            let span = singleton_ident
98                                .span()
99                                .join(punct.span())
100                                .unwrap_or(singleton_ident.span());
101                            singleton_ident.set_span(span.resolved_at(singleton_ident.span()));
102                        }
103                        (map_singleton_fn)(singleton_ident)
104                    } else {
105                        TokenTree::Punct(punct)
106                    }
107                }
108                TokenTree::Literal(lit) => TokenTree::Literal(lit),
109            };
110            Some(out)
111        })
112        .collect()
113}
dfir_lang/process_singletons.rs

dfir_lang/
process_singletons.rs
