Skip to main content

token_goblin_runtime/
ux.rs

1//! Better UX for proc-macro.
2//! Inspired by `crabtime`.
3//!
4//! Allows to receiving inputs and producing outputs in non `TokenStream` way.
5//! This is the boring goblin craft: fewer raw token piles, more typed little bundles.
6//!
7//! E.g. instead of:
8//! ```
9//! # use proc_macro2::TokenStream;
10//! # use syn::parse::Parser;
11//!
12//!
13//! fn foo(input: TokenStream) -> TokenStream {
14//!    let parser = syn::punctuated::Punctuated::<syn::LitStr, syn::Token![,]>::parse_terminated;
15//!    let lit_components = parser.parse2(input).unwrap();
16//!    let components = lit_components.iter().map(|c| c.value()).collect::<Vec<_>>();
17//!    // Handling of `components`
18//!    # todo!()
19//! }
20//! ```
21//!
22//! One could write:
23//! ```
24//! # use proc_macro2::TokenStream;
25//! # use syn::parse::Parser;
26//! # use token_goblin_runtime::prelude::*;
27//!
28//! fn foo(components: CommaSeparated<Token>) -> TokenStream {
29//!    // Handling of `components`
30//!    # todo!()
31//! }
32//! ```
33//!
34//! Since extending `syn::parse::Parse` with std types is not possible due to orphan rule.
35//! We use macro `parse_into!`, that hardcodes checks for specific types.
36//!
37//! Note: having `String` and `Vec<String>` in input params remove span information, and reduce IDE/diagnostics quality.
38//!
39//! Output is a little bit more simple, it expected in three forms:
40//! - `String` - For strings that should be converted to `TokenStream` without input span information
41//! - `TokenStream` - as basic case.
42//! - and in empty form - for cases where output is already emitted as `output_str!`, `output!` macros.
43//!
44//! So we have a trait `IntoTokenStream` that is solely focused on converting specific types into `TokenStream`.
45//!
46//! The user can extend it as well, to support custom types in output.
47
48use core::fmt::{self, Display};
49use std::{cell::RefCell, fmt::Debug, str::FromStr};
50
51use proc_macro2::{Span, TokenStream};
52use quote::ToTokens;
53use syn::parse::{Parse, ParseStream, Parser};
54
55/// Represents a single entry in `snif!` route.
56///
57/// Used to represent a single result of handled `snif!` macro.
58///
59/// Example:
60/// ```no_build
61/// #[derive(token_goblin::Snif)]
62/// struct Foo {
63///     x: i32,
64/// }
65///
66/// token_goblin::snif!(Foo in stringify_our!()); // -> "Foo => { struct Foo { x : i32, } }" is one entry
67/// ```
68///
69#[derive(Clone)]
70pub struct SnifedEntry {
71    /// Path to macro that was used to generate this entry.
72    pub snif_path: syn::Path,
73    arrow: syn::Token![=>],
74    brace: syn::token::Brace,
75    /// Item that was snifed.
76    pub item: syn::Item,
77}
78/// Represents a group of `snif!` entries.
79///
80/// Used to represent a group of `snif!` entries in a macro.
81///
82/// Example:
83/// ```no_build
84///  token_goblin::snif!(Foo, Bar in stringify_our!("extra tokens"));
85///  // -> "Foo => { struct Foo { x : i32, } }" is one entry
86///  // "Bar => { struct Bar { x : i32, } }" is another
87///  //
88///  // "extra tokens" is passed to the macro as input.
89/// ```
90///
91#[derive(Clone)]
92pub struct SnifedEntries {
93    first_group: syn::token::Bracket,
94    pub entries: Vec<SnifedEntry>,
95    second_group: syn::token::Bracket,
96    pub macro_input: TokenStream,
97}
98impl SnifedEntries {
99    #[must_use]
100    pub fn span(&self) -> proc_macro2::Span {
101        self.entries
102            .first()
103            .map_or_else(Span::call_site, SnifedEntry::span)
104    }
105}
106impl SnifedEntry {
107    #[must_use]
108    pub fn span(&self) -> proc_macro2::Span {
109        self.snif_path
110            .segments
111            .first()
112            .map_or_else(Span::call_site, |segment| segment.ident.span())
113    }
114}
115/// Represents a comma separated list of parsable values.
116///
117/// *A tidy little bundle of tokens, comma-sorted by the goblin before it hands them over.*
118///
119/// Can be used to provide a typed interface for input params of `token-goblin` `charms`.
120///
121/// Example:
122/// ```no_build
123/// #[token_goblin::munch]
124/// fn foo(input: CommaSeparated<syn::LitStr>) -> TokenStream {
125///     output_str!("{}", input.0.iter().map(|s| s.value()).collect::<Vec<_>>().join(", "));
126/// }
127///
128/// foo!("foo", "bar", "baz");
129/// // -> "foo, bar, baz"
130/// ```
131///
132pub struct CommaSeparated<T>(pub Vec<T>);
133
134impl From<CommaSeparated<Token>> for Vec<String> {
135    fn from(value: CommaSeparated<Token>) -> Self {
136        value.0.into_iter().map(|t| t.to_string()).collect()
137    }
138}
139
140/// Represents either `Ident` or `LitStr` token.
141///
142/// *A bare ident or a quoted string - the goblin chews both the same.*
143///
144/// Used when macro need a simple interface for input, and user can decide a way to provide string.
145///
146/// Example:
147/// ```no_build
148/// #[token_goblin::munch]
149/// fn foo(input: Token) -> TokenStream {
150///     output_str!("{}", input.to_string());
151/// }
152///
153/// foo!("foo");
154/// // -> foo
155///
156pub enum Token {
157    Ident(syn::Ident),
158    Literal(syn::LitStr),
159}
160impl Display for Token {
161    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
162        match self {
163            Token::Ident(ident) => write!(f, "{ident}"),
164            Token::Literal(literal) => write!(f, "{}", literal.value()),
165        }
166    }
167}
168
169impl Debug for Token {
170    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
171        match self {
172            Token::Ident(ident) => write!(f, "Ident({ident:?})"),
173            Token::Literal(literal) => write!(f, "Literal({:?})", literal.value()),
174        }
175    }
176}
177impl PartialEq<&str> for Token {
178    fn eq(&self, other: &&str) -> bool {
179        match self {
180            Token::Ident(ident) => ident == *other,
181            // creates an owned string (but we don't have an api to compare directly)
182            Token::Literal(literal) => literal.value() == *other,
183        }
184    }
185}
186
187#[doc(hidden)] // auto trait for FromTokenStream
188pub trait TokenStreamInto<T> {
189    fn convert_token_stream(self) -> syn::Result<T>;
190}
191impl<T: syn::parse::Parse> TokenStreamInto<T> for TokenStream {
192    fn convert_token_stream(self) -> syn::Result<T> {
193        T::parse.parse2(self)
194    }
195}
196
197/// Convert specific type into `TokenStream`.
198///
199/// In `token-goblin` it is used to convert output types of `token-goblin` `charms` into `TokenStream`.
200/// We provide default implementations for:
201/// - `String`, `TokenStream`, `()` - so them can be used as output for `charm` fn
202///   out of the box.
203///
204/// For `#[munch] mod {..}` user can provide custom implementation, to support custom types in output.
205pub trait IntoTokenStream {
206    fn into_token_stream(self) -> TokenStream;
207}
208
209impl IntoTokenStream for String {
210    fn into_token_stream(self) -> TokenStream {
211        TokenStream::from_str(&self).unwrap_or_else(|e| {
212            compile_error(&format!("Failed to convert String to TokenStream: {e}"))
213        })
214    }
215}
216impl IntoTokenStream for TokenStream {
217    fn into_token_stream(self) -> TokenStream {
218        self
219    }
220}
221impl IntoTokenStream for () {
222    fn into_token_stream(self) -> TokenStream {
223        TokenStream::new()
224    }
225}
226
227fn compile_error(text: &str) -> TokenStream {
228    quote::quote! {
229        ::core::compile_error!(#text)
230    }
231}
232
233/// Emit formatted string as token stream.
234///
235/// *The goblin's quick spit: hand it a string, it coughs up tokens.*
236///
237/// Example:
238/// ```
239/// # use token_goblin_runtime::prelude::*;
240/// output_str!("foo + 2");
241/// ```
242///
243/// This will spit `foo + 2` token stream (ident, punct, literal) as output of the macro, just before emitting result.
244///
245/// The format of input is the same as in `format!` macro.
246///
247/// Note: If input is invalid `TokenStream` this will emit compile error.
248#[macro_export]
249macro_rules! output_str {
250    ($($tokens:tt)*) => {
251        $crate::ux::push_output(format!($($tokens)*));
252    };
253}
254
255/// Emit quote as token stream.
256///
257/// Example:
258/// ```
259/// # use token_goblin_runtime::prelude::*;
260/// output! {
261///     foo + bar
262/// };
263/// ```
264///
265/// This will spit quoted `TokenStream` as output of the macro, just before emitting result.
266/// The format of input is the same as in `quote!` macro.
267///
268/// Note: that this is different from `output_str!` macro:
269/// ```
270/// # use token_goblin_runtime::prelude::*;
271/// output_str!("foo + 2");
272/// output! {
273///     "foo + 2"
274/// };
275/// ```
276///
277/// The first will emit `foo + 2` token stream (ident, punct, literal) as output of the macro.
278/// But the second one will emit `"foo + 2"` as string literal.
279///
280#[macro_export]
281macro_rules! output {
282    ($($tokens:tt)*) => {
283        $crate::ux::push_output($crate::prelude::quote!($($tokens)*));
284    };
285}
286
287thread_local! {
288    static COLLECTED_OUTPUT: RefCell<TokenStream> = RefCell::new(TokenStream::new());
289}
290
291/// For some usages, user might want to emit output streamingly, like `println!` or `write!` macros.
292///
293/// This function is internall implementation of this feature, it's recommended to use:
294/// `output!`, or `output_str!` macros instead.
295pub fn push_output(output: impl IntoTokenStream) {
296    COLLECTED_OUTPUT.with(|collected_output| {
297        collected_output
298            .borrow_mut()
299            .extend(output.into_token_stream());
300    });
301}
302
303#[doc(hidden)]
304#[must_use]
305pub(crate) fn flush_output(last_part: TokenStream) -> TokenStream {
306    COLLECTED_OUTPUT.with(|collected_output| {
307        let mut collected_output = std::mem::take(&mut *collected_output.borrow_mut());
308        collected_output.extend(last_part);
309        collected_output
310    })
311}
312
313impl Parse for Token {
314    fn parse(input: ParseStream) -> syn::Result<Self> {
315        if input.peek(syn::Ident) {
316            Ok(Token::Ident(input.parse()?))
317        } else if input.peek(syn::LitStr) {
318            Ok(Token::Literal(input.parse()?))
319        } else {
320            Err(syn::Error::new(input.span(), "Expected ident or literal"))
321        }
322    }
323}
324
325impl<T: Parse> Parse for CommaSeparated<T> {
326    fn parse(input: ParseStream) -> syn::Result<Self> {
327        let parser = syn::punctuated::Punctuated::<T, syn::Token![,]>::parse_terminated;
328        let components = parser(input)?;
329        Ok(CommaSeparated(components.into_iter().collect()))
330    }
331}
332
333impl syn::parse::Parse for SnifedEntry {
334    fn parse(input: ParseStream) -> syn::Result<Self> {
335        // Skip ident + `::`, find `=>` in tokenstream. then feed bounded stream into `syn::Path::parse`
336
337        let path = syn::Path::parse_mod_style(input)?;
338
339        let arrow = input.parse()?;
340
341        let content;
342        let brace = syn::braced!(content in input);
343        let item = content.parse()?;
344
345        Ok(SnifedEntry {
346            snif_path: path,
347            arrow,
348            brace,
349            item,
350        })
351    }
352}
353impl syn::parse::Parse for SnifedEntries {
354    fn parse(input: ParseStream) -> syn::Result<Self> {
355        let items_input;
356        let first_group = syn::bracketed!(items_input in input);
357        let mut items = Vec::new();
358        while !items_input.is_empty() {
359            items.push(SnifedEntry::parse(&items_input)?);
360        }
361        let macro_input;
362        let second_group = syn::bracketed!(macro_input in input);
363
364        Ok(SnifedEntries {
365            first_group,
366            entries: items,
367            second_group,
368            macro_input: macro_input.parse()?,
369        })
370    }
371}
372impl ToTokens for SnifedEntry {
373    fn to_tokens(&self, tokens: &mut TokenStream) {
374        self.snif_path.to_tokens(tokens);
375        self.arrow.to_tokens(tokens);
376        self.brace.surround(tokens, |tokens| {
377            self.item.to_tokens(tokens);
378        });
379    }
380}
381impl ToTokens for SnifedEntries {
382    fn to_tokens(&self, tokens: &mut TokenStream) {
383        self.first_group.surround(tokens, |tokens| {
384            for item in &self.entries {
385                item.to_tokens(tokens);
386            }
387        });
388        self.second_group.surround(tokens, |tokens| {
389            self.macro_input.to_tokens(tokens);
390        });
391    }
392}
393#[cfg(test)]
394mod tests {
395    use std::str::FromStr;
396
397    use super::*;
398
399    #[test]
400    fn test_parse_string() {
401        let tokens = TokenStream::from_str(" \"123\" ").unwrap();
402        let into: Token = tokens.convert_token_stream().unwrap();
403        assert_eq!(into.to_string(), "123");
404    }
405    #[test]
406    fn test_parse_vec() {
407        let tokens = TokenStream::from_str(" \"1\", \"2\", \"3\" ").unwrap();
408        let into: CommaSeparated<Token> = tokens.convert_token_stream().unwrap();
409        assert_eq!(into.0, vec!["1", "2", "3"]);
410    }
411
412    #[test]
413    fn test_parse_tts() {
414        let tokens = TokenStream::from_str("123").unwrap();
415        let into: TokenStream = tokens.clone().convert_token_stream().unwrap();
416        assert_eq!(into.to_string(), tokens.to_string());
417    }
418
419    #[test]
420    fn test_parse_syn_type() {
421        let tokens = TokenStream::from_str("asd").unwrap();
422        let into: syn::Ident = tokens.convert_token_stream().unwrap();
423        assert_eq!(into.to_string(), "asd");
424    }
425
426    #[test]
427    fn test_streaming_output() {
428        output_str!("foo");
429        output_str!("bar");
430        output! {
431            "baz" // quote will emit tokens so this becumes string literal
432        };
433        let output = flush_output(TokenStream::from_str("qux").unwrap());
434        assert_eq!(output.to_string(), "foo bar \"baz\" qux");
435    }
436}