Skip to main content

virtue_next/generate/
stream_builder.rs

1use crate::prelude::{
2    Delimiter, Group, Ident, LexError, Literal, Punct, Result, Spacing, Span, TokenStream,
3    TokenTree,
4};
5use std::str::FromStr;
6
7/// A helper struct build around a [TokenStream] to make it easier to build code.
8#[must_use]
9#[derive(Default)]
10pub struct StreamBuilder {
11    pub(crate) stream: TokenStream,
12}
13
14impl StreamBuilder {
15    /// Generate a new StreamBuilder
16    pub fn new() -> Self {
17        Self {
18            stream: TokenStream::new(),
19        }
20    }
21
22    /// Add multiple `TokenTree` items to the stream.
23    pub fn extend(&mut self, item: impl IntoIterator<Item = TokenTree>) -> &mut Self {
24        self.stream.extend(item);
25        self
26    }
27
28    /// Append another StreamBuilder to the current StreamBuilder.
29    pub fn append(&mut self, builder: StreamBuilder) -> &mut Self {
30        self.stream.extend(builder.stream);
31        self
32    }
33
34    /// Push a single token to the stream.
35    pub fn push(&mut self, item: impl Into<TokenTree>) -> &mut Self {
36        self.stream.extend([item.into()]);
37        self
38    }
39
40    /// Attempt to parse the given string as valid Rust code, and append the parsed result to the internal stream.
41    ///
42    /// Currently panics if the string could not be parsed as valid Rust code.
43    pub fn push_parsed(&mut self, item: impl AsRef<str>) -> Result<&mut Self> {
44        let tokens = TokenStream::from_str(item.as_ref()).map_err(|e| PushParseError {
45            error: e,
46            code: item.as_ref().to_string(),
47        })?;
48        self.stream.extend(tokens);
49        Ok(self)
50    }
51
52    /// Push a single ident to the stream. An ident is any word that a code file may contain, e.g. `fn`, `struct`, `where`, names of functions and structs, etc.
53    pub fn ident(&mut self, ident: Ident) -> &mut Self {
54        self.stream.extend([TokenTree::Ident(ident)]);
55        self
56    }
57
58    /// Push a single ident to the stream. An ident is any word that a code file may contain, e.g. `fn`, `struct`, `where`, names of functions and structs, etc.
59    pub fn ident_str(&mut self, ident: impl AsRef<str>) -> &mut Self {
60        self.stream.extend([TokenTree::Ident(Ident::new(
61            ident.as_ref(),
62            Span::call_site(),
63        ))]);
64        self
65    }
66
67    /// Add a group. A group is any block surrounded by `{ .. }`, `[ .. ]` or `( .. )`.
68    ///
69    /// `delim` indicates which group it is. The `inner` callback is used to fill the contents of the group.
70    pub fn group<FN>(&mut self, delim: Delimiter, inner: FN) -> crate::Result<&mut Self>
71    where
72        FN: FnOnce(&mut StreamBuilder) -> crate::Result<()>,
73    {
74        let mut stream = StreamBuilder::new();
75        inner(&mut stream)?;
76        self.stream
77            .extend([TokenTree::Group(Group::new(delim, stream.stream))]);
78        Ok(self)
79    }
80
81    /// Add a single punctuation to the stream. Puncts are single-character tokens like `.`, `<`, `#`, etc
82    ///
83    /// Note that this should not be used for multi-punct constructions like `::` or `->`. For that use [`puncts`] instead.
84    ///
85    /// [`puncts`]: #method.puncts
86    pub fn punct(&mut self, p: char) -> &mut Self {
87        self.stream
88            .extend([TokenTree::Punct(Punct::new(p, Spacing::Alone))]);
89        self
90    }
91
92    /// Add multiple punctuations to the stream. Multi punct tokens are e.g. `::`, `->` and `=>`.
93    ///
94    /// Note that this is the only way to add multi punct tokens.
95    /// If you were to use [`Punct`] to insert `->` it would be inserted as `-` and then `>`, and not form a single token. Rust would interpret this as a "minus sign and then a greater than sign", not as a single arrow.
96    pub fn puncts(&mut self, puncts: &str) -> &mut Self {
97        self.stream.extend(
98            puncts
99                .chars()
100                .map(|char| TokenTree::Punct(Punct::new(char, Spacing::Joint))),
101        );
102        self
103    }
104
105    /// Add a lifetime to the stream.
106    ///
107    /// Note that this is the only way to add lifetimes, if you were to do:
108    /// ```ignore
109    /// builder.punct('\'');
110    /// builder.ident_str("static");
111    /// ```
112    /// It would not add `'static`, but instead it would add `' static` as seperate tokens, and the lifetime would not work.
113    pub fn lifetime(&mut self, lt: Ident) -> &mut Self {
114        self.stream.extend([
115            TokenTree::Punct(Punct::new('\'', Spacing::Joint)),
116            TokenTree::Ident(lt),
117        ]);
118        self
119    }
120
121    /// Add a lifetime to the stream.
122    ///
123    /// Note that this is the only way to add lifetimes, if you were to do:
124    /// ```ignore
125    /// builder.punct('\'');
126    /// builder.ident_str("static");
127    /// ```
128    /// It would not add `'static`, but instead it would add `' static` as seperate tokens, and the lifetime would not work.
129    pub fn lifetime_str(&mut self, lt: &str) -> &mut Self {
130        self.stream.extend([
131            TokenTree::Punct(Punct::new('\'', Spacing::Joint)),
132            TokenTree::Ident(Ident::new(lt, Span::call_site())),
133        ]);
134        self
135    }
136
137    /// Add a literal string (`&'static str`) to the stream.
138    pub fn lit_str(&mut self, str: impl AsRef<str>) -> &mut Self {
139        self.stream
140            .extend([TokenTree::Literal(Literal::string(str.as_ref()))]);
141        self
142    }
143
144    /// Add an `usize` value to the stream.
145    pub fn lit_usize(&mut self, val: usize) -> &mut Self {
146        self.stream
147            .extend([TokenTree::Literal(Literal::usize_unsuffixed(val))]);
148        self
149    }
150
151    /// Set the given span on all tokens in the stream. This span is used by rust for e.g. compiler errors, to indicate the position of the error.
152    ///
153    /// Normally your derive will report an error on the derive, e.g.:
154    ///
155    /// ```text
156    /// #[derive(YourMacro)]
157    ///          ^^^^^^^^^
158    ///          |
159    ///          `self` value is a keyword only available in methods with a `self` parameter
160    /// ```
161    ///
162    /// If you want to improve feedback to the user of your macro, you can use this macro to set the location for a given streambuilder.
163    ///
164    /// A `span` can be obtained from e.g. an ident with `ident.span()`.
165    pub fn set_span_on_all_tokens(&mut self, span: Span) {
166        self.stream = std::mem::take(&mut self.stream)
167            .into_iter()
168            .map(|mut token| {
169                token.set_span(span);
170                token
171            })
172            .collect();
173    }
174}
175
176/// Failed to parse the code passed to [`StreamBuilder::push_parsed`]
177///
178/// [`StreamBuilder::push_parsed`]: struct.StreamBuilder.html#method.push_parsed
179#[derive(Debug)]
180pub struct PushParseError {
181    /// The parsing error
182    pub error: LexError,
183    /// The code that was being parsed
184    pub code: String,
185}